From ff6f0fadaa9dfb832fdb41b6c9cadcea5b68b8dc Mon Sep 17 00:00:00 2001
From: domrichardson <100129001+domrichardson@users.noreply.github.com>
Date: Tue, 16 Jun 2026 10:28:46 +0100
Subject: [PATCH] updates

---
 internal/grpc/pb/keymanager.pb.go |  6 +++++-
 internal/keys/keys.go             | 31 ++++++++++++++++++++++++-------
 internal/sync/sync.go             | 13 ++++++++++---
 3 files changed, 39 insertions(+), 11 deletions(-)

diff --git a/internal/grpc/pb/keymanager.pb.go b/internal/grpc/pb/keymanager.pb.go
index 79aecbd..748535a 100644
--- a/internal/grpc/pb/keymanager.pb.go
+++ b/internal/grpc/pb/keymanager.pb.go
@@ -50,7 +50,11 @@ type ServerCommand struct {
 }
 
 type GenerateKeyCmd struct {
-	Label string `json:"label"`
+	Label      string `json:"label"`
+	KeyType    string `json:"key_type,omitempty"`
+	KeySize    int    `json:"key_size,omitempty"`
+	Passphrase string `json:"passphrase,omitempty"`
+	Comment    string `json:"comment,omitempty"`
 }
 
 type AgentMessage struct {
diff --git a/internal/keys/keys.go b/internal/keys/keys.go
index e9a8cc4..00bff38 100644
--- a/internal/keys/keys.go
+++ b/internal/keys/keys.go
@@ -93,19 +93,36 @@ func fingerprint(pubKey string) string {
 	return "MD5:" + strings.Join(pairs, ":")
 }
 
-// GenerateKeyPair generates an ed25519 SSH keypair and returns the public key.
+// KeyGenOptions controls how ssh-keygen is invoked.
+type KeyGenOptions struct {
+	KeyType    string // ed25519 (default), rsa, ecdsa
+	KeySize    int    // bits; used for rsa and ecdsa
+	Passphrase string // empty = no passphrase
+	Comment    string // embedded in the public key
+}
+
+// GenerateKeyPair generates an SSH keypair and returns the public key.
 // The private key is written to keyPath; keyPath+".pub" holds the public key.
-func GenerateKeyPair(keyPath, comment string) (string, error) {
+func GenerateKeyPair(keyPath string, opts KeyGenOptions) (string, error) {
 	if err := os.MkdirAll(filepath.Dir(keyPath), 0700); err != nil {
 		return "", err
 	}
 
-	args := []string{
-		"-t", "ed25519",
-		"-f", keyPath,
-		"-N", "",
-		"-C", comment,
+	keyType := opts.KeyType
+	if keyType == "" {
+		keyType = "ed25519"
 	}
+
+	args := []string{
+		"-t", keyType,
+		"-f", keyPath,
+		"-N", opts.Passphrase,
+		"-C", opts.Comment,
+	}
+	if opts.KeySize > 0 && keyType != "ed25519" {
+		args = append(args, "-b", fmt.Sprintf("%d", opts.KeySize))
+	}
+
 	cmd := exec.Command("ssh-keygen", args...)
 	out, err := cmd.CombinedOutput()
 	if err != nil {
diff --git a/internal/sync/sync.go b/internal/sync/sync.go
index 2f6df2d..85cccff 100644
--- a/internal/sync/sync.go
+++ b/internal/sync/sync.go
@@ -166,10 +166,17 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error {
 }
 
 func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) {
-	label := cmd.GenerateKey.Label
+	g := cmd.GenerateKey
+	label := g.Label
 	keyPath := fmt.Sprintf("/root/.ssh/keymanager_%s", strings.ReplaceAll(label, " ", "_"))
 
-	pubKey, err := keys.GenerateKeyPair(keyPath, label)
+	opts := keys.KeyGenOptions{
+		KeyType:    g.KeyType,
+		KeySize:    g.KeySize,
+		Passphrase: g.Passphrase,
+		Comment:    g.Comment,
+	}
+	pubKey, err := keys.GenerateKeyPair(keyPath, opts)
 	if err != nil {
 		log.Printf("key generation failed (cmd=%s): %v", cmd.CommandId, err)
 		return
@@ -214,7 +221,7 @@ func GenerateAndUpload(cfg *config.Config, label string) error {
 	defer client.Close()
 
 	keyPath := fmt.Sprintf("/root/.ssh/keymanager_%s", strings.ReplaceAll(label, " ", "_"))
-	pubKey, err := keys.GenerateKeyPair(keyPath, label)
+	pubKey, err := keys.GenerateKeyPair(keyPath, keys.KeyGenOptions{Comment: label})
 	if err != nil {
 		return err
 	}