// Package bans is the ban ledger: the bans smallwebwaf makes on the
// netblocks of clients that break a rate limit or show a clear sign of
// attack, with their notes, as the "Bans" section of SPEC.md describes.
// The bans are kept in memory, and written to bans.json and read from it
// by the state package.
package bans

import (
	"net/netip"
	"slices"
	"strings"
	"sync"
	"time"

	"github.com/hashicorp/golang-lru/v2/simplelru"
)

// The causes of the bans smallwebwaf makes. A ban an admin adds to
// bans.json may have no cause.
const (
	// CauseLimit is a ban for a broken limit.
	CauseLimit = "limit"
	// CauseAttack is a ban for a clear sign of attack.
	CauseAttack = "attack"
)

// repeatFactor is how many times as long as the netblock's last ban a ban
// for a limit broken again within the repeat window lasts.
const repeatFactor = 3

// maxTextBytes is how much of each text in a ban's notes is kept.
const maxTextBytes = 256

// Rules are how long a ban lasts, and how many bans are held.
type Rules struct {
	// LimitBanDuration is how long a first ban for a broken limit lasts.
	LimitBanDuration time.Duration
	// LimitBanRepeatWindow is how soon after the end of the netblock's
	// ban that ended last, other than one for a clear sign of attack, a
	// broken limit counts as a repeat, which bans for repeatFactor times as
	// long as that ban.
	LimitBanRepeatWindow time.Duration
	// MaxBanDuration is the longest ban for a broken limit; one that would
	// be longer is permanent instead.
	MaxBanDuration time.Duration
	// AttackBanDuration is how long a first ban for a clear sign of attack
	// lasts.
	AttackBanDuration time.Duration
	// MaxBans is the most bans held, at least one. Past it, the earliest
	// ban of the netblock that has gone longest without a request is
	// dropped.
	MaxBans int
}

// Ban is a ban on a netblock.
type Ban struct {
	Netblock netip.Prefix
	Start    time.Time
	// Expires is when the ban ends, zero for a permanent ban.
	Expires time.Time
	// Cause is CauseLimit or CauseAttack, or "" for a ban an admin added
	// without one.
	Cause string
	Notes Notes
}

// Permanent reports whether the ban never runs out.
func (b Ban) Permanent() bool {
	return b.Expires.IsZero()
}

// ActiveAt reports whether the ban refuses requests at now.
func (b Ban) ActiveAt(now time.Time) bool {
	return b.Permanent() || now.Before(b.Expires)
}

// Notes are what an admin needs to decide whether to lift a ban. The
// JSON names are those of bans.json.
//
//nolint:tagliatelle // the state files use snake_case, as the request log does
type Notes struct {
	// Country is the client's country, when it was looked up.
	Country string `json:"country"`
	// Limit, Window and Count are, for a ban for a broken limit, the limit
	// that was broken, its window, "minute", "hour" or "day", and the
	// count reached: the client's requests in the window, the one that
	// broke the limit included. These are the requests that counted
	// toward the ban, and the window is the time over which they came.
	Limit  int64   `json:"limit,omitempty"`
	Window string  `json:"window,omitempty"`
	Count  float64 `json:"count,omitempty"`
	// RuleID and Target are, for a ban for a clear sign of attack, the id
	// of the rule file rule that matched, and its target.
	RuleID string `json:"rule_id,omitempty"`
	Target string `json:"target,omitempty"`
	// Request is the request that broke the limit, or that was the clear
	// sign of attack.
	Request Request `json:"request"`
	// Requests is how many requests the netblock has sent since it was
	// first seen, and Refused how many of them the ban has refused so
	// far. Both go up with each request the ban refuses.
	Requests int64 `json:"requests"`
	Refused  int64 `json:"refused"`
	// EarlierBans is how many bans the netblock had before this one, by
	// cause.
	EarlierBans EarlierBans `json:"earlier_bans"`
}

// EarlierBans counts a netblock's bans before a ban, by cause.
//
//nolint:tagliatelle // the state files use snake_case, as the request log does
type EarlierBans struct {
	Limit  int `json:"limit"`
	Attack int `json:"attack"`
	// WithoutCause counts the bans an admin added without a cause.
	WithoutCause int `json:"without_cause"`
}

