Improve public key handling

internal/data/users.go

@@ -1,12 +1,14 @@
 package data
 
 import (
-	"bytes"
 	"encoding/json"
 	"errors"
 	"sync"
 
 	"git.tcp.direct/kayos/common/entropy"
+	"git.tcp.direct/kayos/common/squish"
+	"github.com/davecgh/go-spew/spew"
+	"github.com/gliderlabs/ssh"
 	"github.com/rs/zerolog/log"
 )
 
@@ -31,9 +33,13 @@ func AuthMethodFromMap(m map[string]string) AuthMethod {
 			Password: m["password"],
 		}
 	case "publickey":
+		pkParsed, err := ssh.ParsePublicKey(squish.B64d(m["pubkey"]))
+		if err != nil {
+			return nil
+		}
 		return &PubKey{
 			Username: m["pub_username"],
-			Pub:      []byte(m["pubkey"]),
+			Pub:      pkParsed,
 		}
 	}
 	return nil
@@ -107,8 +113,15 @@ func (up *UserPass) Authenticate() error {
 }
 
 type PubKey struct {
-	Username string `json:"pub_username"`
+	Username string        `json:"pub_username"`
-	Pub      []byte `json:"pubkey"`
+	Pub      ssh.PublicKey `json:"pubkey"`
+}
+
+func NewPubKey(username string, pubkey ssh.PublicKey) *PubKey {
+	return &PubKey{
+		Username: username,
+		Pub:      pubkey,
+	}
 }
 
 func (pk *PubKey) Name() string {
@@ -119,7 +132,7 @@ func (pk *PubKey) Map() map[string]string {
 	return map[string]string{
 		"type":         "publickey",
 		"pub_username": pk.Username,
-		"pubkey":       string(pk.Pub),
+		"pubkey":       squish.B64e(pk.Pub.Marshal()),
 	}
 }
 
@@ -137,7 +150,22 @@ func (pk *PubKey) Authenticate() error {
 	for _, method := range user.AuthMethods {
 		switch method["type"] {
 		case "publickey":
-			if method["pub_username"] == pk.Username && bytes.Equal([]byte(method["pubkey"]), pk.Pub) {
+			if method["pub_username"] != pk.Username {
+				log.Warn().Str("username", pk.Username).Msg("username mismatch")
+				continue
+			}
+			pkdat, ok := method["pubkey"]
+			if !ok {
+				log.Warn().Str("username", pk.Username).Msg("pubkey not found")
+				continue
+			}
+			pubkeyParsed, err := ssh.ParsePublicKey(squish.B64d(pkdat))
+			if err != nil {
+				log.Warn().Err(err).Str("username", pk.Username).Msg("error parsing public key")
+				spew.Dump(method)
+				continue
+			}
+			if ssh.KeysEqual(pubkeyParsed, pk.Pub) {
 				return nil
 			}
 		default:
@@ -187,7 +215,7 @@ func NewUser(username string, authMethods ...AuthMethod) (*User, error) {
 			if len(usableMethod.Username) == 0 {
 				return nil, errors.New("username cannot be empty")
 			}
-			if len(usableMethod.Pub) == 0 {
+			if len(usableMethod.Pub.Marshal()) == 0 {
 				return nil, errors.New("public key cannot be empty")
 			}
 			methods = append(methods, method.Map())
@@ -226,16 +254,19 @@ func DelUser(username string) error {
 	return db.With("users").Delete([]byte(username))
 }
 
-func (user *User) DelPubKey(pubkey []byte) (*User, error) {
+func (user *User) DelPubKey(pubkey ssh.PublicKey) (*User, error) {
 	user.Lock()
 	defer user.Unlock()
 	var found = false
 	var methods []map[string]string
 	for _, method := range user.AuthMethods {
 		m := AuthMethodFromMap(method)
+		if m == nil {
+			continue
+		}
 		if m.Name() == "publickey" {
 			pubKey := m.(*PubKey)
-			if bytes.Equal(pubKey.Pub, pubkey) {
+			if ssh.KeysEqual(pubKey.Pub, pubkey) {
 				found = true
 				continue
 			}
@@ -260,6 +291,9 @@ func (user *User) ChangePassword(newPassword string) (*User, error) {
 	var methods []map[string]string
 	for _, method := range user.AuthMethods {
 		m := AuthMethodFromMap(method)
+		if m == nil {
+			continue
+		}
 		if m.Name() == "password" {
 			ponce.Do(func() {
 				hashed, err := HashPassword(newPassword)

internal/data/users_test.go

@@ -3,8 +3,33 @@ package data
 import (
 	"os"
 	"testing"
+
+	"golang.org/x/crypto/ssh"
 )
 
+var (
+	testPublicKey1 ssh.PublicKey
+	testPublicKey2 ssh.PublicKey
+	testPublicKey3 ssh.PublicKey
+)
+
+func init() {
+	var err error
+	// generate public keys for testing
+	testPublicKey1, _, _, _, err = ssh.ParseAuthorizedKey([]byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIO6EFqmelEJ6MELBPHUEFTGmlJBfhS7Jeq5B5BCrFSun"))
+	if err != nil {
+		panic(err)
+	}
+	testPublicKey2, _, _, _, err = ssh.ParseAuthorizedKey([]byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIH+ZTIMTWwYWHUEJlHfhT7dcYhgETGWgwEpDLdURaTPb"))
+	if err != nil {
+		panic(err)
+	}
+	testPublicKey3, _, _, _, err = ssh.ParseAuthorizedKey([]byte("ssh-ed25519 AAAAC3NzaC1lZDI1NTE5AAAAIHUEFpqqYCfBkVLRwgYlGbZyzgnEcMLpT0o97JUHNpIt"))
+	if err != nil {
+		panic(err)
+	}
+}
+
 func TestUsers(t *testing.T) {
 	testMode()
 	Start()
@@ -48,19 +73,33 @@ func TestUsers(t *testing.T) {
 		if user == nil {
 			t.Fatal("expected user to not be nil")
 		}
-		if user, err = user.AddAuthMethod(&PubKey{Username: "test2", Pub: []byte("pub")}); err != nil {
+		if user, err = user.AddAuthMethod(NewPubKey(user.Username, testPublicKey1)); err != nil {
 			t.Fatal(err)
 		}
 		if len(user.AuthMethods) != 2 {
 			t.Fatalf("expected 2 auth methods, got %d", len(user.AuthMethods))
 		}
-		pk := &PubKey{Username: "test2", Pub: []byte("pub")}
+		pk := NewPubKey("test2", testPublicKey1)
 		if err = pk.Authenticate(); err != nil {
-			t.Fatal("expected pub key to authenticate")
+			t.Fatal("expected pub key 1 to authenticate")
 		}
-		if user, err = user.AddAuthMethod(&PubKey{Username: "test2", Pub: []byte("pub2")}); err != nil {
+		pk = NewPubKey("test2", testPublicKey2)
+		if err = pk.Authenticate(); err == nil {
+			t.Fatal("expected pub key 2 to not authenticate")
+		}
+		if user, err = user.AddAuthMethod(NewPubKey(user.Username, testPublicKey2)); err != nil {
 			t.Fatal(err)
 		}
+
+		pk = NewPubKey("test2", testPublicKey1)
+		if err = pk.Authenticate(); err != nil {
+			t.Fatal("expected pub key 1 to authenticate")
+		}
+
+		pk = NewPubKey("test2", testPublicKey2)
+		if err = pk.Authenticate(); err != nil {
+			t.Fatal("expected pub key 2 to authenticate")
+		}
 		if len(user.AuthMethods) != 3 {
 			t.Fatalf("expected 2 auth methods, got %d", len(user.AuthMethods))
 		}
@@ -70,36 +109,33 @@ func TestUsers(t *testing.T) {
 		if user.AuthMethods[1]["type"] != "publickey" {
 			t.Fatalf("expected auth method to be 'publickey', got '%s'", user.AuthMethods[1]["type"])
 		}
-		auth := &PubKey{
-			Username: "test2",
-			Pub:      []byte("pub"),
-		}
-		if err = auth.Authenticate(); err != nil {
-			t.Fatalf("expected auth to succeed, got: %v", err)
-		}
-		auth.Pub = []byte("asdjfas")
-		if err = auth.Authenticate(); err == nil {
-			t.Fatal("expected auth to fail")
-		}
 	})
 	t.Run("DelPubKey", func(t *testing.T) {
 		user, err := GetUser("test2")
 		if err != nil {
 			t.Fatal(err)
 		}
-		if user, err = user.DelPubKey([]byte("fdsafdas")); err == nil {
+		if user, err = user.DelPubKey(testPublicKey3); err == nil {
 			t.Fatal("expected error deleting non-existent key")
 		}
 		if user == nil {
 			t.Fatal("expected user to not be nil")
 		}
-		if user, err = user.DelPubKey([]byte("pub2")); err != nil {
+		if user, err = user.DelPubKey(testPublicKey2); err != nil {
 			t.Fatal(err)
 		}
 		auth := NewUserPass(false, "test2", "test2")
 		if err = auth.Authenticate(); err != nil {
 			t.Fatalf("expected userpass to still be there after deleting public key, got: %v", err)
 		}
+		pk := &PubKey{"test2", testPublicKey2}
+		if err = pk.Authenticate(); err == nil {
+			t.Fatal("expected public key 2 to be deleted")
+		}
+		pk = &PubKey{"test2", testPublicKey1}
+		if err = pk.Authenticate(); err != nil {
+			t.Fatal("expected public key 1 to not be deleted")
+		}
 	})
 	t.Run("ChangePassword", func(t *testing.T) {
 		user, err := GetUser("test2")

internal/sshui/server.go

@@ -3,6 +3,7 @@ package sshui
 import (
 	"crypto/rand"
 	"crypto/rsa"
+	"fmt"
 	"os"
 	"path/filepath"
 
@@ -12,33 +13,49 @@ import (
 	"git.tcp.direct/kayos/ziggs/internal/data"
 )
 
+func newHostKey() error {
+	privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
+	if err != nil {
+		return err
+	}
+	if err = privateKey.Validate(); err != nil {
+		return err
+	}
+	dir, _ := filepath.Split(config.Filename)
+	newFile := filepath.Join(dir, "host_rsa")
+	if err = os.WriteFile(newFile, encodePrivateKeyToPEM(privateKey), 0600); err != nil {
+		return err
+	}
+	config.Snek.Set("ssh.host_key", newFile)
+	config.SSHHostKey = newFile
+	if err = config.Snek.WriteConfig(); err != nil {
+		return fmt.Errorf("viper config save error: %v", err)
+	}
+	return nil
+}
+
 func ServeSSH() error {
 	var opts []ssh.Option
 
-	switch config.SSHHostKey {
+	if config.SSHHostKey == "" {
-	case "":
+		if err := newHostKey(); err != nil {
-		privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
-		if err != nil {
 			return err
 		}
-		if err = privateKey.Validate(); err != nil {
-			return err
-		}
-		dir, _ := filepath.Split(config.Filename)
-		newFile := filepath.Join(dir, "host_rsa")
-		if err = os.WriteFile(newFile, encodePrivateKeyToPEM(privateKey), 0600); err != nil {
-			return err
-		}
-		config.Snek.Set("ssh.host_key", newFile)
-	default:
-		opts = append(opts, ssh.HostKeyFile(config.SSHHostKey))
 	}
 
+	opts = append(opts, ssh.HostKeyFile(config.SSHHostKey))
+
 	opts = append(opts, ssh.PasswordAuth(func(ctx ssh.Context, password string) bool {
 		attempt := data.NewUserPass(false, ctx.User(), password)
 		err := attempt.Authenticate()
 		return err == nil
 	}))
 
+	opts = append(opts, ssh.PublicKeyAuth(func(ctx ssh.Context, key ssh.PublicKey) bool {
+		attempt := data.NewPubKey(ctx.User(), key)
+		err := attempt.Authenticate()
+		return err == nil
+	}))
+
 	return ssh.ListenAndServe(config.SSHListen, nil, opts...)
 }