ollie-core/pkg/elevate/sshagent.go

package elevate

import (
	"fmt"
	"net"
	"os"
	"time"

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

// SSHAgentProxy proxies SSH agent requests to the real agent,
// intercepting sign requests for approval via the elevation broker.
type SSHAgentProxy struct {
	upstreamPath string
	broker       *Broker
	listener     net.Listener
	logf         func(string, ...any)
}

// SSHAgentConfig configures the SSH agent proxy.
type SSHAgentConfig struct {
	ListenPath   string // path for proxy socket
	UpstreamPath string // real SSH_AUTH_SOCK
	Broker       *Broker
	Logf         func(string, ...any)
}

// NewSSHAgentProxy creates and starts the SSH agent proxy.
func NewSSHAgentProxy(cfg SSHAgentConfig) (*SSHAgentProxy, error) {
	if cfg.Logf == nil {
		cfg.Logf = func(string, ...any) {}
	}
	if cfg.UpstreamPath == "" {
		cfg.UpstreamPath = os.Getenv("SSH_AUTH_SOCK")
	}
	if cfg.UpstreamPath == "" {
		return nil, fmt.Errorf("ssh-agent-proxy: no upstream SSH_AUTH_SOCK")
	}

	// Ensure directory exists
	dir := socketDir(cfg.ListenPath)
	os.MkdirAll(dir, 0700) //nolint:errcheck
	os.Remove(cfg.ListenPath) //nolint:errcheck

	ln, err := net.Listen("unix", cfg.ListenPath)
	if err != nil {
		return nil, fmt.Errorf("ssh-agent-proxy: listen %s: %w", cfg.ListenPath, err)
	}
	os.Chmod(cfg.ListenPath, 0600) //nolint:errcheck

	p := &SSHAgentProxy{
		upstreamPath: cfg.UpstreamPath,
		broker:       cfg.Broker,
		listener:     ln,
		logf:         cfg.Logf,
	}

	go p.acceptLoop()
	cfg.Logf("ssh-agent-proxy: listening on %s (upstream: %s)", cfg.ListenPath, cfg.UpstreamPath)
	return p, nil
}

// Close stops the proxy.
func (p *SSHAgentProxy) Close() {
	p.listener.Close()
}

func (p *SSHAgentProxy) acceptLoop() {
	for {
		conn, err := p.listener.Accept()
		if err != nil {
			return
		}
		go p.handleConn(conn)
	}
}

func (p *SSHAgentProxy) handleConn(conn net.Conn) {
	defer conn.Close()

	// Connect to upstream agent
	upstream, err := net.Dial("unix", p.upstreamPath)
	if err != nil {
		p.logf("ssh-agent-proxy: connect upstream: %v", err)
		return
	}
	defer upstream.Close()

	upstreamAgent := agent.NewClient(upstream)
	proxy := &approvalAgent{
		upstream: upstreamAgent,
		broker:   p.broker,
		logf:     p.logf,
	}

	// Serve the agent protocol on the client connection
	agent.ServeAgent(proxy, conn) //nolint:errcheck
}

// approvalAgent wraps an upstream agent, requiring approval for sign requests.
type approvalAgent struct {
	upstream agent.ExtendedAgent
	broker   *Broker
	logf     func(string, ...any)
}

func (a *approvalAgent) List() ([]*agent.Key, error) {
	return a.upstream.List()
}

func (a *approvalAgent) Sign(key ssh.PublicKey, data []byte) (*ssh.Signature, error) {
	return a.signWithApproval(key, data, 0)
}

func (a *approvalAgent) SignWithFlags(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) {
	return a.signWithApproval(key, data, flags)
}

func (a *approvalAgent) signWithApproval(key ssh.PublicKey, data []byte, flags agent.SignatureFlags) (*ssh.Signature, error) {
	fp := ssh.FingerprintSHA256(key)

	// Try to get key comment for readable display
	comment := ""
	if keys, err := a.upstream.List(); err == nil {
		for _, k := range keys {
			if ssh.FingerprintSHA256(k) == fp {
				comment = k.Comment
				break
			}
		}
	}

	// Check policy
	effective := a.broker.SessionPolicy("")
	if effective.Matches("", fp) {
		a.logf("ssh-agent-proxy: auto-approved sign for %s", fp)
		if flags != 0 {
			return a.upstream.SignWithFlags(key, data, flags)
		}
		return a.upstream.Sign(key, data)
	}

	// Create pending request for approval
	label := comment
	if label == "" {
		label = fp
	}

	req := &Request{
		ID:        nextRequestID(),
		Cmd:       fmt.Sprintf("ssh-sign: %s", label),
		Cwd:       "",
		SessionID: "",
		CreatedAt: timeNow(),
		resolved:  make(chan Resolution, 1),
	}

	a.broker.mu.Lock()
	a.broker.pending[req.ID] = req
	a.broker.mu.Unlock()

	a.broker.notify(req)
	a.logf("ssh-agent-proxy: pending sign approval id=%s fingerprint=%s", req.ID, fp)

	// Wait for resolution or timeout
	var res Resolution
	select {
	case res = <-req.Resolved():
	case <-timeAfter(RequestTTL):
		a.broker.mu.Lock()
		delete(a.broker.pending, req.ID)
		a.broker.mu.Unlock()
		res = ResolveTimeout
	}

	switch res {
	case ResolveApprove:
		a.logf("ssh-agent-proxy: approved sign id=%s", req.ID)
	case ResolvePersist:
		a.logf("ssh-agent-proxy: persisted sign id=%s fingerprint=%s", req.ID, fp)
		a.broker.policy.AddGlobal(Rule{SSH: fp}) //nolint:errcheck
	default:
		a.logf("ssh-agent-proxy: denied sign id=%s reason=%s", req.ID, res)
		return nil, fmt.Errorf("sign request denied")
	}

	if flags != 0 {
		return a.upstream.SignWithFlags(key, data, flags)
	}
	return a.upstream.Sign(key, data)
}

func (a *approvalAgent) Add(key agent.AddedKey) error {
	return fmt.Errorf("add not allowed via proxy")
}

func (a *approvalAgent) Remove(key ssh.PublicKey) error {
	return fmt.Errorf("remove not allowed via proxy")
}

func (a *approvalAgent) RemoveAll() error {
	return fmt.Errorf("remove all not allowed via proxy")
}

func (a *approvalAgent) Lock(passphrase []byte) error {
	return fmt.Errorf("lock not allowed via proxy")
}

func (a *approvalAgent) Unlock(passphrase []byte) error {
	return fmt.Errorf("unlock not allowed via proxy")
}

func (a *approvalAgent) Signers() ([]ssh.Signer, error) {
	return nil, fmt.Errorf("signers not supported via proxy")
}

func (a *approvalAgent) Extension(extensionType string, contents []byte) ([]byte, error) {
	return nil, agent.ErrExtensionUnsupported
}

// For testing: allow mocking time functions
var (
	timeNow   = func() time.Time { return time.Now() }
	timeAfter = func(d time.Duration) <-chan time.Time { return time.After(d) }
)