sshkey.go

package remote

import (
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/pem"
	"fmt"
	"os"
	"path/filepath"
	"strings"

	"golang.org/x/crypto/ssh"
)

// SSHKeyPair holds the paths to a locally generated key pair.
type SSHKeyPair struct {
	PrivateKeyPath string
	PublicKeyPath  string
}

// GenerateKeyPair creates a 4096-bit RSA key pair (id_rsa, id_rsa.pub) in
// sshDir. Only the public key is uploaded to AWS; the private key never
// leaves the local machine.
func GenerateKeyPair(sshDir string) (*SSHKeyPair, error) {
	if err := os.MkdirAll(sshDir, 0o700); err != nil {
		return nil, fmt.Errorf("create ssh dir: %w", err)
	}

	privateKey, err := rsa.GenerateKey(rand.Reader, 4096)
	if err != nil {
		return nil, fmt.Errorf("generate RSA key: %w", err)
	}

	privPath := filepath.Join(sshDir, "id_rsa")
	pubPath := filepath.Join(sshDir, "id_rsa.pub")

	privPEM := &pem.Block{
		Type:  "RSA PRIVATE KEY",
		Bytes: x509.MarshalPKCS1PrivateKey(privateKey),
	}

	// Create with 0600 from the start so the key is never world-readable.
	privFile, err := os.OpenFile(privPath, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0o600)
	if err != nil {
		return nil, fmt.Errorf("create private key file: %w", err)
	}
	defer privFile.Close()

	if err := pem.Encode(privFile, privPEM); err != nil {
		return nil, fmt.Errorf("encode private key: %w", err)
	}
	// OpenFile only applies the mode on creation; chmod covers a
	// pre-existing file with looser permissions.
	if err := os.Chmod(privPath, 0o600); err != nil {
		return nil, fmt.Errorf("chmod private key: %w", err)
	}

	publicKey, err := ssh.NewPublicKey(&privateKey.PublicKey)
	if err != nil {
		return nil, fmt.Errorf("create SSH public key: %w", err)
	}

	pubBytes := ssh.MarshalAuthorizedKey(publicKey)
	if err := os.WriteFile(pubPath, pubBytes, 0o644); err != nil {
		return nil, fmt.Errorf("write public key: %w", err)
	}

	return &SSHKeyPair{PrivateKeyPath: privPath, PublicKeyPath: pubPath}, nil
}

func (k *SSHKeyPair) ReadPublicKey() (string, error) {
	data, err := os.ReadFile(k.PublicKeyPath)
	if err != nil {
		return "", err
	}
	return strings.TrimSpace(string(data)), nil
}

// Cleanup removes both key files. Call it only after the stack is deleted;
// the private key is needed for SSH until then.
func (k *SSHKeyPair) Cleanup() {
	os.Remove(k.PrivateKeyPath)
	os.Remove(k.PublicKeyPath)
}