// Package tlsenum probes a remote endpoint to discover the exact set of
// SSL/TLS protocol versions and cipher suites it accepts.
//
// The Go stdlib's crypto/tls only negotiates a curated subset of modern
// suites and refuses to even offer legacy ones (RC4, 3DES, EXPORT, NULL,
// anonymous, …), so it cannot be used to *audit* what a server accepts.
// Instead we use github.com/refraction-networking/utls to craft a fully
// custom ClientHello carrying a single (version, cipher) pair and let the
// server tell us — by ServerHello or alert — whether it accepts it.
//
// Scope of the minimal version:
//   - TLS 1.0, 1.1, 1.2, 1.3 (negotiated via the SupportedVersions extension).
//   - Direct TLS only; STARTTLS upgrade is the caller's responsibility for
//     now (the existing checker package owns those dialect handlers).
//   - SSLv3 and SSLv2 are deliberately out of scope; SSLv2 has a different
//     wire format and would require either raw byte crafting or a legacy
//     OpenSSL sidecar.
package tlsenum

import (
	"context"
	"errors"
	"fmt"
	"net"
	"strconv"
	"time"

	utls "github.com/refraction-networking/utls"
)

// AllVersions is the set of protocol versions Probe knows how to offer.
var AllVersions = []uint16{
	utls.VersionTLS10,
	utls.VersionTLS11,
	utls.VersionTLS12,
	utls.VersionTLS13,
}

// VersionName returns a human-readable label for a TLS protocol version.
func VersionName(v uint16) string {
	switch v {
	case utls.VersionTLS10:
		return "TLS 1.0"
	case utls.VersionTLS11:
		return "TLS 1.1"
	case utls.VersionTLS12:
		return "TLS 1.2"
	case utls.VersionTLS13:
		return "TLS 1.3"
	default:
		return "0x" + strconv.FormatUint(uint64(v), 16)
	}
}

// ProbeResult is the outcome of a single (version, cipher) attempt.
type ProbeResult struct {
	OfferedVersion uint16
	OfferedCipher  uint16

	// Accepted is true when the server completed enough of the handshake to
	// echo back a ServerHello with our offered version and cipher. We do not
	// require a fully successful handshake (certificate verification can fail
	// for unrelated reasons); ServerHello acceptance is what we measure.
	Accepted bool

	// NegotiatedVersion / NegotiatedCipher are populated when Accepted is
	// true. They should match the offered values; if they differ, the server
	// is misbehaving (or downgrading).
	NegotiatedVersion uint16
	NegotiatedCipher  uint16

	// Err is the underlying error from the dial or handshake. For a clean
	// "server rejected this combination" outcome it will typically be a TLS
	// alert (handshake_failure, protocol_version, insufficient_security…).
	Err error
}

// ProbeOptions controls a single Probe call.
type ProbeOptions struct {
	// Timeout bounds dial + (optional) upgrade + handshake. A zero value
	// means no deadline beyond the parent context's.
	Timeout time.Duration

	// Upgrader, when non-nil, is invoked on the freshly-dialed connection
	// before the TLS ClientHello is sent. It is the injection point for
	// STARTTLS dialect handlers (SMTP, IMAP, POP3, …): the callback drives
	// the plaintext exchange that requests the upgrade and returns nil once
	// the connection is ready for tls.Client. tlsenum stays agnostic of the
	// dialect; the caller owns that knowledge.
	Upgrader func(net.Conn) error
}

// Probe attempts a TLS handshake against addr offering exactly one protocol
// version and one cipher suite. It never panics; transport / handshake errors
// are reported on the returned ProbeResult.
//
// addr must be host:port. sni is the SNI to send (pass the host if unsure).
func Probe(ctx context.Context, addr, sni string, version, cipher uint16, opts ProbeOptions) ProbeResult {
	res := ProbeResult{OfferedVersion: version, OfferedCipher: cipher}

	dialCtx := ctx
	if opts.Timeout > 0 {
		var cancel context.CancelFunc
		dialCtx, cancel = context.WithTimeout(ctx, opts.Timeout)
		defer cancel()
	}

	d := &net.Dialer{}
	raw, err := d.DialContext(dialCtx, "tcp", addr)
	if err != nil {
		res.Err = fmt.Errorf("dial: %w", err)
		return res
	}
	defer raw.Close()
	if dl, ok := dialCtx.Deadline(); ok {
		_ = raw.SetDeadline(dl)
	}

	if opts.Upgrader != nil {
		if err := opts.Upgrader(raw); err != nil {
			res.Err = fmt.Errorf("upgrade: %w", err)
			return res
		}
	}

	cfg := &utls.Config{
		ServerName:         sni,
		InsecureSkipVerify: true, // #nosec G402 -- enumeration; we only care about handshake outcome
	}
	uc := utls.UClient(raw, cfg, utls.HelloCustom)
	spec := buildSpec(version, cipher, sni)
	if err := uc.ApplyPreset(&spec); err != nil {
		res.Err = fmt.Errorf("apply-preset: %w", err)
		return res
	}

	err = uc.Handshake()
	state := uc.ConnectionState()
	if err == nil {
		res.Accepted = true
		res.NegotiatedVersion = state.Version
		res.NegotiatedCipher = state.CipherSuite
		return res
	}

	// Some servers complete ServerHello (so we know they accepted version +
	// cipher) but fail later — for example, certificate-mismatch or the
	// client failing to verify. If state has a non-zero Version/CipherSuite
	// matching what we offered, we still count it as accepted.
	if state.Version == version && state.CipherSuite == cipher && state.CipherSuite != 0 {
		res.Accepted = true
		res.NegotiatedVersion = state.Version
		res.NegotiatedCipher = state.CipherSuite
	}
	res.Err = err
	return res
}

