feat: support multi keys

.gitignore

@@ -1,3 +1,5 @@
+allowed_keys
+simple-ssh-server
 .idea/
 # ---> Go
 # Binaries for programs and plugins

go.mod

@@ -3,6 +3,7 @@ module git.hatter.ink/simple-ssh-server
 go 1.13
 
 require (
-	github.com/creack/pty v1.1.13 // indirect
+	github.com/creack/pty v1.1.13
 	github.com/gliderlabs/ssh v0.3.3
+	golang.org/x/crypto v0.0.0-20210616213533-5ff15b29337e
 )

go.sum

@@ -11,6 +11,7 @@ golang.org/x/sys v0.0.0-20201119102817-f84b799fce68/go.mod h1:h1NjWce9XRLGQEsW7w
 golang.org/x/sys v0.0.0-20210615035016-665e8c7367d1/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
 golang.org/x/sys v0.0.0-20210616094352-59db8d763f22 h1:RqytpXGR1iVNX7psjB3ff8y7sNFinVFvkx1c8SjBkio=
 golang.org/x/sys v0.0.0-20210616094352-59db8d763f22/go.mod h1:oPkhp1MJrh7nUepCBck5+mAzfO9JrbApNNgaTdGDITg=
+golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1 h1:v+OssWQX+hTHEmOBgwxdZxK4zHq3yOs8F9J7mk0PY8E=
 golang.org/x/term v0.0.0-20201126162022-7de9c90e9dd1/go.mod h1:bj7SfCRtBDWHUb9snDiAeCFNEtKQo2Wmx5Cou7ajbmo=
 golang.org/x/text v0.3.3/go.mod h1:5Zoc/QRtKVWzQhOtBMvqHzDpF6irO9z98xDceosuGiQ=
 golang.org/x/tools v0.0.0-20180917221912-90fa682c2a6e/go.mod h1:n7NCudcB/nEzxVGmLbDWY5pfWTLqBcC2KZ6jyYvM4mQ=

simple-ssh-server.go

