first commit

agent/internal/config/config.go

@@ -0,0 +1,45 @@
+package config
+
+import (
+	"os"
+	"time"
+
+	"gopkg.in/yaml.v3"
+)
+
+const ConfigPath = "/etc/keymanager/config.yaml"
+
+type Config struct {
+	ServerURL    string        `yaml:"server_url"`
+	ServerID     string        `yaml:"server_id"`
+	PreRegToken  string        `yaml:"pre_reg_token"`
+	AgentToken   string        `yaml:"agent_token"`
+	PollInterval time.Duration `yaml:"poll_interval"`
+	TLS          bool          `yaml:"tls"`
+}
+
+func Load() (*Config, error) {
+	data, err := os.ReadFile(ConfigPath)
+	if err != nil {
+		return nil, err
+	}
+	var cfg Config
+	if err := yaml.Unmarshal(data, &cfg); err != nil {
+		return nil, err
+	}
+	if cfg.PollInterval == 0 {
+		cfg.PollInterval = 30 * time.Second
+	}
+	return &cfg, nil
+}
+
+func Save(cfg *Config) error {
+	data, err := yaml.Marshal(cfg)
+	if err != nil {
+		return err
+	}
+	if err := os.MkdirAll("/etc/keymanager", 0700); err != nil {
+		return err
+	}
+	return os.WriteFile(ConfigPath, data, 0600)
+}

agent/internal/grpc/client.go

@@ -0,0 +1,100 @@
+package grpcclient
+
+import (
+	"context"
+	"crypto/tls"
+	"time"
+
+	"github.com/mrhid6/keymanager/agent/internal/grpc/pb"
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/credentials"
+	"google.golang.org/grpc/credentials/insecure"
+	"google.golang.org/grpc/encoding"
+)
+
+func init() {
+	encoding.RegisterCodec(JSONCodec{})
+}
+
+type Client struct {
+	conn   *grpc.ClientConn
+	client pb.KeyManagerClient
+}
+
+func New(serverURL string, useTLS bool) (*Client, error) {
+	var dialOpts []grpc.DialOption
+
+	if useTLS {
+		tlsCfg := &tls.Config{
+			InsecureSkipVerify: false,
+		}
+		creds := credentials.NewTLS(tlsCfg)
+		dialOpts = append(dialOpts, grpc.WithTransportCredentials(creds))
+	} else {
+		dialOpts = append(dialOpts, grpc.WithTransportCredentials(insecure.NewCredentials()))
+	}
+
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	conn, err := grpc.DialContext(ctx, serverURL, dialOpts...)
+	if err != nil {
+		return nil, err
+	}
+
+	return &Client{
+		conn:   conn,
+		client: pb.NewKeyManagerClient(conn),
+	}, nil
+}
+
+func (c *Client) Close() error {
+	return c.conn.Close()
+}
+
+func (c *Client) Register(serverID, preRegToken, hostname, ipAddress, osInfo string) (string, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	resp, err := c.client.Register(ctx, &pb.RegisterRequest{
+		ServerId:    serverID,
+		PreRegToken: preRegToken,
+		Hostname:    hostname,
+		IpAddress:   ipAddress,
+		OsInfo:      osInfo,
+	})
+	if err != nil {
+		return "", err
+	}
+	return resp.AgentToken, nil
+}
+
+func (c *Client) SyncKeys(serverID, agentToken string) ([]string, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	resp, err := c.client.SyncKeys(ctx, &pb.SyncRequest{
+		ServerId:   serverID,
+		AgentToken: agentToken,
+	})
+	if err != nil {
+		return nil, err
+	}
+	return resp.PublicKeys, nil
+}
+
+func (c *Client) UploadGeneratedKey(serverID, agentToken, publicKey, label string) (string, error) {
+	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
+	defer cancel()
+
+	resp, err := c.client.UploadGeneratedKey(ctx, &pb.UploadKeyRequest{
+		ServerId:   serverID,
+		AgentToken: agentToken,
+		PublicKey:  publicKey,
+		Label:      label,
+	})
+	if err != nil {
+		return "", err
+	}
+	return resp.KeyId, nil
+}

