diff --git a/agent/internal/grpc/client.go b/agent/internal/grpc/client.go index 55e047e..c3d8199 100644 --- a/agent/internal/grpc/client.go +++ b/agent/internal/grpc/client.go @@ -87,7 +87,7 @@ func (c *Client) SyncKeys(serverID, agentToken string) ([]string, error) { return resp.PublicKeys, nil } -func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, label string) (string, error) { +func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, privateKey, label string) (string, error) { ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second) defer cancel() @@ -95,6 +95,7 @@ func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, label strin ServerId: serverID, AgentToken: agentToken, PublicKey: publicKey, + PrivateKey: privateKey, Label: label, }) if err != nil { diff --git a/agent/internal/grpc/pb/keymanager.pb.go b/agent/internal/grpc/pb/keymanager.pb.go index 748535a..514c8e9 100644 --- a/agent/internal/grpc/pb/keymanager.pb.go +++ b/agent/internal/grpc/pb/keymanager.pb.go @@ -36,6 +36,7 @@ type UploadKeyRequest struct { AgentToken string `json:"agent_token"` PublicKey string `json:"public_key"` Label string `json:"label"` + PrivateKey string `json:"private_key,omitempty"` } type UploadKeyResponse struct { @@ -47,6 +48,11 @@ type UploadKeyResponse struct { type ServerCommand struct { CommandId string `json:"command_id"` GenerateKey *GenerateKeyCmd `json:"generate_key,omitempty"` + DeleteKey *DeleteKeyCmd `json:"delete_key,omitempty"` +} + +type DeleteKeyCmd struct { + Label string `json:"label"` } type GenerateKeyCmd struct { diff --git a/agent/internal/keys/keys.go b/agent/internal/keys/keys.go index 00bff38..56ed6c4 100644 --- a/agent/internal/keys/keys.go +++ b/agent/internal/keys/keys.go @@ -11,6 +11,9 @@ import ( ) const authorizedKeysPath = "/root/.ssh/authorized_keys" +const sshConfigPath = "/root/.ssh/config" +const managedConfigPath = "/root/.ssh/keymanager.conf" +const includeDirective = "Include /root/.ssh/keymanager.conf" func ReadAuthorizedKeys() ([]string, error) { data, err := os.ReadFile(authorizedKeysPath) @@ -135,3 +138,91 @@ func GenerateKeyPair(keyPath string, opts KeyGenOptions) (string, error) { } return strings.TrimSpace(string(pubData)), nil } + +// AddSSHIdentity writes an IdentityFile entry for keyPath into the managed +// keymanager.conf include file, and ensures ~/.ssh/config includes it. +func AddSSHIdentity(keyPath string) error { + if err := os.MkdirAll(filepath.Dir(sshConfigPath), 0700); err != nil { + return fmt.Errorf("mkdir .ssh: %w", err) + } + + if err := ensureIncludeDirective(); err != nil { + return err + } + + // Read existing managed config (it may not exist yet). + var existing string + data, err := os.ReadFile(managedConfigPath) + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("read %s: %w", managedConfigPath, err) + } + existing = string(data) + + line := "IdentityFile " + keyPath + for _, l := range strings.Split(existing, "\n") { + if strings.TrimSpace(l) == line { + return nil // already present + } + } + + if existing != "" && !strings.HasSuffix(existing, "\n") { + existing += "\n" + } + updated := existing + line + "\n" + + if err := os.WriteFile(managedConfigPath, []byte(updated), 0600); err != nil { + return fmt.Errorf("write %s: %w", managedConfigPath, err) + } + return nil +} + +// RemoveSSHIdentity removes the IdentityFile entry for keyPath from the managed config. +func RemoveSSHIdentity(keyPath string) error { + data, err := os.ReadFile(managedConfigPath) + if os.IsNotExist(err) { + return nil + } + if err != nil { + return fmt.Errorf("read %s: %w", managedConfigPath, err) + } + + line := "IdentityFile " + keyPath + var kept []string + for _, l := range strings.Split(strings.TrimRight(string(data), "\n"), "\n") { + if strings.TrimSpace(l) != line { + kept = append(kept, l) + } + } + + content := strings.Join(kept, "\n") + if len(kept) > 0 { + content += "\n" + } + if err := os.WriteFile(managedConfigPath, []byte(content), 0600); err != nil { + return fmt.Errorf("write %s: %w", managedConfigPath, err) + } + return nil +} + +// ensureIncludeDirective adds "Include /root/.ssh/keymanager.conf" to the top +// of ~/.ssh/config if it is not already present. The Include must appear before +// any Host stanzas to be effective for all connections. +func ensureIncludeDirective() error { + data, err := os.ReadFile(sshConfigPath) + if err != nil && !os.IsNotExist(err) { + return fmt.Errorf("read %s: %w", sshConfigPath, err) + } + + for _, l := range strings.Split(string(data), "\n") { + if strings.TrimSpace(l) == includeDirective { + return nil // already present + } + } + + // Prepend the Include directive so it takes effect before any Host blocks. + updated := includeDirective + "\n" + string(data) + if err := os.WriteFile(sshConfigPath, []byte(updated), 0600); err != nil { + return fmt.Errorf("write %s: %w", sshConfigPath, err) + } + return nil +} diff --git a/agent/internal/sync/sync.go b/agent/internal/sync/sync.go index 85cccff..c815245 100644 --- a/agent/internal/sync/sync.go +++ b/agent/internal/sync/sync.go @@ -162,9 +162,28 @@ func connectAndHandleStream(ctx context.Context, cfg *config.Config) error { if cmd.GenerateKey != nil { go handleGenerateKey(cfg, cmd) } + if cmd.DeleteKey != nil { + go handleDeleteKey(cmd) + } } } +func handleDeleteKey(cmd *pb.ServerCommand) { + label := cmd.DeleteKey.Label + keyPath := fmt.Sprintf("/root/.ssh/keymanager_%s", strings.ReplaceAll(label, " ", "_")) + + if err := keys.RemoveSSHIdentity(keyPath); err != nil { + log.Printf("remove ssh identity failed (cmd=%s): %v", cmd.CommandId, err) + } + + for _, path := range []string{keyPath, keyPath + ".pub"} { + if err := os.Remove(path); err != nil && !os.IsNotExist(err) { + log.Printf("delete key file %s (cmd=%s): %v", path, cmd.CommandId, err) + } + } + log.Printf("deleted local key files for %q (cmd=%s)", label, cmd.CommandId) +} + func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) { g := cmd.GenerateKey label := g.Label @@ -182,6 +201,12 @@ func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) { return } + privKeyData, err := os.ReadFile(keyPath) + if err != nil { + log.Printf("read private key failed (cmd=%s): %v", cmd.CommandId, err) + return + } + client, err := grpcclient.New(cfg.ServerURL, cfg.TLS) if err != nil { log.Printf("dial for key upload failed (cmd=%s): %v", cmd.CommandId, err) @@ -189,11 +214,15 @@ func handleGenerateKey(cfg *config.Config, cmd *pb.ServerCommand) { } defer client.Close() - keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, label) + keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, string(privKeyData), label) if err != nil { log.Printf("key upload failed (cmd=%s): %v", cmd.CommandId, err) return } + + if err := keys.AddSSHIdentity(keyPath); err != nil { + log.Printf("add ssh identity failed (cmd=%s): %v", cmd.CommandId, err) + } log.Printf("generated and uploaded key %q (key_id=%s, cmd=%s)", label, keyID, cmd.CommandId) } @@ -226,10 +255,19 @@ func GenerateAndUpload(cfg *config.Config, label string) error { return err } - keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, label) + privKeyData, err := os.ReadFile(keyPath) + if err != nil { + return fmt.Errorf("read private key: %w", err) + } + + keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, string(privKeyData), label) if err != nil { return err } + + if err := keys.AddSSHIdentity(keyPath); err != nil { + log.Printf("add ssh identity: %v", err) + } log.Printf("uploaded generated key %s (key_id=%s)", label, keyID) return nil } diff --git a/proto/keymanager/v1/keymanager.proto b/proto/keymanager/v1/keymanager.proto index 754b93d..a456a8b 100644 --- a/proto/keymanager/v1/keymanager.proto +++ b/proto/keymanager/v1/keymanager.proto @@ -38,6 +38,7 @@ message UploadKeyRequest { string agent_token = 2; string public_key = 3; string label = 4; + string private_key = 5; } message UploadKeyResponse { @@ -67,9 +68,14 @@ message ServerCommand { string command_id = 1; oneof command { GenerateKeyCmd generate_key = 2; + DeleteKeyCmd delete_key = 3; } } +message DeleteKeyCmd { + string label = 1; +} + message GenerateKeyCmd { string label = 1; string key_type = 2; // ed25519 | rsa | ecdsa (default: ed25519) diff --git a/server/internal/grpc/pb/keymanager.pb.go b/server/internal/grpc/pb/keymanager.pb.go index 93b94fd..35dcba1 100644 --- a/server/internal/grpc/pb/keymanager.pb.go +++ b/server/internal/grpc/pb/keymanager.pb.go @@ -39,6 +39,7 @@ type UploadKeyRequest struct { AgentToken string `json:"agent_token"` PublicKey string `json:"public_key"` Label string `json:"label"` + PrivateKey string `json:"private_key,omitempty"` } type UploadKeyResponse struct { @@ -50,6 +51,11 @@ type UploadKeyResponse struct { type ServerCommand struct { CommandId string `json:"command_id"` GenerateKey *GenerateKeyCmd `json:"generate_key,omitempty"` + DeleteKey *DeleteKeyCmd `json:"delete_key,omitempty"` +} + +type DeleteKeyCmd struct { + Label string `json:"label"` } type GenerateKeyCmd struct { diff --git a/server/internal/grpc/server.go b/server/internal/grpc/server.go index 691fae4..73e24ba 100644 --- a/server/internal/grpc/server.go +++ b/server/internal/grpc/server.go @@ -54,7 +54,7 @@ func (s *keyManagerServer) UploadGeneratedKey(ctx context.Context, req *pb.Uploa return nil, status.Errorf(codes.Unauthenticated, "invalid agent token") } - key, err := services.CreateKey(req.Label, req.PublicKey, "generated", srv.ServerID, "") + key, err := services.CreateKey(req.Label, req.PublicKey, "generated", srv.ServerID, req.PrivateKey) if err != nil { return nil, status.Errorf(codes.Internal, "failed to store key: %v", err) } diff --git a/server/internal/services/dispatch.go b/server/internal/services/dispatch.go index 033e7ae..5a63a26 100644 --- a/server/internal/services/dispatch.go +++ b/server/internal/services/dispatch.go @@ -67,6 +67,22 @@ type KeyGenParams struct { Comment string } +// DispatchDeleteKey sends a delete-key command to the named server's agent. +// It is best-effort: if the agent is offline the local files will remain until next connection. +func DispatchDeleteKey(serverID, label string) { + if !Dispatcher.IsConnected(serverID) { + return + } + cmd := &pb.ServerCommand{ + CommandId: uuid.New().String(), + DeleteKey: &pb.DeleteKeyCmd{Label: label}, + } + if err := Dispatcher.dispatch(serverID, cmd); err != nil { + // Non-fatal: agent will clean up files on next manual intervention or reinstall. + _ = err + } +} + // DispatchGenerateKey sends a generate-key command to the named server's agent. // Returns the command ID that can be used to correlate the agent's result. func DispatchGenerateKey(serverID string, p KeyGenParams) (string, error) { diff --git a/server/internal/services/keys.go b/server/internal/services/keys.go index af40c35..5ebc708 100644 --- a/server/internal/services/keys.go +++ b/server/internal/services/keys.go @@ -126,11 +126,22 @@ func DeleteKey(keyID string) error { ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second) defer cancel() + var key models.Key + if err := db.Col("keys").FindOne(ctx, bson.M{"key_id": keyID}).Decode(&key); err != nil { + return err + } + if _, err := db.Col("keys").DeleteOne(ctx, bson.M{"key_id": keyID}); err != nil { return err } - _, err := db.Col("assignments").DeleteMany(ctx, bson.M{"key_id": keyID}) - return err + if _, err := db.Col("assignments").DeleteMany(ctx, bson.M{"key_id": keyID}); err != nil { + return err + } + + if key.Source == "generated" && key.GeneratedByServerID != "" { + DispatchDeleteKey(key.GeneratedByServerID, key.Label) + } + return nil } func AssignKey(keyID, serverID string) (*models.Assignment, error) {