// Request is a request in a ban's notes. Each text is cut to 256 bytes.
//
//nolint:tagliatelle // the state files use snake_case, as the request log does
type Request struct {
	Time   time.Time `json:"time"`
	Method string    `json:"method"`
	Host   string    `json:"host"`
	// Path is the path with its query string.
	Path string `json:"path"`
	// Status is what the client was sent, 0 if nothing was.
	Status    int    `json:"status"`
	UserAgent string `json:"user_agent"`
}

// Ledger holds the bans. It is safe for concurrent use.
type Ledger struct {
	rules Rules
	// changed receives a value when a ban is made, unless one is waiting
	// already.
	changed chan struct{}

	mu sync.Mutex
	// netblocks holds each banned netblock's bans, oldest first. Check and
	// Find make each netblock they find the most recently seen.
	netblocks *simplelru.LRU[netip.Prefix, *[]Ban]
	// held is how many bans netblocks holds, at most rules.MaxBans.
	held int
	// made is how many bans the ledger has made since the start, by cause.
	made map[string]int
	// v4Lengths and v6Lengths are the lengths of the IPv4 and IPv6
	// netblocks that have been banned. Check looks for a ban at each of
	// them, so that a ban read from bans.json refuses every client in its
	// netblock even when it was made with another SWWAF_BAN_SCOPE_V4_PREFIX,
	// or another length of an IPv6 client's netblock.
	v4Lengths, v6Lengths []int
}

// New returns a Ledger with no ban yet.
func New(rules Rules) *Ledger {
	// Every netblock held has a ban, so there are never more netblocks
	// than rules.MaxBans, and the LRU never drops one itself.
	netblocks, err := simplelru.NewLRU[netip.Prefix, *[]Ban](rules.MaxBans, nil)
	if err != nil {
		panic(err) // NewLRU fails only for a size below one
	}

	return &Ledger{
		rules:     rules,
		changed:   make(chan struct{}, 1),
		netblocks: netblocks,
		made:      map[string]int{},
	}
}

// Changed receives a value after a ban is made or made permanent, so that
// bans.json can be written. Several changes before it is read leave one
// value.
func (l *Ledger) Changed() <-chan struct{} {
	return l.changed
}

// Check is called for a request from client, at now. It reports whether
// a ban on a netblock client is in is active, and returns that ban, with
// the request counted among those it refused. A ban for a clear sign of
// attack is made permanent by the request: the netblock is malicious.
func (l *Ledger) Check(client netip.Addr, now time.Time) (Ban, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	ban := l.active(client, now)
	if ban == nil {
		return Ban{}, false
	}

	ban.Notes.Requests++
	ban.Notes.Refused++

	if ban.Cause == CauseAttack && !ban.Permanent() {
		ban.Expires = time.Time{}

		l.markChanged()
	}

	return *ban, true
}

// Find is Check without counting the request among those the ban
// refused: in observe mode a ban refuses nothing.
func (l *Ledger) Find(client netip.Addr, now time.Time) (Ban, bool) {
	l.mu.Lock()
	defer l.mu.Unlock()

	ban := l.active(client, now)
	if ban == nil {
		return Ban{}, false
	}

	return *ban, true
}

// activeBan returns the ban in bans, a netblock's bans oldest first, that
// is active at now, or nil when none is. If several are, it returns the
// one that started last. Every ban is looked at, since a ban an admin adds
// to bans.json can start before the netblock's others and outlast them.
func activeBan(bans []Ban, now time.Time) *Ban {
	for i := len(bans) - 1; i >= 0; i-- {
		if bans[i].ActiveAt(now) {
			return &bans[i]
		}
	}

	return nil
}

// BanForLimit bans netblock at now for a broken limit, with notes, and
// returns the ban. A first ban lasts LimitBanDuration. A ban made within
// LimitBanRepeatWindow after the netblock's ban that ended last, other
// than one for a clear sign of attack, lasts repeatFactor times as long as
// that one. A ban that would be longer than MaxBanDuration is permanent
// instead. If a ban on netblock is still active, as when two of its
// requests break a limit at once, that ban is returned and no other is
// made. The ledger fills in the notes' Refused and EarlierBans itself.
func (l *Ledger) BanForLimit(netblock netip.Prefix, now time.Time, notes Notes) Ban {
	return l.ban(netblock, now, CauseLimit, notes)
}

