sshkeeper/internal/ssh/planner.go

package ssh

import (
	"fmt"
	"os"
	"sort"
	"strings"

	"github.com/mirivlad/sshkeeper/internal/model"
)

// ProfileResolver resolves a stable sshkeeper server ID for route planning.
type ProfileResolver func(serverID int64) (*model.Server, error)

type PlannedHop struct {
	Server *model.Server
	Raw    string
}

type ConnectionPlan struct {
	Target       *model.Server
	Hops         []PlannedHop
	UsesProfiles bool
}

func PlanConnection(target *model.Server, resolve ProfileResolver) (*ConnectionPlan, error) {
	if target == nil {
		return nil, fmt.Errorf("target server is required")
	}
	plan := &ConnectionPlan{Target: target}
	stack := map[int64]bool{}
	if target.ID > 0 {
		stack[target.ID] = true
	}
	seenProfiles := map[int64]bool{}
	seenRaw := map[string]bool{}
	hops, err := flattenRoute(target, resolve, stack, seenProfiles, seenRaw)
	if err != nil {
		return nil, err
	}
	plan.Hops = hops
	for _, hop := range hops {
		if hop.Server != nil {
			plan.UsesProfiles = true
			break
		}
	}
	return plan, nil
}

func flattenRoute(owner *model.Server, resolve ProfileResolver, stack, seenProfiles map[int64]bool, seenRaw map[string]bool) ([]PlannedHop, error) {
	var result []PlannedHop
	for _, hop := range owner.Route.Hops {
		if hop.Profile() {
			if hop.ServerID <= 0 {
				return nil, fmt.Errorf("route profile %q has no stable ID; edit and re-save the route", hop.Alias)
			}
			if resolve == nil {
				return nil, fmt.Errorf("route profile %s requires sshkeeper profile resolution", hop.DisplayName())
			}
			if stack[hop.ServerID] {
				return nil, fmt.Errorf("route cycle detected at %s", hop.DisplayName())
			}
			server, err := resolve(hop.ServerID)
			if err != nil {
				return nil, fmt.Errorf("resolve route profile %s: %w", hop.DisplayName(), err)
			}
			if server.AuthMethod == model.AuthPassword || server.AuthMethod == model.AuthKeyPassphrase {
				return nil, fmt.Errorf("jump profile %s uses %s authentication; password/passphrase jump profiles are not supported by the current OpenSSH vault flow", server.Alias, server.AuthMethod)
			}
			stack[server.ID] = true
			nested, err := flattenRoute(server, resolve, stack, seenProfiles, seenRaw)
			delete(stack, server.ID)
			if err != nil {
				return nil, err
			}
			result = append(result, nested...)
			if seenProfiles[server.ID] {
				return nil, fmt.Errorf("route resolves to duplicate profile hop %s", server.Alias)
			}
			seenProfiles[server.ID] = true
			result = append(result, PlannedHop{Server: server})
			continue
		}
		raw := strings.TrimSpace(hop.Raw)
		if raw == "" {
			return nil, fmt.Errorf("route contains an empty raw hop")
		}
		if seenRaw[raw] {
			return nil, fmt.Errorf("route resolves to duplicate raw hop %q", raw)
		}
		seenRaw[raw] = true
		result = append(result, PlannedHop{Raw: raw})
	}
	return result, nil
}

func syntheticProfileHost(id int64) string {
	return fmt.Sprintf("sshkeeper-profile-%d", id)
}

func syntheticTargetHost(target *model.Server) string {
	if target.ID > 0 {
		return fmt.Sprintf("sshkeeper-target-%d", target.ID)
	}
	return "sshkeeper-target-unsaved"
}

