// Package discover finds services listening on this machine and, only when asked to, on configured
// private subnets, so Hammerhead can offer them as suggested targets. It is opt-in and deliberately
// gentle: connections are plain TCP connects, probes are rate limited, only private address ranges
// of a bounded size are ever scanned, and nothing is sent beyond one HTTP HEAD (or a TLS handshake)
// to a port that stayed silent.
package discover

import (
	"bufio"
	"context"
	"crypto/tls"
	"fmt"
	"net"
	"net/netip"
	"os"
	"runtime"
	"sort"
	"strconv"
	"strings"
	"sync"
	"time"

	"golang.org/x/time/rate"

	"git.veillax.com/Veillax/reef/model"
)

const (
	// MaxSubnetAddrs caps how many addresses one subnet may cover (a /22).
	MaxSubnetAddrs = 1024
	// MaxResults caps what is reported, so a noisy network cannot bloat every report.
	MaxResults = 500

	DefaultRate     = 50 // new connections per second
	DefaultInterval = 15 * time.Minute
)

// DefaultPorts are the TCP ports a subnet scan looks at when none are configured.
var DefaultPorts = []int{22, 80, 443, 445, 1883, 3000, 3306, 5000, 5432, 6379, 8000, 8080, 8096, 8123, 8443, 8888, 9000, 9090, 32400}

// wellKnown names a likely service from the port alone, used when a service says nothing.
var wellKnown = map[int]string{
	21: "ftp", 22: "ssh", 25: "smtp", 80: "http", 110: "pop3", 143: "imap", 443: "https", 445: "smb",
	587: "smtp", 1883: "mqtt", 3306: "mysql", 3389: "rdp", 5432: "postgres", 5900: "vnc", 6379: "redis",
	8080: "http", 8443: "https", 27017: "mongodb",
}

type Config struct {
	Local    bool           // list this machine's listening sockets
	Subnets  []netip.Prefix // private subnets to scan; empty = no scanning
	Ports    []int          // ports to try on each scanned address (default DefaultPorts)
	Rate     float64        // probes per second across everything (default DefaultRate)
	Interval time.Duration  // between sweeps (default DefaultInterval)
	Exclude  []int          // local listening ports to leave out (Bamboo's own relay, for instance)

	// the next fields exist for tests
	dial      func(ctx context.Context, addr string) (net.Conn, error)
	listening func() ([]netip.AddrPort, error)
	timeout   time.Duration // per probe (default 1.5s)
}

// ParseSubnets turns "192.168.1.0/24, 10.0.0.0/24" into prefixes and refuses anything that is not
// a private, link-local or loopback range of at most MaxSubnetAddrs addresses.
func ParseSubnets(s string) ([]netip.Prefix, error) {
	var out []netip.Prefix
	for _, f := range strings.FieldsFunc(s, func(r rune) bool { return r == ',' || r == ' ' || r == ';' }) {
		p, err := netip.ParsePrefix(f)
		if err != nil {
			return nil, fmt.Errorf("bad subnet %q: %v", f, err)
		}
		p = p.Masked()
		if err := checkSubnet(p); err != nil {
			return nil, err
		}
		out = append(out, p)
	}
	return out, nil
}

func checkSubnet(p netip.Prefix) error {
	a := p.Addr()
	if !(a.IsPrivate() || a.IsLoopback() || a.IsLinkLocalUnicast()) {
		return fmt.Errorf("subnet %s is not a private range; only private networks are scanned", p)
	}
	bits := a.BitLen() - p.Bits()
	if bits > 10 {
		return fmt.Errorf("subnet %s is too large (at most %d addresses, a /%d for IPv4)", p, MaxSubnetAddrs, 32-10)
	}
	return nil
}

