socks5 setupProxyDialer refactoring

prot0s34 · Feb 7, 2024 · 038c08a · 038c08a
1 parent df24441
commit 038c08a
Show file tree

Hide file tree

Showing 2 changed files with 51 additions and 82 deletions.
diff --git a/ssm/session.go b/ssm/session.go
@@ -3,16 +3,11 @@ package main
 import (
 	"fmt"
 	"golang.org/x/crypto/ssh"
-	"golang.org/x/net/proxy"
 	"log"
-	"os"
-	"os/exec"
-	"time"
 )
 
 func executeSSHCommand(targetUsername, targetPassword, targetHost string) {
 	log.Printf("Configuring SSH client for %s...\n", targetHost)
-
 	config := &ssh.ClientConfig{
 		User: targetUsername,
 		Auth: []ssh.AuthMethod{
@@ -22,14 +17,12 @@ func executeSSHCommand(targetUsername, targetPassword, targetHost string) {
 	}
 
 	log.Printf("Connecting to SSH server %s...\n", targetHost)
-
 	client, err := ssh.Dial("tcp", targetHost+":22", config)
 	if err != nil {
 		log.Fatalf("Failed to dial: %s", err)
 	}
 	defer client.Close()
 	log.Println("SSH server connection established.")
-
 	log.Println("Creating new SSH session...")
 
 	session, err := client.NewSession()
@@ -98,48 +91,11 @@ func executeSSHJumpCommand(jumpUsername, jumpPassword, jumpHost, targetUsername,
 }
 
 func executeSSHKubeCommand(kubeconfigPath, namespace, podName, targetUsername, targetPassword, targetHost string) {
-	localPort := 49152
-	targetPort := 1080
-
-	if isPortOpen(localPort) {
-		log.Printf("Local port %d is already open. Attempting to use the existing forwarding...\n", localPort)
-	} else {
-		log.Println("Starting port forwarding...")
-		portForwardCmd := exec.Command("kubectl", "port-forward", "svc/dante", fmt.Sprintf("%d:%d", localPort, targetPort), "-n", namespace, "--kubeconfig", kubeconfigPath)
-		portForwardCmd.Stderr = os.Stderr
-
-		if err := portForwardCmd.Start(); err != nil {
-			log.Fatalf("Failed to start port-forwarding: %s", err)
-		}
-		log.Println("Port forwarding started.")
 
-		defer func() {
-			log.Println("Terminating port forwarding...")
-			if err := portForwardCmd.Process.Kill(); err != nil {
-				log.Printf("Failed to kill port-forwarding process: %s", err)
-			}
-			log.Println("Port forwarding terminated.")
-		}()
-
-		log.Println("Waiting for port forwarding to establish...")
-		if !waitForPortOpen(localPort, 10*time.Second) {
-			log.Fatalf("Timeout reached, port %d did not open", localPort)
-		}
-	}
-
-	log.Println("Creating SOCKS5 dialer...")
-	dialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("localhost:%d", localPort), nil, proxy.Direct)
+	conn, err := setupProxyDialer(kubeconfigPath, namespace, 49152, 1080, targetHost)
 	if err != nil {
-		log.Fatalf("Failed to create SOCKS5 dialer: %s", err)
+		log.Fatalf("Error setting up port forwarding and dialing: %v", err)
 	}
-	log.Println("SOCKS5 dialer created.")
-
-	log.Printf("Dialing SSH server %s via SOCKS5 proxy...\n", targetHost)
-	conn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", targetHost, 22))
-	if err != nil {
-		log.Fatalf("Failed to dial SSH server via SOCKS5 proxy: %s", err)
-	}
-	log.Println("SSH server dialed.")
 
 	log.Println("Setting up SSH connection...")
 	ncc, chans, reqs, err := ssh.NewClientConn(conn, fmt.Sprintf("%s:%d", targetHost, 22), &ssh.ClientConfig{
@@ -168,41 +124,6 @@ func executeSSHKubeCommand(kubeconfigPath, namespace, podName, targetUsername, t
 }
 
 func executeSSHKubeJumpCommand(kubeconfigPath, namespace, podName, jumpHost, jumpUsername, jumpPassword, targetUsername, targetPassword, targetHost string) {
-	localPort := 49152
-	targetPort := 1080
-
-	if isPortOpen(localPort) {
-		log.Printf("Local port %d is already open. Attempting to use the existing forwarding...\n", localPort)
-	} else {
-		log.Println("Starting port forwarding...")
-		portForwardCmd := exec.Command("kubectl", "port-forward", "svc/dante", fmt.Sprintf("%d:%d", localPort, targetPort), "-n", namespace, "--kubeconfig", kubeconfigPath)
-		portForwardCmd.Stderr = os.Stderr
-
-		if err := portForwardCmd.Start(); err != nil {
-			log.Fatalf("Failed to start port-forwarding: %s", err)
-		}
-		log.Println("Port forwarding started.")
-
-		defer func() {
-			log.Println("Terminating port forwarding...")
-			if err := portForwardCmd.Process.Kill(); err != nil {
-				log.Printf("Failed to kill port-forwarding process: %s", err)
-			}
-			log.Println("Port forwarding terminated.")
-		}()
-
-		log.Println("Waiting for port forwarding to establish...")
-		if !waitForPortOpen(localPort, 10*time.Second) {
-			log.Fatalf("Timeout reached, port %d did not open", localPort)
-		}
-	}
-
-	log.Println("Creating SOCKS5 dialer...")
-	dialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("localhost:%d", localPort), nil, proxy.Direct)
-	if err != nil {
-		log.Fatalf("Failed to create SOCKS5 dialer: %s", err)
-	}
-	log.Println("SOCKS5 dialer created.")
 
 	log.Printf("Connecting to jump host %s...\n", jumpHost)
 	jumpHostConfig := &ssh.ClientConfig{
@@ -213,7 +134,7 @@ func executeSSHKubeJumpCommand(kubeconfigPath, namespace, podName, jumpHost, jum
 		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
 	}
 
-	jumpHostConn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", jumpHost, 22))
+	jumpHostConn, err := setupProxyDialer(kubeconfigPath, namespace, 49152, 1080, jumpHost)
 	if err != nil {
 		log.Fatalf("Failed to dial jump host: %s", err)
 	}

diff --git a/ssm/socks5.go b/ssm/socks5.go
@@ -2,7 +2,11 @@ package main
 
 import (
 	"fmt"
+	"golang.org/x/net/proxy"
+	"log"
 	"net"
+	"os"
+	"os/exec"
 	"time"
 )
 
@@ -29,3 +33,47 @@ func waitForPortOpen(port int, timeout time.Duration) bool {
 	}
 	return false
 }
+
+func setupProxyDialer(kubeconfigPath, namespace string, localPort, targetPort int, targetHost string) (net.Conn, error) {
+	if isPortOpen(localPort) {
+		log.Printf("Local port %d is already open. Attempting to use the existing forwarding...\n", localPort)
+	} else {
+		log.Println("Starting port forwarding...")
+		portForwardCmd := exec.Command("kubectl", "port-forward", "svc/dante", fmt.Sprintf("%d:%d", localPort, targetPort), "-n", namespace, "--kubeconfig", kubeconfigPath)
+		portForwardCmd.Stderr = os.Stderr
+
+		if err := portForwardCmd.Start(); err != nil {
+			return nil, fmt.Errorf("failed to start port-forwarding: %w", err)
+		}
+		log.Println("Port forwarding started.")
+
+		defer func() {
+			log.Println("Terminating port forwarding...")
+			if err := portForwardCmd.Process.Kill(); err != nil {
+				log.Printf("Failed to kill port-forwarding process: %s", err)
+			}
+			log.Println("Port forwarding terminated.")
+		}()
+
+		log.Println("Waiting for port forwarding to establish...")
+		if !waitForPortOpen(localPort, 10*time.Second) {
+			return nil, fmt.Errorf("timeout reached, port %d did not open", localPort)
+		}
+	}
+
+	log.Println("Creating SOCKS5 dialer...")
+	dialer, err := proxy.SOCKS5("tcp", fmt.Sprintf("localhost:%d", localPort), nil, proxy.Direct)
+	if err != nil {
+		return nil, fmt.Errorf("failed to create SOCKS5 dialer: %w", err)
+	}
+	log.Println("SOCKS5 dialer created.")
+
+	log.Printf("Dialing SSH server %s via SOCKS5 proxy...\n", targetHost)
+	conn, err := dialer.Dial("tcp", fmt.Sprintf("%s:%d", targetHost, 22))
+	if err != nil {
+		return nil, fmt.Errorf("failed to dial SSH server via SOCKS5 proxy: %w", err)
+	}
+	log.Println("SSH server dialed.")
+
+	return conn, nil
+}