func appendHostBlock(b *strings.Builder, hostAlias string, server *model.Server) {
	fmt.Fprintf(b, "Host %s\n", hostAlias)
	fmt.Fprintf(b, "    HostName %s\n", server.Host)
	port := server.Port
	if port == 0 {
		port = 22
	}
	fmt.Fprintf(b, "    Port %d\n", port)
	if server.User != "" {
		fmt.Fprintf(b, "    User %s\n", server.User)
	}
	if server.IdentityFile != "" && server.AuthMethod != model.AuthPassword && server.AuthMethod != model.AuthAgent {
		fmt.Fprintf(b, "    IdentityFile %s\n", server.IdentityFile)
	}
	fmt.Fprintln(b, "    StrictHostKeyChecking accept-new")
	fmt.Fprintln(b)
}

// OpenSSHConfig renders the deterministic temporary config used when a route
// references sshkeeper profiles. The user's normal config is included first so
// raw OpenSSH jump targets keep their existing configuration.
func (p *ConnectionPlan) OpenSSHConfig() string {
	var b strings.Builder
	b.WriteString("# Temporary config generated by sshkeeper\n")
	b.WriteString("Include ~/.ssh/config\n\n")
	profiles := map[int64]*model.Server{}
	for _, hop := range p.Hops {
		if hop.Server != nil {
			profiles[hop.Server.ID] = hop.Server
		}
	}
	ids := make([]int64, 0, len(profiles))
	for id := range profiles {
		ids = append(ids, id)
	}
	sort.Slice(ids, func(i, j int) bool { return ids[i] < ids[j] })
	for _, id := range ids {
		appendHostBlock(&b, syntheticProfileHost(id), profiles[id])
	}
	appendHostBlock(&b, syntheticTargetHost(p.Target), p.Target)
	if len(p.Hops) > 0 {
		parts := make([]string, 0, len(p.Hops))
		for _, hop := range p.Hops {
			if hop.Server != nil {
				parts = append(parts, syntheticProfileHost(hop.Server.ID))
			} else {
				parts = append(parts, hop.Raw)
			}
		}
		// OpenSSH uses the first value it obtains. Append a target-specific
		// stanza after the generic one with ProxyJump before any competing rule.
		fmt.Fprintf(&b, "Host %s\n", syntheticTargetHost(p.Target))
		fmt.Fprintf(&b, "    ProxyJump %s\n\n", strings.Join(parts, ","))
	}
	return b.String()
}

type PreparedInvocation struct {
	Args       []string
	ConfigPath string
}

func (p *PreparedInvocation) Cleanup() {
	if p != nil && p.ConfigPath != "" {
		_ = os.Remove(p.ConfigPath)
		p.ConfigPath = ""
	}
}

func enabledForwards(forwards []*model.Forward) []*model.Forward {
	result := make([]*model.Forward, 0, len(forwards))
	for _, forward := range forwards {
		if forward != nil && forward.Enabled {
			result = append(result, forward)
		}
	}
	return result
}

func PrepareSSHInvocation(server *model.Server, forwards []*model.Forward, forwardOnly bool, resolve ProfileResolver) (*PreparedInvocation, error) {
	plan, err := PlanConnection(server, resolve)
	if err != nil {
		return nil, err
	}
	active := enabledForwards(forwards)
	if !plan.UsesProfiles {
		return &PreparedInvocation{Args: BuildSSHArgs(server, active, forwardOnly)}, nil
	}
	file, err := os.CreateTemp("", "sshkeeper-*.conf")
	if err != nil {
		return nil, fmt.Errorf("create temporary ssh config: %w", err)
	}
	path := file.Name()
	if err := file.Chmod(0600); err != nil {
		file.Close()
		os.Remove(path)
		return nil, err
	}
	if _, err := file.WriteString(plan.OpenSSHConfig()); err != nil {
		file.Close()
		os.Remove(path)
		return nil, err
	}
	if err := file.Close(); err != nil {
		os.Remove(path)
		return nil, err
	}
	args := []string{"-F", path}
	if len(active) > 0 {
		args = append(args, BuildForwardArgs(active, true)...)
	}
	if forwardOnly {
		args = append(args, "-N")
	}
	args = append(args, syntheticTargetHost(server))
	return &PreparedInvocation{Args: args, ConfigPath: path}, nil
}