// ParsePorts turns "80,443,8000-8010" into a sorted, de-duplicated list.
func ParsePorts(s string) ([]int, error) {
	seen := map[int]bool{}
	var out []int
	for _, f := range strings.FieldsFunc(s, func(r rune) bool { return r == ',' || r == ' ' }) {
		lo, hi := f, f
		if i := strings.IndexByte(f, '-'); i > 0 {
			lo, hi = f[:i], f[i+1:]
		}
		a, err1 := strconv.Atoi(lo)
		b, err2 := strconv.Atoi(hi)
		if err1 != nil || err2 != nil || a < 1 || b > 65535 || a > b || b-a > 1000 {
			return nil, fmt.Errorf("bad port or range %q", f)
		}
		for p := a; p <= b; p++ {
			if !seen[p] {
				seen[p] = true
				out = append(out, p)
			}
		}
	}
	sort.Ints(out)
	return out, nil
}

// Enabled reports whether this configuration will look at anything.
func (c Config) Enabled() bool { return c.Local || len(c.Subnets) > 0 }

// Discoverer runs sweeps and keeps the latest results.
type Discoverer struct {
	cfg Config
	lim *rate.Limiter

	mu      sync.Mutex
	results []model.DiscoveredService
}

func New(cfg Config) (*Discoverer, error) {
	for _, p := range cfg.Subnets {
		if err := checkSubnet(p); err != nil { // a Config built in code gets the same checks as flags
			return nil, err
		}
	}
	if len(cfg.Ports) == 0 {
		cfg.Ports = DefaultPorts
	}
	if cfg.Rate <= 0 {
		cfg.Rate = DefaultRate
	}
	if cfg.Interval <= 0 {
		cfg.Interval = DefaultInterval
	}
	if cfg.timeout <= 0 {
		cfg.timeout = 1500 * time.Millisecond
	}
	if cfg.dial == nil {
		cfg.dial = func(ctx context.Context, addr string) (net.Conn, error) {
			var d net.Dialer
			return d.DialContext(ctx, "tcp", addr)
		}
	}
	if cfg.listening == nil {
		cfg.listening = listeningSockets
	}
	return &Discoverer{cfg: cfg, lim: rate.NewLimiter(rate.Limit(cfg.Rate), 1)}, nil
}

// Results returns the latest sweep's findings (a copy).
func (d *Discoverer) Results() []model.DiscoveredService {
	d.mu.Lock()
	defer d.mu.Unlock()
	return append([]model.DiscoveredService(nil), d.results...)
}

// Run sweeps now and then every Interval until ctx ends. logf may be nil.
func (d *Discoverer) Run(ctx context.Context, logf func(string, ...any)) {
	for {
		res := d.Sweep(ctx)
		if ctx.Err() != nil {
			return
		}
		d.mu.Lock()
		d.results = res
		d.mu.Unlock()
		if logf != nil {
			logf("discovery: %d service(s) found", len(res))
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(d.cfg.Interval):
		}
	}
}

// Sweep does one pass: listening sockets first, then each configured subnet.
func (d *Discoverer) Sweep(ctx context.Context) []model.DiscoveredService {
	var out []model.DiscoveredService
	seen := map[netip.AddrPort]bool{}
	add := func(s model.DiscoveredService) {
		ap, err := netip.ParseAddrPort(net.JoinHostPort(s.Host, strconv.Itoa(s.Port)))
		if err == nil && !seen[ap] && len(out) < MaxResults {
			seen[ap] = true
			out = append(out, s)
		}
	}
	if d.cfg.Local {
		eps, err := d.cfg.listening()
		if err != nil {
			eps = nil
		}
		for _, ap := range eps {
			if ctx.Err() != nil {
				break
			}
			if d.excluded(int(ap.Port())) {
				continue
			}
			host := ap.Addr()
			if host.IsUnspecified() {
				host = netip.AddrFrom4([4]byte{127, 0, 0, 1})
			}
			if s, ok := d.probe(ctx, netip.AddrPortFrom(host, ap.Port()), "local"); ok {
				add(s)
			}
		}
	}
	for _, sub := range d.cfg.Subnets {
		for a := sub.Addr(); sub.Contains(a) && ctx.Err() == nil; a = a.Next() {
			if !a.Is4() || !a.IsValid() {
				continue
			}
			if sub.Bits() < 31 && (a == sub.Addr() || a == lastAddr(sub)) {
				continue // network and broadcast addresses
			}
			for _, port := range d.cfg.Ports {
				if s, ok := d.probe(ctx, netip.AddrPortFrom(a, uint16(port)), "scan"); ok {
					add(s)
				}
			}
		}
	}
	sort.Slice(out, func(i, j int) bool {
		if out[i].Host != out[j].Host {
			return out[i].Host < out[j].Host
		}
		return out[i].Port < out[j].Port
	})
	return out
}