agent/internal/grpc/codec.go

@@ -0,0 +1,17 @@
+package grpcclient
+
+import "encoding/json"
+
+type JSONCodec struct{}
+
+func (JSONCodec) Marshal(v interface{}) ([]byte, error) {
+	return json.Marshal(v)
+}
+
+func (JSONCodec) Unmarshal(data []byte, v interface{}) error {
+	return json.Unmarshal(data, v)
+}
+
+func (JSONCodec) Name() string {
+	return "proto"
+}

agent/internal/grpc/pb/keymanager.pb.go

@@ -0,0 +1,93 @@
+// Hand-written gRPC bindings for keymanager.proto (agent side, JSON codec).
+
+package pb
+
+import (
+	"context"
+
+	"google.golang.org/grpc"
+	"google.golang.org/grpc/codes"
+	"google.golang.org/grpc/status"
+)
+
+type RegisterRequest struct {
+	ServerId    string `json:"server_id"`
+	PreRegToken string `json:"pre_reg_token"`
+	Hostname    string `json:"hostname"`
+	IpAddress   string `json:"ip_address"`
+	OsInfo      string `json:"os_info"`
+}
+
+type RegisterResponse struct {
+	AgentToken string `json:"agent_token"`
+}
+
+type SyncRequest struct {
+	ServerId   string `json:"server_id"`
+	AgentToken string `json:"agent_token"`
+}
+
+type SyncResponse struct {
+	PublicKeys []string `json:"public_keys"`
+}
+
+type UploadKeyRequest struct {
+	ServerId   string `json:"server_id"`
+	AgentToken string `json:"agent_token"`
+	PublicKey  string `json:"public_key"`
+	Label      string `json:"label"`
+}
+
+type UploadKeyResponse struct {
+	KeyId string `json:"key_id"`
+}
+
+type KeyManagerClient interface {
+	Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error)
+	SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error)
+	UploadGeneratedKey(ctx context.Context, in *UploadKeyRequest, opts ...grpc.CallOption) (*UploadKeyResponse, error)
+}
+
+type UnimplementedKeyManagerServer struct{}
+
+func (UnimplementedKeyManagerServer) Register(context.Context, *RegisterRequest) (*RegisterResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "not implemented")
+}
+func (UnimplementedKeyManagerServer) SyncKeys(context.Context, *SyncRequest) (*SyncResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "not implemented")
+}
+func (UnimplementedKeyManagerServer) UploadGeneratedKey(context.Context, *UploadKeyRequest) (*UploadKeyResponse, error) {
+	return nil, status.Errorf(codes.Unimplemented, "not implemented")
+}
+
+type keyManagerClient struct {
+	cc grpc.ClientConnInterface
+}
+
+func NewKeyManagerClient(cc grpc.ClientConnInterface) KeyManagerClient {
+	return &keyManagerClient{cc}
+}
+
+func (c *keyManagerClient) Register(ctx context.Context, in *RegisterRequest, opts ...grpc.CallOption) (*RegisterResponse, error) {
+	out := new(RegisterResponse)
+	if err := c.cc.Invoke(ctx, "/keymanager.v1.KeyManager/Register", in, out, opts...); err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *keyManagerClient) SyncKeys(ctx context.Context, in *SyncRequest, opts ...grpc.CallOption) (*SyncResponse, error) {
+	out := new(SyncResponse)
+	if err := c.cc.Invoke(ctx, "/keymanager.v1.KeyManager/SyncKeys", in, out, opts...); err != nil {
+		return nil, err
+	}
+	return out, nil
+}
+
+func (c *keyManagerClient) UploadGeneratedKey(ctx context.Context, in *UploadKeyRequest, opts ...grpc.CallOption) (*UploadKeyResponse, error) {
+	out := new(UploadKeyResponse)
+	if err := c.cc.Invoke(ctx, "/keymanager.v1.KeyManager/UploadGeneratedKey", in, out, opts...); err != nil {
+		return nil, err
+	}
+	return out, nil
+}

