// Package abuseguard automatically blacklists IPs that rack up too many failed
// SMTP/IMAP auth attempts, and rejects connections from already-blacklisted IPs before
// the SMTP/IMAP banner is ever sent. Deliberately separate from the web admin/webmail
// login lockout (internal/webui/ratelimit.go) and from the relay-authorization
// whitelist (esrv_whitelisted_ips) — see the [Security] section of settings.ini.
package abuseguard

import (
	"net"
	"sync"
	"time"

	"gopkg.in/ini.v1"

	"mailgoserver/internal/db"
	"mailgoserver/internal/toolbox"
)

// RecordFailureAndMaybeBlacklist should be called after every failed SMTP AUTH or IMAP
// login. It counts recent failures from ip and blacklists it once the configured
// threshold is hit. Fails open (does nothing) on a DB error rather than blocking auth
// over a transient issue.
func RecordFailureAndMaybeBlacklist(database *db.DB, cfg *ini.File, logger *toolbox.Logger, ip string) {
	if ip == "" || cfg == nil {
		return
	}
	sec := cfg.Section("Security")
	if !sec.Key("abuse_detection_enabled").MustBool(true) {
		return
	}
	if whitelisted, err := database.IsIPAbuseWhitelisted(ip); err != nil || whitelisted {
		return
	}

	threshold := sec.Key("abuse_failure_threshold").MustInt(8)
	windowMinutes := sec.Key("abuse_detection_window_minutes").MustInt(10)
	since := time.Now().Add(-time.Duration(windowMinutes) * time.Minute)

	n, err := database.CountFailedAuthAttemptsByIP(ip, since)
	if err != nil || n < threshold {
		return
	}

	baseHours := sec.Key("abuse_blacklist_base_hours").MustInt(12)
	maxHours := sec.Key("abuse_blacklist_max_hours").MustInt(168)
	reason := "automatic: too many failed SMTP/IMAP auth attempts"
	if err := database.BlacklistIP(ip, reason, baseHours, maxHours); err != nil && logger != nil {
		logger.Error("abuseguard: failed to blacklist %s: %v", ip, err)
		return
	}
	if logger != nil {
		logger.Warning("abuseguard: blacklisted %s after %d failed attempts in %dm", ip, n, windowMinutes)
	}
}

// guardedListener wraps a net.Listener so Accept() silently drops connections from
// blacklisted IPs (never returning them to the caller) before any protocol banner is
// written, and keeps looping rather than returning an error. Also caps concurrent
// connections per source IP — a resource-exhaustion guard distinct from the
// failed-auth-triggered blacklist above (a connection flood doesn't need to fail auth
// even once to tie up every worker/file-descriptor this server has).
type guardedListener struct {
	net.Listener
	database *db.DB
	logger   *toolbox.Logger
	maxPerIP int // <=0 means unlimited

	mu     sync.Mutex
	counts map[string]int
}

// GuardListener wraps inner so every accepted connection is checked against the IP
// blacklist (skipping the check entirely for abuse-whitelisted IPs) and the per-IP
// concurrent-connection cap ([Security] max_connections_per_ip, default 20) before the
// caller ever sees it.
func GuardListener(inner net.Listener, database *db.DB, cfg *ini.File, logger *toolbox.Logger) net.Listener {
	maxPerIP := 20
	if cfg != nil {
		maxPerIP = cfg.Section("Security").Key("max_connections_per_ip").MustInt(20)
	}
	return &guardedListener{Listener: inner, database: database, logger: logger, maxPerIP: maxPerIP, counts: make(map[string]int)}
}

func (g *guardedListener) Accept() (net.Conn, error) {
	for {
		conn, err := g.Listener.Accept()
		if err != nil {
			return nil, err
		}
		host, _, splitErr := net.SplitHostPort(conn.RemoteAddr().String())
		if splitErr != nil {
			host = conn.RemoteAddr().String()
		}
		abuseWhitelisted := false
		if whitelisted, wErr := g.database.IsIPAbuseWhitelisted(host); wErr == nil && whitelisted {
			abuseWhitelisted = true
		}
		if !abuseWhitelisted {
			blocked, bErr := g.database.IsIPBlacklisted(host)
			if bErr == nil && blocked {
				if g.logger != nil {
					g.logger.Warning("abuseguard: rejected connection from blacklisted IP %s", host)
				}
				conn.Close()
				continue
			}
		}
		// The concurrent-connection cap applies even to an abuse-whitelisted IP —
		// whitelisting exempts an IP from being auto-blacklisted over failed auth, not
		// from basic resource-exhaustion protection, a different concern.
		if g.maxPerIP > 0 {
			g.mu.Lock()
			if g.counts[host] >= g.maxPerIP {
				g.mu.Unlock()
				if g.logger != nil {
					g.logger.Warning("abuseguard: rejected connection from %s: at the concurrent-connection limit (%d)", host, g.maxPerIP)
				}
				conn.Close()
				continue
			}
			g.counts[host]++
			g.mu.Unlock()
			conn = &countedConn{Conn: conn, g: g, host: host}
		}
		return conn, nil
	}
}

// countedConn decrements guardedListener's per-IP counter exactly once, however Close
// ends up getting called (explicitly, via a defer, or both).
type countedConn struct {
	net.Conn
	g    *guardedListener
	host string
	once sync.Once
}

func (c *countedConn) Close() error {
	c.once.Do(func() {
		c.g.mu.Lock()
		c.g.counts[c.host]--
		if c.g.counts[c.host] <= 0 {
			delete(c.g.counts, c.host)
		}
		c.g.mu.Unlock()
	})
	return c.Conn.Close()
}