func lastAddr(p netip.Prefix) netip.Addr {
	b := p.Addr().As4()
	host := 32 - p.Bits()
	v := uint32(b[0])<<24 | uint32(b[1])<<16 | uint32(b[2])<<8 | uint32(b[3])
	v |= (1 << host) - 1
	return netip.AddrFrom4([4]byte{byte(v >> 24), byte(v >> 16), byte(v >> 8), byte(v)})
}

func (d *Discoverer) excluded(port int) bool {
	for _, p := range d.cfg.Exclude {
		if p == port {
			return true
		}
	}
	return false
}

// probe connects to one address (waiting its turn on the rate limiter) and identifies what answers.
func (d *Discoverer) probe(ctx context.Context, ap netip.AddrPort, source string) (model.DiscoveredService, bool) {
	if d.lim.Wait(ctx) != nil {
		return model.DiscoveredService{}, false
	}
	cctx, cancel := context.WithTimeout(ctx, d.cfg.timeout)
	defer cancel()
	c, err := d.cfg.dial(cctx, ap.String())
	if err != nil {
		return model.DiscoveredService{}, false
	}
	c.Close()
	svc, detail := d.identify(ctx, ap)
	return model.DiscoveredService{Host: ap.Addr().String(), Port: int(ap.Port()), Proto: "tcp", Service: svc, Detail: detail, Source: source}, true
}

// identify guesses the protocol: first by listening for a banner (SSH, SMTP, FTP and others speak
// first), then by a TLS handshake, then by one HTTP HEAD, and finally by the port number.
func (d *Discoverer) identify(ctx context.Context, ap netip.AddrPort) (service, detail string) {
	hint := wellKnown[int(ap.Port())]
	if c := d.connect(ctx, ap); c != nil {
		c.SetReadDeadline(time.Now().Add(400 * time.Millisecond))
		if line, ok := readLine(c); ok {
			c.Close()
			if s := bannerService(line); s != "" {
				return s, line
			}
			return hint, line
		}
		c.Close()
	}
	if c := d.connect(ctx, ap); c != nil {
		c.SetDeadline(time.Now().Add(d.cfg.timeout))
		t := tls.Client(c, &tls.Config{InsecureSkipVerify: true, ServerName: ap.Addr().String()}) // identification only; no data is sent
		if t.HandshakeContext(ctx) == nil {
			t.Close()
			return "https", ""
		}
		c.Close()
	}
	if c := d.connect(ctx, ap); c != nil {
		c.SetDeadline(time.Now().Add(d.cfg.timeout))
		fmt.Fprintf(c, "HEAD / HTTP/1.0\r\nHost: %s\r\nUser-Agent: bamboo-discovery\r\n\r\n", ap.Addr())
		br := bufio.NewReaderSize(c, 2048)
		line, _ := br.ReadString('\n')
		if line = trim(line); strings.HasPrefix(line, "HTTP/") {
			server := headerValue(br, "Server")
			c.Close()
			return "http", join(line, server)
		}
		c.Close()
	}
	return hint, ""
}

func (d *Discoverer) connect(ctx context.Context, ap netip.AddrPort) net.Conn {
	if d.lim.Wait(ctx) != nil {
		return nil
	}
	cctx, cancel := context.WithTimeout(ctx, d.cfg.timeout)
	defer cancel()
	c, err := d.cfg.dial(cctx, ap.String())
	if err != nil {
		return nil
	}
	return c
}