// BanForAttack bans netblock at now for a clear sign of attack, with
// notes, and returns the ban, as BanForLimit does. A first ban lasts
// AttackBanDuration; once the netblock has had one, the next is
// permanent.
func (l *Ledger) BanForAttack(netblock netip.Prefix, now time.Time, notes Notes) Ban {
	return l.ban(netblock, now, CauseAttack, notes)
}

// Bans returns the bans held on netblock, oldest first. It is not a
// request from netblock, and leaves when it was last seen unchanged.
func (l *Ledger) Bans(netblock netip.Prefix) []Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	bans, found := l.netblocks.Peek(netblock)
	if !found {
		return nil
	}

	return slices.Clone(*bans)
}

// Made returns how many bans for cause the ledger has made since the
// start; bans read from bans.json are not among them.
func (l *Ledger) Made(cause string) int {
	l.mu.Lock()
	defer l.mu.Unlock()

	return l.made[cause]
}

// Count returns how many of the bans held are active at now, and how many
// are permanent.
func (l *Ledger) Count(now time.Time) (int, int) {
	l.mu.Lock()
	defer l.mu.Unlock()

	active, permanent := 0, 0

	for _, bans := range l.netblocks.Values() {
		for _, ban := range *bans {
			if ban.ActiveAt(now) {
				active++
			}

			if ban.Permanent() {
				permanent++
			}
		}
	}

	return active, permanent
}

// Snapshot returns every ban held, sorted by netblock, and each
// netblock's bans oldest first, as bans.json lists them.
func (l *Ledger) Snapshot() []Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	held := make([]Ban, 0, l.held)
	for _, bans := range l.netblocks.Values() {
		held = append(held, *bans...)
	}

	slices.SortStableFunc(held, func(a, b Ban) int {
		return a.Netblock.Compare(b.Netblock)
	})

	return held
}

// Load puts bans read from bans.json into the ledger, in place of the
// bans it holds, in the order they started, so that a netblock whose last
// ban started latest counts as the most recently seen. Each netblock is
// masked to its length, so that 203.0.113.9/24 is 203.0.113.0/24, and
// each text in the notes is cut to 256 bytes. Past MaxBans the earliest
// bans are dropped, as when they are made.
func (l *Ledger) Load(bans []Ban) {
	bans = slices.Clone(bans)
	slices.SortStableFunc(bans, func(a, b Ban) int {
		return a.Start.Compare(b.Start)
	})

	l.mu.Lock()
	defer l.mu.Unlock()

	l.netblocks.Purge()
	l.held = 0
	l.v4Lengths, l.v6Lengths = nil, nil

	for _, ban := range bans {
		ban.Netblock = ban.Netblock.Masked()
		ban.Notes.Request = ban.Notes.Request.cut()
		l.add(ban)
	}
}

// ban bans netblock at now for cause, with notes, as BanForLimit and
// BanForAttack describe, and returns the ban.
func (l *Ledger) ban(
	netblock netip.Prefix, now time.Time, cause string, notes Notes,
) Ban {
	l.mu.Lock()
	defer l.mu.Unlock()

	// held are the netblock's bans, none of them active.
	var held []Ban

	bans, found := l.netblocks.Get(netblock)
	if found {
		active := activeBan(*bans, now)
		if active != nil {
			return *active
		}

		held = *bans
		notes.EarlierBans = earlierBans(held)
	}

	notes.Request = notes.Request.cut()
	ban := Ban{Netblock: netblock, Start: now, Cause: cause, Notes: notes}

	if cause == CauseAttack {
		ban.Expires = l.attackExpiry(held, now)
	} else {
		ban.Expires = l.limitExpiry(held, now)
	}

	l.add(ban)
	l.made[cause]++
	l.markChanged()

	return ban
}

// earlierBans returns how many bans a netblock with the bans held, oldest
// first, has had, by cause: the first ban held counts the bans the
// netblock had before that one, since dropped to make room, and each ban
// held adds one.
func earlierBans(held []Ban) EarlierBans {
	earlier := held[0].Notes.EarlierBans

	for _, ban := range held {
		switch ban.Cause {
		case CauseLimit:
			earlier.Limit++
		case CauseAttack:
			earlier.Attack++
		default:
			earlier.WithoutCause++
		}
	}

	return earlier
}

