// This file is part of the happyDomain (R) project.
// Copyright (c) 2020-2026 happyDomain
// Authors: Pierre-Olivier Mercier, et al.
//
// This program is offered under a commercial and under the AGPL license.
// For commercial licensing, contact us at <contact@happydomain.org>.
//
// For AGPL licensing:
// This program is free software: you can redistribute it and/or modify
// it under the terms of the GNU Affero General Public License as published by
// the Free Software Foundation, either version 3 of the License, or
// (at your option) any later version.
//
// This program is distributed in the hope that it will be useful,
// but WITHOUT ANY WARRANTY; without even the implied warranty of
// MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
// GNU Affero General Public License for more details.
//
// You should have received a copy of the GNU Affero General Public License
// along with this program.  If not, see <https://www.gnu.org/licenses/>.

package checker

import (
	"bufio"
	"context"
	"crypto/sha1"
	"crypto/sha256"
	"encoding/hex"
	"errors"
	"net"
	"strconv"
	"strings"
	"time"

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

// probeEndpoint runs the full probe flow on a single (host, ip, port)
// triple. It never returns a Go error: every failure mode is recorded
// as a raw field on SSHProbe (Stage + Error). Severity / pass/fail
// classification is performed later by CheckRule.Evaluate, never here.
func probeEndpoint(ctx context.Context, host, ip string, port uint16, timeout time.Duration, includeAuthProbe bool, sshfp SSHFPSummary) SSHProbe {
	start := time.Now()
	addr := net.JoinHostPort(ip, strconv.Itoa(int(port)))
	p := SSHProbe{
		Host:    host,
		Port:    port,
		Address: addr,
		IP:      ip,
		IsIPv6:  strings.Contains(ip, ":"),
		Stage:   "dial",
	}

	dialCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()

	d := &net.Dialer{}
	conn, err := d.DialContext(dialCtx, "tcp", addr)
	if err != nil {
		p.Error = "dial: " + err.Error()
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}
	defer conn.Close()

	p.TCPConnected = true
	p.Stage = "banner"
	if deadline, ok := dialCtx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}

	// Phase 1: protocol banner exchange.
	br := bufio.NewReader(conn)
	banner, err := readBanner(br)
	if err != nil {
		p.Error = "banner: " + err.Error()
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}
	p.Banner = banner
	p.ProtoVer, p.SoftVer, p.Vendor = parseBanner(banner)
	p.Stage = "banner_write"

	if err := writeBanner(conn); err != nil {
		p.Error = "write-banner: " + err.Error()
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}

	// Phase 2: exchange KEXINIT.
	p.Stage = "kexinit_read"
	srvPayload, err := readPacket(br)
	if err != nil {
		p.Error = "kexinit-read: " + err.Error()
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}
	p.Stage = "kexinit_parse"
	srvKex, err := parseKexInit(srvPayload)
	if err != nil {
		p.Error = "kexinit-parse: " + err.Error()
		p.ElapsedMS = time.Since(start).Milliseconds()
		return p
	}

	p.KEX = srvKex.KexAlgorithms
	p.HostKey = srvKex.ServerHostKeyAlgorithms
	p.CiphersC2S = srvKex.EncryptionAlgorithmsClientToSvr
	p.CiphersS2C = srvKex.EncryptionAlgorithmsSvrToClient
	p.MACsC2S = srvKex.MACAlgorithmsClientToSvr
	p.MACsS2C = srvKex.MACAlgorithmsSvrToClient
	p.CompC2S = srvKex.CompressionAlgorithmsClientToSv
	p.CompS2C = srvKex.CompressionAlgorithmsSvrToClt
	p.Stage = "kexinit_ok"

	// We intentionally don't proceed with KEX here: algorithm posture is
	// already captured. Closing now is friendlier than triggering a full
	// exchange that might never terminate.
	_ = conn.Close()

	// We hand off to Go's ssh package for the full handshake. HostKeyCallback lets us
	// capture each presented key without reimplementing DH/curve25519/kyber ourselves.
	p.HostKeys = probeHostKeys(ctx, addr, host, srvKex.ServerHostKeyAlgorithms, timeout)
	for i := range p.HostKeys {
		p.HostKeys[i].applySSHFP(sshfp)
	}
	if len(p.HostKeys) > 0 {
		p.Stage = "handshake_ok"
	}

	if includeAuthProbe {
		p.AuthProbeAttempted = true
		methods, err := probeAuthMethods(ctx, addr, timeout)
		if err == nil {
			p.AuthMethods = methods
			for _, m := range methods {
				switch m {
				case "password":
					p.PasswordAuth = true
				case "keyboard-interactive":
					p.KeyboardInteractive = true
				case "publickey":
					p.PublicKeyAuth = true
				}
			}
		}
	}

	p.ElapsedMS = time.Since(start).Milliseconds()
	return p
}