agent/internal/keys/keys.go

@@ -0,0 +1,120 @@
+package keys
+
+import (
+	"crypto/md5"
+	"encoding/base64"
+	"fmt"
+	"os"
+	"os/exec"
+	"path/filepath"
+	"strings"
+)
+
+const authorizedKeysPath = "/root/.ssh/authorized_keys"
+
+func ReadAuthorizedKeys() ([]string, error) {
+	data, err := os.ReadFile(authorizedKeysPath)
+	if os.IsNotExist(err) {
+		return nil, nil
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	var lines []string
+	for _, line := range strings.Split(string(data), "\n") {
+		line = strings.TrimSpace(line)
+		if line != "" && !strings.HasPrefix(line, "#") {
+			lines = append(lines, line)
+		}
+	}
+	return lines, nil
+}
+
+func WriteAuthorizedKeys(keys []string) error {
+	dir := filepath.Dir(authorizedKeysPath)
+	if err := os.MkdirAll(dir, 0700); err != nil {
+		return fmt.Errorf("mkdir %s: %w", dir, err)
+	}
+
+	content := strings.Join(keys, "\n")
+	if len(keys) > 0 {
+		content += "\n"
+	}
+
+	tmpPath := authorizedKeysPath + ".tmp"
+	if err := os.WriteFile(tmpPath, []byte(content), 0600); err != nil {
+		return fmt.Errorf("write tmp: %w", err)
+	}
+
+	if err := os.Rename(tmpPath, authorizedKeysPath); err != nil {
+		os.Remove(tmpPath)
+		return fmt.Errorf("rename: %w", err)
+	}
+
+	return os.Chmod(authorizedKeysPath, 0600)
+}
+
+func FingerprintLines(lines []string) map[string]bool {
+	fp := make(map[string]bool, len(lines))
+	for _, line := range lines {
+		fp[fingerprint(line)] = true
+	}
+	return fp
+}
+
+func StateChanged(current, desired []string) bool {
+	if len(current) != len(desired) {
+		return true
+	}
+	cur := FingerprintLines(current)
+	for _, line := range desired {
+		if !cur[fingerprint(line)] {
+			return true
+		}
+	}
+	return false
+}
+
+func fingerprint(pubKey string) string {
+	parts := strings.Fields(pubKey)
+	if len(parts) < 2 {
+		return pubKey
+	}
+	raw, err := base64.StdEncoding.DecodeString(parts[1])
+	if err != nil {
+		return pubKey
+	}
+	sum := md5.Sum(raw)
+	var pairs []string
+	for _, b := range sum {
+		pairs = append(pairs, fmt.Sprintf("%02x", b))
+	}
+	return "MD5:" + strings.Join(pairs, ":")
+}
+
+// GenerateKeyPair generates an ed25519 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) {
+	if err := os.MkdirAll(filepath.Dir(keyPath), 0700); err != nil {
+		return "", err
+	}
+
+	args := []string{
+		"-t", "ed25519",
+		"-f", keyPath,
+		"-N", "",
+		"-C", comment,
+	}
+	cmd := exec.Command("ssh-keygen", args...)
+	out, err := cmd.CombinedOutput()
+	if err != nil {
+		return "", fmt.Errorf("ssh-keygen: %w: %s", err, out)
+	}
+
+	pubData, err := os.ReadFile(keyPath + ".pub")
+	if err != nil {
+		return "", fmt.Errorf("read pubkey: %w", err)
+	}
+	return strings.TrimSpace(string(pubData)), nil
+}

agent/internal/sync/sync.go