func readLine(c net.Conn) (string, bool) {
	line, _ := bufio.NewReaderSize(c, 512).ReadString('\n')
	line = trim(line)
	return line, line != ""
}

func trim(s string) string {
	s = strings.TrimSpace(s)
	if len(s) > 120 {
		s = s[:120]
	}
	return s
}

// headerValue scans the remaining response headers for one header.
func headerValue(br *bufio.Reader, name string) string {
	for i := 0; i < 40; i++ {
		l, err := br.ReadString('\n')
		l = strings.TrimSpace(l)
		if k, v, ok := strings.Cut(l, ":"); ok && strings.EqualFold(k, name) {
			return trim(v)
		}
		if err != nil || l == "" {
			break
		}
	}
	return ""
}

func join(a, b string) string {
	if b == "" {
		return a
	}
	return a + "; Server: " + b
}

func bannerService(line string) string {
	l := strings.ToUpper(line)
	switch {
	case strings.HasPrefix(l, "SSH-"):
		return "ssh"
	case strings.HasPrefix(l, "HTTP/"):
		return "http"
	case strings.HasPrefix(l, "220") && strings.Contains(l, "FTP"):
		return "ftp"
	case strings.HasPrefix(l, "220") && (strings.Contains(l, "SMTP") || strings.Contains(l, "ESMTP")):
		return "smtp"
	case strings.HasPrefix(l, "+OK"):
		return "pop3"
	case strings.HasPrefix(l, "* OK"):
		return "imap"
	case strings.HasPrefix(l, "RFB "):
		return "vnc"
	}
	return ""
}

// ParseProcNetTCP reads the text of /proc/net/tcp or /proc/net/tcp6 and returns listening endpoints.
func ParseProcNetTCP(text string, v6 bool) []netip.AddrPort {
	var out []netip.AddrPort
	sc := bufio.NewScanner(strings.NewReader(text))
	sc.Scan() // header
	for sc.Scan() {
		f := strings.Fields(sc.Text())
		if len(f) < 4 || f[3] != "0A" { // 0A = LISTEN
			continue
		}
		ip, port, ok := strings.Cut(f[1], ":")
		if !ok {
			continue
		}
		p, err := strconv.ParseUint(port, 16, 16)
		if err != nil {
			continue
		}
		if a, ok := parseProcAddr(ip, v6); ok {
			out = append(out, netip.AddrPortFrom(a, uint16(p)))
		}
	}
	return out
}

// parseProcAddr decodes the kernel's hex address: each 32-bit word is written in host (little-endian) order.
func parseProcAddr(h string, v6 bool) (netip.Addr, bool) {
	n := 8
	if v6 {
		n = 32
	}
	if len(h) != n {
		return netip.Addr{}, false
	}
	b := make([]byte, n/2)
	for i := range b {
		v, err := strconv.ParseUint(h[2*i:2*i+2], 16, 8)
		if err != nil {
			return netip.Addr{}, false
		}
		b[i] = byte(v)
	}
	for w := 0; w+4 <= len(b); w += 4 {
		b[w], b[w+1], b[w+2], b[w+3] = b[w+3], b[w+2], b[w+1], b[w]
	}
	if v6 {
		a := netip.AddrFrom16([16]byte(b))
		return a.Unmap(), true
	}
	return netip.AddrFrom4([4]byte(b)), true
}

// listeningSockets lists this machine's TCP listeners. Only Linux exposes them cheaply; elsewhere
// it returns nothing and the caller relies on scanning.
func listeningSockets() ([]netip.AddrPort, error) {
	if runtime.GOOS != "linux" {
		return nil, nil
	}
	var out []netip.AddrPort
	for _, f := range []struct {
		path string
		v6   bool
	}{{"/proc/net/tcp", false}, {"/proc/net/tcp6", true}} {
		b, err := os.ReadFile(f.path)
		if err != nil {
			continue
		}
		out = append(out, ParseProcNetTCP(string(b), f.v6)...)
	}
	return out, nil
}