// Most deployments expose at most two or three key families (ed25519, rsa, ecdsa),
// so connecting once per family stays cheap.
func probeHostKeys(ctx context.Context, addr, host string, algos []string, timeout time.Duration) []HostKeyInfo {
	wantFamilies := pickHostKeyFamilies(algos)

	seen := map[string]bool{} // by sha256 hex, dedupe across families
	var out []HostKeyInfo

	for _, algo := range wantFamilies {
		key, err := fetchHostKey(ctx, addr, host, algo, timeout)
		if err != nil || key == nil {
			continue
		}
		info := describeHostKey(key)
		if seen[info.SHA256] {
			continue
		}
		seen[info.SHA256] = true
		out = append(out, info)
	}

	return out
}

// rsa-sha2-512 and rsa-sha2-256 both return the same RSA key, so we collapse by family.
func pickHostKeyFamilies(algos []string) []string {
	var out []string
	families := map[string]bool{}
	add := func(family, algo string) {
		if families[family] {
			return
		}
		families[family] = true
		out = append(out, algo)
	}
	for _, a := range algos {
		switch a {
		case "ssh-ed25519", "ssh-ed25519-cert-v01@openssh.com":
			add("ed25519", "ssh-ed25519")
		case "rsa-sha2-512", "rsa-sha2-256", "ssh-rsa":
			add("rsa", a)
		case "ecdsa-sha2-nistp256", "ecdsa-sha2-nistp384", "ecdsa-sha2-nistp521":
			add("ecdsa", a)
		case "ssh-dss":
			add("dsa", "ssh-dss")
		}
	}
	return out
}

// Offering no auth methods aborts the handshake at the auth step, which is enough
// to capture the host key without completing a full session.
func fetchHostKey(ctx context.Context, addr, host, algo string, timeout time.Duration) (ssh.PublicKey, error) {
	var captured ssh.PublicKey
	cfg := &ssh.ClientConfig{
		User: "happydomain-checker",
		Auth: nil,
		HostKeyCallback: func(_ string, _ net.Addr, key ssh.PublicKey) error {
			captured = key
			return nil
		},
		HostKeyAlgorithms: []string{algo},
		Timeout:           timeout,
		ClientVersion:     sshClientBanner,
	}

	dialCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	d := &net.Dialer{}
	conn, err := d.DialContext(dialCtx, "tcp", addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	if deadline, ok := dialCtx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}

	_, _, _, err = ssh.NewClientConn(conn, host, cfg)
	if err != nil && captured == nil {
		return nil, err
	}
	return captured, nil
}

// probeAuthMethods opens a fresh connection, completes the KEX, and
// then sends a "none" authentication request (RFC 4252 §5.2). The
// server's failure response carries the list of methods it would
// actually accept: exactly what we need.
func probeAuthMethods(ctx context.Context, addr string, timeout time.Duration) ([]string, error) {
	cfg := &ssh.ClientConfig{
		User:            "happydomain-checker",
		Auth:            []ssh.AuthMethod{}, // forces a "none" attempt
		HostKeyCallback: ssh.InsecureIgnoreHostKey(),
		Timeout:         timeout,
		ClientVersion:   sshClientBanner,
	}

	dialCtx, cancel := context.WithTimeout(ctx, timeout)
	defer cancel()
	d := &net.Dialer{}
	conn, err := d.DialContext(dialCtx, "tcp", addr)
	if err != nil {
		return nil, err
	}
	defer conn.Close()
	if deadline, ok := dialCtx.Deadline(); ok {
		_ = conn.SetDeadline(deadline)
	}

	_, _, _, err = ssh.NewClientConn(conn, addr, cfg)
	if err == nil {
		// A server that lets us through "none" is unusual but possible
		// (anonymous SSH for git-serve-style deployments); report that
		// upstream by returning an empty list.
		return nil, nil
	}
	return extractMethodsFromAuthError(err), nil
}

