// Package sshx establishes the in-process SSH connection the SFTP browser runs
// over. haul is a GUI, so it never hands the terminal to the system ssh client;
// it opens the transport itself.
//
// There are two ways in. Connect resolves a ~/.ssh/config alias by asking the
// system ssh client via `ssh -G <alias>`, which prints the fully resolved
// effective configuration (HostName, User, Port, IdentityFile, ProxyJump,
// known-hosts files, …) — so ssh_config semantics are never re-implemented here.
// ConnectParams takes an explicit target instead, for quick connect. Both then
// authenticate with the ssh-agent, any identity files, and optionally a
// password; verify the host key against known_hosts; and dial directly or
// through a ProxyJump chain.
package sshx

import (
	"bufio"
	"errors"
	"fmt"
	"net"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"time"

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

// Params is the subset of resolved ssh_config that dial needs to connect.
type Params struct {
	Alias          string
	HostName       string
	User           string
	Port           string
	IdentityFiles  []string
	ProxyJump      string // "" or an alias/host (possibly comma-separated)
	KnownHostsFile []string
	StrictHostKey  string // "yes" | "no" | "accept-new" | "ask"
	IdentityAgent  string // socket path, or "" / "none"
	IdentitiesOnly bool   // only offer the configured identity files
}

// Options tunes a connection attempt. Every callback may block on user input:
// they are invoked from the connecting goroutine, never from the UI thread.
type Options struct {
	// Passphrase is called to unlock an encrypted identity file that the agent
	// doesn't already hold. It may return an empty slice to skip the key.
	Passphrase func(keyfile string) ([]byte, error)
	// Password, if set, enables password (and keyboard-interactive)
	// authentication as a fallback after any keys have been offered. It is the
	// quick-connect path: a host typed into the GUI has no configured identity.
	Password func(user, host string) ([]byte, error)
	// AcceptHostKey, if set, is consulted when the server's key is not in
	// known_hosts, and the key is recorded when it returns true. Without it an
	// unknown host is refused outright. A *changed* key is always refused and
	// never reaches this callback.
	AcceptHostKey func(host string, key ssh.PublicKey) (bool, error)
	// Logf, if set, receives human-readable diagnostics about which auth methods
	// were assembled — printed by the caller when a connection fails.
	Logf func(format string, args ...any)
}

func (o Options) logf(format string, args ...any) {
	if o.Logf != nil {
		o.Logf(format, args...)
	}
}

// Addr returns host:port for dialing.
func (p Params) Addr() string {
	port := p.Port
	if port == "" {
		port = "22"
	}
	return net.JoinHostPort(p.HostName, port)
}

const dialTimeout = 20 * time.Second

// Resolve runs `ssh -G alias` and parses the effective configuration.
func Resolve(alias string) (Params, error) {
	out, err := exec.Command("ssh", "-G", alias).Output()
	if err != nil {
		return Params{}, fmt.Errorf("resolving ssh config for %q (is ssh installed?): %w", alias, err)
	}
	return parseSSHG(alias, string(out)), nil
}

// parseSSHG parses `ssh -G` key/value output. Keys are case-insensitive and
// some (identityfile, userknownhostsfile) may repeat.
func parseSSHG(alias, out string) Params {
	p := Params{Alias: alias}
	sc := bufio.NewScanner(strings.NewReader(out))
	for sc.Scan() {
		line := strings.TrimSpace(sc.Text())
		if line == "" {
			continue
		}
		key, val, ok := strings.Cut(line, " ")
		if !ok {
			continue
		}
		val = strings.TrimSpace(val)
		switch strings.ToLower(key) {
		case "hostname":
			p.HostName = val
		case "user":
			p.User = val
		case "port":
			p.Port = val
		case "identityfile":
			p.IdentityFiles = append(p.IdentityFiles, expandTilde(val))
		case "proxyjump":
			if val != "none" {
				p.ProxyJump = val
			}
		case "userknownhostsfile":
			for _, f := range strings.Fields(val) {
				p.KnownHostsFile = append(p.KnownHostsFile, expandTilde(f))
			}
		case "stricthostkeychecking":
			p.StrictHostKey = val
		case "identityagent":
			p.IdentityAgent = expandTilde(val)
		case "identitiesonly":
			p.IdentitiesOnly = strings.EqualFold(val, "yes")
		}
	}
	if p.HostName == "" {
		p.HostName = alias
	}
	return p
}

// Connect opens an SSH client to the given alias, following a ProxyJump chain if
// configured. The caller is responsible for Close.
func Connect(alias string, opts Options) (*ssh.Client, error) {
	p, err := Resolve(alias)
	if err != nil {
		return nil, err
	}
	return connect(p, 0, opts)
}

// ConnectParams opens an SSH client to an explicitly specified target without
// consulting ssh_config — haul's quick-connect path, where the user typed a
// host, port and username into the GUI. Auth still uses the agent and any
// IdentityFiles set on p, falling back to Options.Password.
func ConnectParams(p Params, opts Options) (*ssh.Client, error) {
	if p.HostName == "" {
		return nil, errors.New("no host given")
	}
	if p.Alias == "" {
		p.Alias = p.HostName
	}
	return connect(p, 0, opts)
}

const maxJumps = 8

func connect(p Params, depth int, opts Options) (*ssh.Client, error) {
	if depth > maxJumps {
		return nil, errors.New("ProxyJump chain too deep")
	}
	cfg, err := clientConfig(p, opts)
	if err != nil {
		return nil, err
	}

	if p.ProxyJump == "" {
		client, err := ssh.Dial("tcp", p.Addr(), cfg)
		if err != nil {
			return nil, fmt.Errorf("connecting to %s: %w", p.Alias, err)
		}
		return client, nil
	}

	// Dial through the jump host. Only the first hop of a comma-separated list
	// is resolved here; each hop recurses, so a chain still works.
	jumpAlias := strings.Split(p.ProxyJump, ",")[0]
	jp, err := Resolve(jumpAlias)
	if err != nil {
		return nil, err
	}
	jumpClient, err := connect(jp, depth+1, opts)
	if err != nil {
		return nil, fmt.Errorf("connecting via jump host %s: %w", jumpAlias, err)
	}
	conn, err := jumpClient.Dial("tcp", p.Addr())
	if err != nil {
		jumpClient.Close()
		return nil, fmt.Errorf("dialing %s through %s: %w", p.Alias, jumpAlias, err)
	}
	ncc, chans, reqs, err := ssh.NewClientConn(conn, p.Addr(), cfg)
	if err != nil {
		jumpClient.Close()
		return nil, fmt.Errorf("ssh handshake with %s: %w", p.Alias, err)
	}
	return ssh.NewClient(ncc, chans, reqs), nil
}

func clientConfig(p Params, opts Options) (*ssh.ClientConfig, error) {
	user := p.User
	if user == "" {
		if u := os.Getenv("USER"); u != "" {
			user = u
		}
	}
	auth, err := authMethods(p, opts)
	if err != nil {
		return nil, err
	}
	hkCallback, err := hostKeyCallback(p, opts)
	if err != nil {
		return nil, err
	}
	return &ssh.ClientConfig{
		User:            user,
		Auth:            auth,
		HostKeyCallback: hkCallback,
		Timeout:         dialTimeout,
	}, nil
}

// authMethods builds the public-key auth, mirroring how the system ssh client
// chooses keys so it stays under the server's MaxAuthTries limit:
//
//   - It offers only the keys the config's IdentityFile lines point to, not
//     every key in the agent. Offering the whole agent plus all default
//     identity files can exceed MaxAuthTries (typically 6) and the server drops
//     the connection before the right key is tried.
//   - For each configured identity, it prefers a matching signer already held
//     by the agent (so passphrase-protected keys loaded via ssh-add work without
//     re-prompting), then falls back to loading the key file (prompting for a
//     passphrase if the key is encrypted and not in the agent).
//   - If no configured identity resolves, it falls back to offering the agent's
//     own keys so agent-only setups still work.
func authMethods(p Params, opts Options) ([]ssh.AuthMethod, error) {
	agentSigners := dialAgent(p, opts)
	byPub := map[string]ssh.Signer{}
	for _, s := range agentSigners {
		byPub[string(s.PublicKey().Marshal())] = s
	}
	opts.logf("ssh-agent: %d key(s) available", len(agentSigners))

	var ordered []ssh.Signer
	seen := map[string]bool{}
	add := func(s ssh.Signer) {
		k := string(s.PublicKey().Marshal())
		if !seen[k] {
			seen[k] = true
			ordered = append(ordered, s)
		}
	}

	for _, f := range p.IdentityFiles {
		// Prefer the agent's copy of this identity, matched by its public key.
		if pub, err := readPublicKey(f + ".pub"); err == nil {
			if s, ok := byPub[string(pub.Marshal())]; ok {
				opts.logf("identity %s: using agent-held key", f)
				add(s)
				continue
			}
		}
		b, err := os.ReadFile(f)
		if err != nil {
			continue // identity file not present; skip quietly
		}
		signer, err := ssh.ParsePrivateKey(b)
		if err != nil {
			var missing *ssh.PassphraseMissingError
			if errors.As(err, &missing) && opts.Passphrase != nil {
				pass, perr := opts.Passphrase(f)
				if perr == nil && len(pass) > 0 {
					if signer, err = ssh.ParsePrivateKeyWithPassphrase(b, pass); err == nil {
						opts.logf("identity %s: loaded (passphrase)", f)
						add(signer)
						continue
					}
				}
			}
			opts.logf("identity %s: skipped (%v)", f, err)
			continue
		}
		opts.logf("identity %s: loaded", f)
		add(signer)
	}

	var methods []ssh.AuthMethod
	switch {
	case len(ordered) > 0:
		// Offer exactly the configured identities — few, correct, in order.
		methods = append(methods, ssh.PublicKeys(ordered...))
	case !p.IdentitiesOnly && len(agentSigners) > 0:
		// No configured identity resolved; fall back to the agent's keys.
		opts.logf("no configured identity resolved; offering all agent keys")
		methods = append(methods, ssh.PublicKeys(agentSigners...))
	}

	// Password auth goes last so a working key is always preferred. Servers that
	// only advertise keyboard-interactive for passwords need the second method;
	// both draw on the same callback, which the GUI answers from one prompt.
	if opts.Password != nil {
		ask := func() (string, error) {
			b, err := opts.Password(p.User, p.HostName)
			return string(b), err
		}
		methods = append(methods, ssh.PasswordCallback(ask))
		methods = append(methods, ssh.KeyboardInteractive(
			func(name, instruction string, questions []string, echos []bool) ([]string, error) {
				answers := make([]string, len(questions))
				pass, err := ask()
				if err != nil {
					return nil, err
				}
				for i := range questions {
					answers[i] = pass
				}
				return answers, nil
			}))
		opts.logf("password auth available as a fallback")
	}

	if len(methods) == 0 {
		return nil, errors.New("no usable SSH key: load your key into ssh-agent (ssh-add) " +
			"or ensure the configured IdentityFile exists")
	}
	return methods, nil
}

// dialAgent connects to the ssh-agent and returns its signers, or nil if the
// agent is unreachable.
func dialAgent(p Params, opts Options) []ssh.Signer {
	sock := p.IdentityAgent
	if sock == "" || sock == "SSH_AUTH_SOCK" {
		sock = os.Getenv("SSH_AUTH_SOCK")
	}
	if sock == "" || sock == "none" {
		opts.logf("ssh-agent: not configured (SSH_AUTH_SOCK unset)")
		return nil
	}
	conn, err := net.Dial("unix", sock)
	if err != nil {
		opts.logf("ssh-agent: unreachable (%v)", err)
		return nil
	}
	signers, err := agent.NewClient(conn).Signers()
	if err != nil {
		opts.logf("ssh-agent: error listing keys (%v)", err)
		return nil
	}
	return signers
}

// readPublicKey reads and parses an OpenSSH ".pub" file.
func readPublicKey(path string) (ssh.PublicKey, error) {
	b, err := os.ReadFile(path)
	if err != nil {
		return nil, err
	}
	pub, _, _, _, err := ssh.ParseAuthorizedKey(b)
	return pub, err
}

// hostKeyCallback verifies the server against known_hosts. If strict checking
// is disabled it accepts any key. An unknown host is offered to
// Options.AcceptHostKey (and recorded if accepted) — quick connect has to be
// able to reach a host that has never been ssh'd to. A key that is present but
// *different* is always fatal: that is the attack this check exists for, and no
// prompt should make it dismissable.
func hostKeyCallback(p Params, opts Options) (ssh.HostKeyCallback, error) {
	if strings.EqualFold(p.StrictHostKey, "no") {
		return ssh.InsecureIgnoreHostKey(), nil //nolint:gosec — user opted out via config
	}

	files := p.KnownHostsFile
	if len(files) == 0 {
		if home, err := os.UserHomeDir(); err == nil {
			files = []string{filepath.Join(home, ".ssh", "known_hosts")}
		}
	}
	if len(files) == 0 {
		return nil, errors.New("cannot locate a known_hosts file")
	}
	var existing []string
	for _, f := range files {
		if _, err := os.Stat(f); err == nil {
			existing = append(existing, f)
		}
	}
	// New keys are appended to the first configured file, which is the one ssh
	// itself would write to.
	writeTo := files[0]

	if len(existing) == 0 {
		if opts.AcceptHostKey == nil {
			return nil, errors.New("no known_hosts file found; ssh to this host once first to record its key")
		}
		// Nothing recorded yet: every host is unknown, so go straight to the prompt.
		return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
			return acceptAndRecord(writeTo, hostname, key, opts)
		}, nil
	}

	known, err := knownhosts.New(existing...)
	if err != nil {
		return nil, fmt.Errorf("reading known_hosts: %w", err)
	}
	return func(hostname string, remote net.Addr, key ssh.PublicKey) error {
		err := known(hostname, remote, key)
		if err == nil {
			return nil
		}
		var ke *knownhosts.KeyError
		// A KeyError with no Want entries means the host is simply unrecorded;
		// with Want entries it means the recorded key does not match.
		if errors.As(err, &ke) && len(ke.Want) == 0 && opts.AcceptHostKey != nil {
			return acceptAndRecord(writeTo, hostname, key, opts)
		}
		return err
	}, nil
}