@@ -32,13 +32,24 @@ const WELCOME = ",   .     .                     .
 	"                                                                     \n"
 
 func main() {
-	sshPublicKey, sshPublicKeyErr := parseSshPubkey()
+	args := os.Args
-	if sshPublicKeyErr != nil {
+	log.Println("Arguments: ", args)
-		log.Fatal("Parse sk ecdsa public key failed: ", sshPublicKeyErr)
+	port := ":2222"
+	if len(args) > 1 {
+		port = ":" + args[1]
+	}
+	log.Println("Use port: ", port)
+
+	allowedSshPublicKeys, allowedSshPublicKeysErr := parseAllowedSshPubkeysAsStrings()
+	if allowedSshPublicKeysErr != nil {
+		log.Fatal("Parse sk ecdsa public key(s) failed: ", allowedSshPublicKeysErr)
 		return
 	}
-	marshalSshPublicKey := marshalKey(sshPublicKey)
+	log.Println("Found sk ecdsa public keys: ", len(allowedSshPublicKeys))
-	log.Println("Found sk ecdsa public key: ", marshalSshPublicKey)
+	for i, k := range allowedSshPublicKeys {
+		log.Println(i, ">>", k)
+	}
+
 	hostKeyBytes, hostKeyBytesErr := readHostKey()
 	if hostKeyBytesErr != nil {
 		log.Fatal("Load host key failed: ", hostKeyBytesErr)
@@ -79,12 +90,18 @@ func main() {
 	publicKeyOption := ssh.PublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
 		marshalPubKey := marshalKey(key)
 		log.Println("Auth public key: ", marshalPubKey, ", from: ", ctx.RemoteAddr())
-		isAllowed := marshalPubKey == marshalSshPublicKey
+
-		log.Println("Key allowed: ", isAllowed)
+		for idxAllowedSshPublicKey, allowedSshPublicKey := range allowedSshPublicKeys {
-		return isAllowed
+			if marshalPubKey == allowedSshPublicKey {
+				log.Println("Key allowed: ", idxAllowedSshPublicKey)
+				return true
+			}
+		}
+		log.Println("Key NOT allowed")
+		return false
 	})
-	log.Println("Listening :2222...")
+	log.Println("Listening ", port, "...")
-	log.Fatal(ssh.ListenAndServe(":2222", nil, hostKeyOption, publicKeyOption))
+	log.Fatal(ssh.ListenAndServe(port, nil, hostKeyOption, publicKeyOption))
 }
 
 func setWinsize(f *os.File, w, h int) {
@@ -92,34 +109,63 @@ func setWinsize(f *os.File, w, h int) {
 		uintptr(unsafe.Pointer(&struct{ h, w, x, y uint16 }{uint16(h), uint16(w), 0, 0})))
 }
 
-func parseSshPubkey() (gossh.PublicKey, error) {
+func parseAllowedSshPubkeysAsStrings() ([]string, error) {
+	pubkeys, err := parseAllowedSshPubkeys()
+	if err != nil {
+		return nil, err
+	}
+	var pubkeysAsStrings []string
+	for _, pubkey := range pubkeys {
+		pubkeysAsStrings = append(pubkeysAsStrings, marshalKey(pubkey))
+	}
+	return pubkeysAsStrings, nil
+}
+
+func parseAllowedSshPubkeys() ([]gossh.PublicKey, error) {
 	pubkeyBytes, pubkeyErr := ioutil.ReadFile("allowed_keys")
 	if pubkeyErr != nil {
 		return nil, pubkeyErr
 	}
-	pubkey := strings.TrimSpace(string(pubkeyBytes))
+	var pubkeys []gossh.PublicKey
-	if strings.Contains(pubkey, " ") {
+	pubkeyLines := string(pubkeyBytes)
-		pubkey = strings.Split(pubkey, " ")[1]
+	pubkeySplitedLines := strings.Split(pubkeyLines, "\n")
+
+	for _, pubkeyLine := range pubkeySplitedLines {
+		pubkey := strings.TrimSpace(pubkeyLine)
+		if len(pubkey) > 0 && !strings.HasPrefix(pubkey, "#") {
+			pubkey = strings.Split(pubkey, " ")[1]
+			pubkeyBytes, pubkeyBytesErr := base64.StdEncoding.DecodeString(pubkey)
+			if pubkeyBytesErr != nil {
+				return nil, pubkeyBytesErr
+			}
+			publicKey, publicKeyErr := gossh.ParsePublicKey(pubkeyBytes)
+			if publicKeyErr != nil {
+				return nil, publicKeyErr
+			}
+			pubkeys = append(pubkeys, publicKey)
+		}
 	}
-	pubkeyBytes, pubkeyBytesErr := base64.StdEncoding.DecodeString(pubkey)
+	return pubkeys, nil
-	if pubkeyBytesErr != nil {
-		return nil, pubkeyBytesErr
-	}
-	publicKey, publicKeyErr := gossh.ParsePublicKey(pubkeyBytes)
-	if publicKeyErr != nil {
-		return nil, publicKeyErr
-	}
-	return publicKey, nil
 }
 
 func readHostKey() ([]byte, error) {
+	hostKeyEcdsaFile := "/etc/ssh/ssh_host_ecdsa_key"
+	hostKeyEcdsaFileBytes, hostKeyEcdsaFileBytesErr := ioutil.ReadFile(hostKeyEcdsaFile)
+	if hostKeyEcdsaFileBytesErr == nil {
+		log.Println("Found host key: ", hostKeyEcdsaFile)
+		return hostKeyEcdsaFileBytes, nil
+	}
+
 	hostKeyFile := "/etc/ssh/ssh_host_rsa_key"
 	hostKeyFileBytes, hostKeyFileBytesErr := ioutil.ReadFile(hostKeyFile)
 	if hostKeyFileBytesErr == nil {
+		log.Println("Found host key: ", hostKeyFile)
 		return hostKeyFileBytes, nil
 	}
+
 	tempHostKeyFileBytes, tempHostKeyFileBytesErr := ioutil.ReadFile("/Users/hatterjiang/.ssh/id_rsa")
 	if tempHostKeyFileBytesErr == nil {
+		log.Println("!!WARN!! Found host key: ", "~/.ssh/id_rsa")
 		return tempHostKeyFileBytes, nil
 	}
 	return nil, errors.New("Canot read any host key from file")