// x/crypto/ssh does not expose offered auth methods via a typed accessor; string
// parsing is the officially documented path.
func extractMethodsFromAuthError(err error) []string {
	if err == nil {
		return nil
	}
	msg := err.Error()
	start := strings.Index(msg, "attempted methods [")
	if start < 0 {
		return nil
	}
	start += len("attempted methods [")
	end := strings.Index(msg[start:], "]")
	if end < 0 {
		return nil
	}
	raw := strings.Fields(msg[start : start+end])
	var out []string
	for _, m := range raw {
		if m == "none" {
			continue
		}
		out = append(out, m)
	}
	return out
}

func describeHostKey(key ssh.PublicKey) HostKeyInfo {
	marshaled := key.Marshal()
	sha2 := sha256.Sum256(marshaled)
	sha1sum := sha1.Sum(marshaled)
	info := HostKeyInfo{
		Type:   key.Type(),
		SHA256: hex.EncodeToString(sha2[:]),
		SHA1:   hex.EncodeToString(sha1sum[:]),
	}
	info.SSHFPAlgo = sshfpAlgoForKeyType(info.Type)
	info.Bits = keyBits(key)
	return info
}

// keyBits returns a key-family-specific size estimate. It is advisory:
// we only use it in the report, and a server that ships an RSA key
// smaller than 2048 bits is the sort of red flag we want to show.
func keyBits(key ssh.PublicKey) int {
	switch k := key.(type) {
	case ssh.CryptoPublicKey:
		type bitSizer interface{ Size() int }
		switch p := k.CryptoPublicKey().(type) {
		case bitSizer:
			return p.Size() * 8
		default:
			_ = p
		}
	}
	return 0
}

// sshfpAlgoForKeyType maps an SSH host-key type string to the SSHFP
// algorithm number defined in RFC 4255 / RFC 6594 / RFC 7479.
func sshfpAlgoForKeyType(t string) uint8 {
	switch t {
	case "ssh-rsa", "rsa-sha2-256", "rsa-sha2-512":
		return 1
	case "ssh-dss":
		return 2
	case "ecdsa-sha2-nistp256", "ecdsa-sha2-nistp384", "ecdsa-sha2-nistp521":
		return 3
	case "ssh-ed25519":
		return 4
	}
	return 0
}

// parseBanner splits an "SSH-2.0-OpenSSH_9.3p1 Debian-1" banner into
// (protocolVersion, softwareVersion, vendorComment). The grammar is
// RFC 4253 §4.2: "SSH-<protoversion>-<softwareversion> <comments>".
func parseBanner(b string) (proto, soft, vendor string) {
	// SSH- prefix is guaranteed by readBanner.
	rest := strings.TrimPrefix(b, "SSH-")
	dash := strings.IndexByte(rest, '-')
	if dash < 0 {
		return rest, "", ""
	}
	proto = rest[:dash]
	rest = rest[dash+1:]
	if sp := strings.IndexByte(rest, ' '); sp >= 0 {
		soft = rest[:sp]
		vendor = strings.TrimSpace(rest[sp+1:])
	} else {
		soft = rest
	}
	return
}

// applySSHFP fills in the SSHFPMatchSHA* flags based on the declared
// SSHFP records for this key's algorithm family. These are raw
// observations (the record matched this key fingerprint); any
// severity verdict about coverage lives in the SSHFP rule.
func (h *HostKeyInfo) applySSHFP(s SSHFPSummary) {
	for _, rr := range s.Records {
		if rr.Algorithm != h.SSHFPAlgo {
			continue
		}
		want := strings.ToLower(rr.Fingerprint)
		switch rr.Type {
		case 1:
			if want == h.SHA1 {
				h.SSHFPMatchSHA1 = true
			}
		case 2:
			if want == h.SHA256 {
				h.SSHFPMatchSHA256 = true
			}
		}
	}
}

// errNoHostKey is returned by fetchHostKey when the callback never
// fired (e.g. transport-level error before the host key was received).
// Currently only used internally for readability.
var errNoHostKey = errors.New("no host key observed")