vantage-agent/internal/keys/keys.go

package keys

import (
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
)

const authorizedKeysPath = "/root/.ssh/authorized_keys"
const sshConfigPath = "/root/.ssh/config"
const managedConfigPath = "/root/.ssh/vantage.conf"
const includeDirective = "Include /root/.ssh/vantage.conf"

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, ":")
}


type KeyGenOptions struct {
	KeyType    string 
	KeySize    int    
	Passphrase string 
	Comment    string 
}



func GenerateKeyPair(keyPath string, opts KeyGenOptions) (string, error) {
	if err := os.MkdirAll(filepath.Dir(keyPath), 0700); err != nil {
		return "", err
	}

	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 {
		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
}



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
	}

	
	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 
		}
	}

	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
}


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
}




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 
		}
	}

	
	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
}