// markChanged has Changed receive a value, unless one is waiting already.
func (l *Ledger) markChanged() {
	select {
	case l.changed <- struct{}{}:
	default: // a value is waiting already
	}
}

// active returns the ban active at now on a netblock client is in, or
// nil.
func (l *Ledger) active(client netip.Addr, now time.Time) *Ban {
	lengths := l.v6Lengths
	if client.Is4() {
		lengths = l.v4Lengths
	}

	for _, length := range lengths {
		bans, found := l.netblocks.Get(netip.PrefixFrom(client, length).Masked())
		if !found {
			continue
		}

		ban := activeBan(*bans, now)
		if ban != nil {
			return ban
		}
	}

	return nil
}

// add adds ban to its netblock's bans, after the last, and makes its
// netblock the most recently seen. With MaxBans held, it drops one first.
func (l *Ledger) add(ban Ban) {
	if l.held == l.rules.MaxBans {
		l.dropOne()
	}

	// dropOne can have dropped the netblock's last ban, and the netblock
	// with it.
	bans, found := l.netblocks.Get(ban.Netblock)
	if !found {
		bans = &[]Ban{}
		l.netblocks.Add(ban.Netblock, bans)
	}

	*bans = append(*bans, ban)
	l.held++

	lengths := &l.v6Lengths
	if ban.Netblock.Addr().Is4() {
		lengths = &l.v4Lengths
	}

	if !slices.Contains(*lengths, ban.Netblock.Bits()) {
		*lengths = append(*lengths, ban.Netblock.Bits())
	}
}

// limitExpiry returns when a ban for a broken limit made at now ends, or
// zero when it is permanent. held are the netblock's bans, none of them
// active, of which the one that ended last, other than a ban for a clear
// sign of attack, can make the new ban longer. A ban an admin adds to
// bans.json can start after another and end before it, so that one is
// looked for among them all.
func (l *Ledger) limitExpiry(held []Ban, now time.Time) time.Time {
	length := l.rules.LimitBanDuration

	var last *Ban

	for i, ban := range held {
		if ban.Cause != CauseAttack && (last == nil || ban.Expires.After(last.Expires)) {
			last = &held[i]
		}
	}

	if last != nil && now.Sub(last.Expires) <= l.rules.LimitBanRepeatWindow {
		lastLength := last.Expires.Sub(last.Start)
		// This is repeatFactor * lastLength > MaxBanDuration, written so
		// that it cannot overflow.
		if lastLength > l.rules.MaxBanDuration/repeatFactor {
			return time.Time{}
		}

		length = repeatFactor * lastLength
	}

	if length > l.rules.MaxBanDuration {
		return time.Time{}
	}

	return now.Add(length)
}

// attackExpiry returns when a ban for a clear sign of attack made at now
// ends. held are the netblock's bans, none of them active: if one of them
// is for a clear sign of attack too, the new ban is permanent, and its
// end zero; otherwise it ends AttackBanDuration later.
func (l *Ledger) attackExpiry(held []Ban, now time.Time) time.Time {
	for _, ban := range held {
		if ban.Cause == CauseAttack {
			return time.Time{}
		}
	}

	return now.Add(l.rules.AttackBanDuration)
}

// dropOne drops the earliest ban of the netblock that has gone longest
// without a request, and the netblock with it if that was its only ban.
func (l *Ledger) dropOne() {
	netblock, bans, _ := l.netblocks.GetOldest()
	if len(*bans) == 1 {
		l.netblocks.Remove(netblock)
	} else {
		*bans = slices.Delete(*bans, 0, 1)
	}

	l.held--
}

// cut returns r with each text cut to maxTextBytes and copied, so that
// the notes do not keep the rest of the request in memory.
func (r Request) cut() Request {
	r.Method = cutText(r.Method)
	r.Host = cutText(r.Host)
	r.Path = cutText(r.Path)
	r.UserAgent = cutText(r.UserAgent)

	return r
}

// cutText returns a copy of the first maxTextBytes of text.
func cutText(text string) string {
	return strings.Clone(text[:min(len(text), maxTextBytes)])
}