From 5bcf4164fb8633b9537debaca7980d5ae85aa93a Mon Sep 17 00:00:00 2001
From: DaneEveritt <dane@daneeveritt.com>
Date: Sun, 15 May 2022 16:01:52 -0400
Subject: [PATCH] Add support for public key based auth

---
 remote/types.go | 20 ++++++++++++++------
 sftp/server.go  | 23 ++++++++++++++---------
 2 files changed, 28 insertions(+), 15 deletions(-)

diff --git a/remote/types.go b/remote/types.go
index f37baf0..9a85484 100644
--- a/remote/types.go
+++ b/remote/types.go
@@ -11,6 +11,11 @@ import (
 	"github.com/pterodactyl/wings/parser"
 )
 
+const (
+	SftpAuthPassword  = SftpAuthRequestType("password")
+	SftpAuthPublicKey = SftpAuthRequestType("public_key")
+)
+
 // A generic type allowing for easy binding use when making requests to API
 // endpoints that only expect a singular argument or something that would not
 // benefit from being a typed struct.
@@ -63,14 +68,17 @@ type RawServerData struct {
 	ProcessConfiguration json.RawMessage `json:"process_configuration"`
 }
 
+type SftpAuthRequestType string
+
 // SftpAuthRequest defines the request details that are passed along to the Panel
 // when determining if the credentials provided to Wings are valid.
 type SftpAuthRequest struct {
-	User          string `json:"username"`
-	Pass          string `json:"password"`
-	IP            string `json:"ip"`
-	SessionID     []byte `json:"session_id"`
-	ClientVersion []byte `json:"client_version"`
+	Type          SftpAuthRequestType `json:"type"`
+	User          string              `json:"username"`
+	Pass          string              `json:"password"`
+	IP            string              `json:"ip"`
+	SessionID     []byte              `json:"session_id"`
+	ClientVersion []byte              `json:"client_version"`
 }
 
 // SftpAuthResponse is returned by the Panel when a pair of SFTP credentials
@@ -79,7 +87,7 @@ type SftpAuthRequest struct {
 // user for the SFTP subsystem.
 type SftpAuthResponse struct {
 	Server      string   `json:"server"`
-	Token       string   `json:"token"`
+	PublicKeys  []string `json:"public_keys"`
 	Permissions []string `json:"permissions"`
 }
 
diff --git a/sftp/server.go b/sftp/server.go
index 8a7a194..387a466 100644
--- a/sftp/server.go
+++ b/sftp/server.go
@@ -68,9 +68,14 @@ func (c *SFTPServer) Run() error {
 	}
 
 	conf := &ssh.ServerConfig{
-		NoClientAuth:     false,
-		MaxAuthTries:     6,
-		PasswordCallback: c.passwordCallback,
+		NoClientAuth: false,
+		MaxAuthTries: 6,
+		PasswordCallback: func(conn ssh.ConnMetadata, password []byte) (*ssh.Permissions, error) {
+			return c.makeCredentialsRequest(conn, remote.SftpAuthPassword, string(password))
+		},
+		PublicKeyCallback: func(conn ssh.ConnMetadata, key ssh.PublicKey) (*ssh.Permissions, error) {
+			return c.makeCredentialsRequest(conn, remote.SftpAuthPublicKey, string(key.Marshal()))
+		},
 	}
 	conf.AddHostKey(private)
 
@@ -177,17 +182,17 @@ func (c *SFTPServer) generateED25519PrivateKey() error {
 	return nil
 }
 
-// A function capable of validating user credentials with the Panel API.
-func (c *SFTPServer) passwordCallback(conn ssh.ConnMetadata, pass []byte) (*ssh.Permissions, error) {
+func (c *SFTPServer) makeCredentialsRequest(conn ssh.ConnMetadata, t remote.SftpAuthRequestType, p string) (*ssh.Permissions, error) {
 	request := remote.SftpAuthRequest{
+		Type:          t,
 		User:          conn.User(),
-		Pass:          string(pass),
+		Pass:          p,
 		IP:            conn.RemoteAddr().String(),
 		SessionID:     conn.SessionID(),
 		ClientVersion: conn.ClientVersion(),
 	}
 
-	logger := log.WithFields(log.Fields{"subsystem": "sftp", "username": conn.User(), "ip": conn.RemoteAddr().String()})
+	logger := log.WithFields(log.Fields{"subsystem": "sftp", "method": request.Type, "username": request.User, "ip": request.IP})
 	logger.Debug("validating credentials for SFTP connection")
 
 	if !validUsernameRegexp.MatchString(request.User) {
@@ -206,7 +211,7 @@ func (c *SFTPServer) passwordCallback(conn ssh.ConnMetadata, pass []byte) (*ssh.
 	}
 
 	logger.WithField("server", resp.Server).Debug("credentials validated and matched to server instance")
-	sshPerm := &ssh.Permissions{
+	permissions := ssh.Permissions{
 		Extensions: map[string]string{
 			"uuid":        resp.Server,
 			"user":        conn.User(),
@@ -214,7 +219,7 @@ func (c *SFTPServer) passwordCallback(conn ssh.ConnMetadata, pass []byte) (*ssh.
 		},
 	}
 
-	return sshPerm, nil
+	return &permissions, nil
 }
 
 // PrivateKeyPath returns the path the host private key for this server instance.