// acceptAndRecord asks the user about an unrecorded host key and appends it to
// known_hosts when they accept.
func acceptAndRecord(file, hostname string, key ssh.PublicKey, opts Options) error {
	ok, err := opts.AcceptHostKey(hostname, key)
	if err != nil {
		return err
	}
	if !ok {
		return fmt.Errorf("host key for %s was not accepted", hostname)
	}
	if err := appendKnownHost(file, hostname, key); err != nil {
		// Recording is best-effort: failing to write the file shouldn't sink a
		// connection the user just explicitly approved.
		opts.logf("could not record host key in %s: %v", file, err)
	}
	return nil
}

func appendKnownHost(file, hostname string, key ssh.PublicKey) error {
	if err := os.MkdirAll(filepath.Dir(file), 0o700); err != nil {
		return err
	}
	f, err := os.OpenFile(file, os.O_APPEND|os.O_CREATE|os.O_WRONLY, 0o600)
	if err != nil {
		return err
	}
	defer f.Close()
	line := knownhosts.Line([]string{knownhosts.Normalize(hostname)}, key)
	_, err = fmt.Fprintln(f, line)
	return err
}

// FingerprintSHA256 renders a host key the way ssh presents it in its
// "authenticity of host ... can't be established" prompt, so a user can compare
// the two by eye.
func FingerprintSHA256(key ssh.PublicKey) string {
	return ssh.FingerprintSHA256(key)
}

func expandTilde(path string) string {
	if path == "~" || strings.HasPrefix(path, "~/") {
		if home, err := os.UserHomeDir(); err == nil {
			if path == "~" {
				return home
			}
			return home + path[1:]
		}
	}
	return path
}