@@ -0,0 +1,129 @@
+package agentsync
+
+import (
+	"fmt"
+	"log"
+	"net"
+	"os"
+	"runtime"
+	"strings"
+	"time"
+
+	"github.com/mrhid6/keymanager/agent/internal/config"
+	grpcclient "github.com/mrhid6/keymanager/agent/internal/grpc"
+	"github.com/mrhid6/keymanager/agent/internal/keys"
+)
+
+func Run(cfg *config.Config) error {
+	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
+	if err != nil {
+		return fmt.Errorf("dial grpc: %w", err)
+	}
+	defer client.Close()
+
+	// Register if we have a pre-reg token
+	if cfg.PreRegToken != "" {
+		log.Println("registering with server...")
+		hostname, _ := os.Hostname()
+		ipAddress := localIP()
+		osInfo := fmt.Sprintf("%s %s", runtime.GOOS, runtime.GOARCH)
+
+		agentToken, err := client.Register(cfg.ServerID, cfg.PreRegToken, hostname, ipAddress, osInfo)
+		if err != nil {
+			return fmt.Errorf("registration failed: %w", err)
+		}
+
+		cfg.AgentToken = agentToken
+		cfg.PreRegToken = ""
+		if err := config.Save(cfg); err != nil {
+			return fmt.Errorf("save config: %w", err)
+		}
+		log.Println("registration successful")
+
+		// Reconnect with potentially updated state
+		client.Close()
+		client, err = grpcclient.New(cfg.ServerURL, cfg.TLS)
+		if err != nil {
+			return fmt.Errorf("reconnect: %w", err)
+		}
+	}
+
+	if cfg.AgentToken == "" {
+		return fmt.Errorf("no agent token available — registration required")
+	}
+
+	ticker := time.NewTicker(cfg.PollInterval)
+	defer ticker.Stop()
+
+	// Run immediately on startup
+	if err := poll(client, cfg); err != nil {
+		log.Printf("poll error: %v", err)
+	}
+
+	for range ticker.C {
+		if err := poll(client, cfg); err != nil {
+			log.Printf("poll error: %v", err)
+		}
+	}
+	return nil
+}
+
+func poll(client *grpcclient.Client, cfg *config.Config) error {
+	desired, err := client.SyncKeys(cfg.ServerID, cfg.AgentToken)
+	if err != nil {
+		return fmt.Errorf("SyncKeys: %w", err)
+	}
+
+	current, err := keys.ReadAuthorizedKeys()
+	if err != nil {
+		return fmt.Errorf("read authorized_keys: %w", err)
+	}
+
+	if !keys.StateChanged(current, desired) {
+		log.Println("authorized_keys unchanged, skipping write")
+		return nil
+	}
+
+	if err := keys.WriteAuthorizedKeys(desired); err != nil {
+		return fmt.Errorf("write authorized_keys: %w", err)
+	}
+	log.Printf("authorized_keys updated (%d keys)", len(desired))
+	return nil
+}
+
+func localIP() string {
+	addrs, err := net.InterfaceAddrs()
+	if err != nil {
+		return ""
+	}
+	for _, addr := range addrs {
+		if ipNet, ok := addr.(*net.IPNet); ok && !ipNet.IP.IsLoopback() {
+			if ipNet.IP.To4() != nil {
+				return ipNet.IP.String()
+			}
+		}
+	}
+	return ""
+}
+
+// GenerateAndUpload generates an SSH keypair and uploads the public key to the server.
+func GenerateAndUpload(cfg *config.Config, label string) error {
+	client, err := grpcclient.New(cfg.ServerURL, cfg.TLS)
+	if err != nil {
+		return err
+	}
+	defer client.Close()
+
+	keyPath := fmt.Sprintf("/root/.ssh/keymanager_%s", strings.ReplaceAll(label, " ", "_"))
+	pubKey, err := keys.GenerateKeyPair(keyPath, label)
+	if err != nil {
+		return err
+	}
+
+	keyID, err := client.UploadGeneratedKey(cfg.ServerID, cfg.AgentToken, pubKey, label)
+	if err != nil {
+		return err
+	}
+	log.Printf("uploaded generated key %s (key_id=%s)", label, keyID)
+	return nil
+}