//go:build linux || freebsd
// +build linux freebsd

package sftp

import (
	"crypto/sha256"
	"encoding/base64"
	"fmt"
	"net"
	"os"
	"path/filepath"
	"strings"

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

// hostKeyCallback builds the ssh.HostKeyCallback used for dialing.
//
// Precedence:
//
//  1. pinnedHostKey — when set, only a host whose key matches the
//     given fingerprint is accepted. known_hosts and --ssh-insecure
//     are both ignored because pinning is a stricter guarantee.
//  2. insecure --ssh-insecure skips verification entirely.
//  3. known_hosts file (default ~/.ssh/known_hosts).
func hostKeyCallback(knownHostsPath string, insecure bool, pinnedHostKey string) (ssh.HostKeyCallback, error) {
	pinned := normalizePinnedHostKey(pinnedHostKey)
	if pinned != "" {
		return func(_ string, _ net.Addr, key ssh.PublicKey) error {
			actual := sshFingerprintSHA256(key)
			if actual == pinned {
				return nil
			}
			return fmt.Errorf("host key pinning failed: server key %s does not match pinned %s", actual, pinned)
		}, nil
	}

	if insecure {
		return ssh.InsecureIgnoreHostKey(), nil
	}

	if knownHostsPath == "" {
		home, err := os.UserHomeDir()
		if err != nil {
			return nil, fmt.Errorf("known_hosts: %w", err)
		}
		knownHostsPath = filepath.Join(home, ".ssh", "known_hosts")
	}

	if _, err := os.Stat(knownHostsPath); err != nil {
		return nil, fmt.Errorf("known_hosts file not found: %s (use --ssh-insecure to skip host key check)", knownHostsPath)
	}

	return knownhosts.New(knownHostsPath)
}

// sshFingerprintSHA256 returns the OpenSSH-style SHA-256 fingerprint
// of an SSH public key, e.g. "SHA256:abc123...". This matches the
// output of `ssh-keygen -lf <keyfile>`.
func sshFingerprintSHA256(key ssh.PublicKey) string {
	hash := sha256.Sum256(key.Marshal())
	return "SHA256:" + base64.StdEncoding.EncodeToString(hash[:])
}

// normalizePinnedHostKey accepts either the OpenSSH fingerprint form
// ("SHA256:<base64>") or a raw lowercase-hex SHA-256 of the key wire
// format, and returns the canonical OpenSSH fingerprint form. An empty
// input returns an empty string.
func normalizePinnedHostKey(s string) string {
	s = strings.TrimSpace(s)
	if s == "" {
		return ""
	}
	if strings.HasPrefix(s, "SHA256:") {
		return s
	}
	// Raw hex form — convert to base64 to match the OpenSSH fingerprint.
	if decoded := tryHexDecode(s); decoded != nil {
		return "SHA256:" + base64.StdEncoding.EncodeToString(decoded)
	}
	return s
}

// tryHexDecode decodes a lowercase-or-uppercase hex string. Returns nil
// on any error or odd length.
func tryHexDecode(s string) []byte {
	if len(s)%2 != 0 {
		return nil
	}
	out := make([]byte, len(s)/2)
	for i := 0; i < len(s); i += 2 {
		hi, ok1 := hexNibble(s[i])
		lo, ok2 := hexNibble(s[i+1])
		if !ok1 || !ok2 {
			return nil
		}
		out[i/2] = hi<<4 | lo
	}
	return out
}

func hexNibble(c byte) (byte, bool) {
	switch {
	case c >= '0' && c <= '9':
		return c - '0', true
	case c >= 'a' && c <= 'f':
		return c - 'a' + 10, true
	case c >= 'A' && c <= 'F':
		return c - 'A' + 10, true
	}
	return 0, false
}