// EnumerateOptions controls Enumerate.
type EnumerateOptions struct {
	// Timeout for each individual probe. Defaults to 5s when zero.
	ProbeTimeout time.Duration
	// Versions to try. Defaults to AllVersions when nil.
	Versions []uint16
	// Ciphers to try. Defaults to AllCiphers() when nil. The TLS13 flag is
	// honored: TLS 1.3 ciphers are only offered with TLS 1.3 probes, and
	// vice-versa.
	Ciphers []CipherSuite
	// Upgrader, when non-nil, is forwarded to every sub-probe (see
	// ProbeOptions.Upgrader). It is invoked on a freshly-dialed connection
	// before each ClientHello, so STARTTLS dialect handlers run once per
	// probe, not once for the whole sweep.
	Upgrader func(net.Conn) error
}

// EnumerationResult is the aggregate outcome of an enumeration sweep.
type EnumerationResult struct {
	// SupportedVersions lists protocol versions for which at least one
	// cipher was accepted.
	SupportedVersions []uint16
	// CiphersByVersion lists, per accepted version, the cipher suites the
	// server agreed to negotiate.
	CiphersByVersion map[uint16][]CipherSuite
}

// Enumerate sweeps a (version × cipher) matrix against addr and returns what
// the server actually accepts. Probes are performed sequentially; concurrency
// can be added later but tends to upset some middleboxes when probing too
// hard.
func Enumerate(ctx context.Context, addr, sni string, opts EnumerateOptions) (EnumerationResult, error) {
	if opts.ProbeTimeout == 0 {
		opts.ProbeTimeout = 5 * time.Second
	}
	versions := opts.Versions
	if versions == nil {
		versions = AllVersions
	}
	ciphers := opts.Ciphers
	if ciphers == nil {
		ciphers = AllCiphers()
	}

	out := EnumerationResult{
		CiphersByVersion: make(map[uint16][]CipherSuite),
	}
	seenVersion := make(map[uint16]bool)

	for _, v := range versions {
		isTLS13 := v == utls.VersionTLS13
		for _, c := range ciphers {
			if c.TLS13 != isTLS13 {
				continue
			}
			if err := ctx.Err(); err != nil {
				return out, err
			}
			r := Probe(ctx, addr, sni, v, c.ID, ProbeOptions{
				Timeout:  opts.ProbeTimeout,
				Upgrader: opts.Upgrader,
			})
			if !r.Accepted {
				continue
			}
			out.CiphersByVersion[v] = append(out.CiphersByVersion[v], c)
			if !seenVersion[v] {
				seenVersion[v] = true
				out.SupportedVersions = append(out.SupportedVersions, v)
			}
		}
	}
	return out, nil
}

// buildSpec assembles a ClientHelloSpec offering exactly one cipher and one
// protocol version. For TLS 1.3 the legacy version field stays at TLS 1.2 and
// the real version is signalled through the SupportedVersions extension, per
// RFC 8446 §4.1.2 / §4.2.1.
func buildSpec(version, cipher uint16, sni string) utls.ClientHelloSpec {
	tlsVersMin := version
	tlsVersMax := version
	if version == utls.VersionTLS13 {
		// utls inspects TLSVersMax to decide whether to drive TLS 1.3
		// machinery; the on-the-wire legacy_version stays TLS 1.2.
		tlsVersMin = utls.VersionTLS12
	}

	exts := []utls.TLSExtension{
		&utls.SNIExtension{ServerName: sni},
		&utls.SupportedCurvesExtension{Curves: []utls.CurveID{
			utls.X25519, utls.CurveP256, utls.CurveP384, utls.CurveP521,
		}},
		&utls.SupportedPointsExtension{SupportedPoints: []byte{0}}, // uncompressed
		&utls.SignatureAlgorithmsExtension{SupportedSignatureAlgorithms: []utls.SignatureScheme{
			utls.ECDSAWithP256AndSHA256, utls.ECDSAWithP384AndSHA384, utls.ECDSAWithP521AndSHA512,
			utls.PSSWithSHA256, utls.PSSWithSHA384, utls.PSSWithSHA512,
			utls.PKCS1WithSHA256, utls.PKCS1WithSHA384, utls.PKCS1WithSHA512,
			utls.PKCS1WithSHA1, utls.ECDSAWithSHA1,
		}},
		&utls.RenegotiationInfoExtension{Renegotiation: utls.RenegotiateOnceAsClient},
	}

	if version == utls.VersionTLS13 {
		exts = append(exts,
			&utls.SupportedVersionsExtension{Versions: []uint16{utls.VersionTLS13}},
			&utls.KeyShareExtension{KeyShares: []utls.KeyShare{
				{Group: utls.X25519},
			}},
			&utls.PSKKeyExchangeModesExtension{Modes: []uint8{utls.PskModeDHE}},
		)
	}

	return utls.ClientHelloSpec{
		TLSVersMin:         tlsVersMin,
		TLSVersMax:         tlsVersMax,
		CipherSuites:       []uint16{cipher},
		CompressionMethods: []byte{0}, // null
		Extensions:         exts,
	}
}

// ErrNoVersions is returned when an enumeration request asks for an empty set
// of versions or ciphers.
var ErrNoVersions = errors.New("tlsenum: no versions or ciphers to probe")