package main

import (
	"bufio"
	"compress/gzip"
	"encoding/json"
	"errors"
	"flag"
	"fmt"
	"io"
	"net/netip"
	"os"
	"os/exec"
	"path/filepath"
	"strconv"
	"strings"
	"time"
)

const (
	tableName             = "oci_l4_relay_smtp_deny"
	activePolicyStateFile = "active-state.json"
	maximumBundleBytes    = 16 << 20
	absoluteMaxEntries    = 100000
	absoluteMaxLife       = 7 * 24 * time.Hour
)

var transportReasons = map[string]struct{}{
	"connection-churn": {},
	"connection-rate":  {},
	"syn-rate":         {},
}

type denyEntry struct {
	Start              string    `json:"start"`
	End                string    `json:"end"`
	Reason             string    `json:"reason"`
	ObservationStarted time.Time `json:"observation_started"`
	ObservationEnded   time.Time `json:"observation_ended"`
	ObservationCount   uint64    `json:"observation_count"`
	ExpiresAt          time.Time `json:"expires_at"`
}

type bundle struct {
	SchemaVersion       int         `json:"schema_version"`
	GeneratedAt         time.Time   `json:"generated_at"`
	ExpiresAt           time.Time   `json:"expires_at"`
	DeniedIPv4          []denyEntry `json:"denied_ipv4"`
	DeniedIntervalCount int         `json:"denied_interval_count"`
	DeniedAddressCount  uint64      `json:"denied_address_count"`
}

type parsedEntry struct {
	start     uint32
	end       uint32
	expiresAt time.Time
}

type validatedBundle struct {
	bundle
	entries []parsedEntry
}

type activePolicyState struct {
	ActivatedAt         time.Time `json:"activated_at"`
	Source              string    `json:"source"`
	BundleGeneratedAt   time.Time `json:"bundle_generated_at,omitempty"`
	BundleExpiresAt     time.Time `json:"bundle_expires_at,omitempty"`
	LastRemoteCurrentAt time.Time `json:"last_remote_current_at,omitempty"`
}

type activationConfig struct {
	publicInterface          string
	stateDirectory           string
	nft                      string
	minimumObservations      uint64
	maximumAddresses         uint64
	maximumIntervals         int
	maximumLifetime          time.Duration
	maximumObservationWindow time.Duration
}

func main() {
	if len(os.Args) < 2 {
		fatal(errors.New("expected apply, empty, restore, validate, or render"))
	}
	switch os.Args[1] {
	case "apply":
		flags := flag.NewFlagSet("apply", flag.ExitOnError)
		configuration := activationConfig{}
		input := flags.String("input", "", "downloaded SMTP deny bundle")
		addActivationFlags(flags, &configuration)
		_ = flags.Parse(os.Args[2:])
		if *input == "" {
			fatal(errors.New("input is required"))
		}
		now := time.Now().UTC()
		if err := activate(now, configuration, *input, true, "current"); err != nil {
			applyErr := err
			if restoreErr := restore(now, configuration); restoreErr != nil {
				if failOpenErr := failOpen(now, configuration); failOpenErr != nil {
					fatal(errors.Join(applyErr, restoreErr, failOpenErr))
				}
				fatal(errors.Join(applyErr, restoreErr))
			}
			fmt.Println("local_apply=false cache_restored=true")
			fatal(applyErr)
		}
	case "empty":
		flags := flag.NewFlagSet("empty", flag.ExitOnError)
		output := flags.String("output", "", "empty SMTP deny bundle output")
		lifetime := flags.Duration("lifetime", 24*time.Hour, "empty bundle lifetime")
		_ = flags.Parse(os.Args[2:])
		if *output == "" || *lifetime <= 0 || *lifetime > absoluteMaxLife {
			fatal(errors.New("output and a bounded lifetime are required"))
		}
		now := time.Now().UTC().Truncate(time.Second)
		document := bundle{
			SchemaVersion: 1,
			GeneratedAt:   now,
			ExpiresAt:     now.Add(*lifetime),
			DeniedIPv4:    []denyEntry{},
		}
		if err := writeBundle(*output, document); err != nil {
			fatal(err)
		}
	case "restore":
		flags := flag.NewFlagSet("restore", flag.ExitOnError)
		configuration := activationConfig{}
		addActivationFlags(flags, &configuration)
		_ = flags.Parse(os.Args[2:])
		now := time.Now().UTC()
		if err := restore(now, configuration); err != nil {
			if failOpenErr := failOpen(now, configuration); failOpenErr != nil {
				fatal(errors.Join(err, failOpenErr))
			}
			fatal(err)
		}
	case "validate", "render":
		flags := flag.NewFlagSet(os.Args[1], flag.ExitOnError)
		configuration := activationConfig{}
		input := flags.String("input", "", "SMTP deny bundle")
		previousInput := flags.String("previous-input", "", "previous SMTP deny bundle")
		summaryOutput := flags.String("summary-output", "", "sanitized validation summary")
		addActivationFlags(flags, &configuration)
		_ = flags.Parse(os.Args[2:])
		now := time.Now().UTC()
		validated, err := readBundle(now, configuration, *input)
		if err != nil {
			fatal(err)
		}
		removedIntervals := 0
		if *previousInput != "" {
			previous, err := readBundle(now, configuration, *previousInput)
			if err != nil {
				fatal(err)
			}
			removedIntervals = removedIntervalCount(previous, validated)
		}
		if *summaryOutput != "" {
			summary := struct {
				GeneratedAt          time.Time `json:"generated_at"`
				ExpiresAt            time.Time `json:"expires_at"`
				DeniedIntervalCount  int       `json:"denied_interval_count"`
				DeniedAddressCount   uint64    `json:"denied_address_count"`
				RemovedIntervalCount int       `json:"removed_interval_count"`
			}{
				GeneratedAt:          validated.GeneratedAt,
				ExpiresAt:            validated.ExpiresAt,
				DeniedIntervalCount:  validated.DeniedIntervalCount,
				DeniedAddressCount:   validated.DeniedAddressCount,
				RemovedIntervalCount: removedIntervals,
			}
			if err := writeJSONFile(*summaryOutput, summary); err != nil {
				fatal(err)
			}
		}
		if os.Args[1] == "render" {
			if err := render(os.Stdout, now, validated, configuration.publicInterface); err != nil {
				fatal(err)
			}
		}
	default:
		fatal(fmt.Errorf("unknown command %q", os.Args[1]))
	}
}

func writeBundle(path string, document bundle) error {
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	temporary, err := os.CreateTemp(filepath.Dir(path), ".bundle-")
	if err != nil {
		return err
	}
	name := temporary.Name()
	defer func() { _ = os.Remove(name) }()
	writer := io.Writer(temporary)
	var compressed *gzip.Writer
	if strings.HasSuffix(path, ".gz") {
		compressed = gzip.NewWriter(temporary)
		compressed.ModTime = time.Time{}
		writer = compressed
	}
	encoder := json.NewEncoder(writer)
	if err := encoder.Encode(document); err != nil {
		_ = temporary.Close()
		return err
	}
	if compressed != nil {
		if err := compressed.Close(); err != nil {
			_ = temporary.Close()
			return err
		}
	}
	if err := temporary.Chmod(0600); err != nil {
		_ = temporary.Close()
		return err
	}
	if err := temporary.Sync(); err != nil {
		_ = temporary.Close()
		return err
	}
	if err := temporary.Close(); err != nil {
		return err
	}
	return os.Rename(name, path)
}

func writeJSONFile(path string, value any) error {
	contents, err := json.Marshal(value)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(filepath.Dir(path), 0700); err != nil {
		return err
	}
	return os.WriteFile(path, append(contents, '\n'), 0644)
}

func removedIntervalCount(previous, current validatedBundle) int {
	retained := make(map[[2]uint32]struct{}, len(current.entries))
	for _, entry := range current.entries {
		retained[[2]uint32{entry.start, entry.end}] = struct{}{}
	}
	removed := 0
	for _, entry := range previous.entries {
		if _, exists := retained[[2]uint32{entry.start, entry.end}]; !exists {
			removed++
		}
	}
	return removed
}

func addActivationFlags(flags *flag.FlagSet, configuration *activationConfig) {
	flags.StringVar(&configuration.publicInterface, "public-interface", "", "public network interface")
	flags.StringVar(&configuration.stateDirectory, "state-directory", "/var/lib/oci-l4-relay/smtp-deny", "state directory")
	flags.StringVar(&configuration.nft, "nft", "/usr/sbin/nft", "nft path")
	flags.Uint64Var(&configuration.minimumObservations, "minimum-observations", 1, "minimum observations for each entry")
	flags.Uint64Var(&configuration.maximumAddresses, "maximum-addresses", 1, "maximum covered IPv4 addresses")
	flags.IntVar(&configuration.maximumIntervals, "maximum-intervals", 1, "maximum IPv4 intervals")
	flags.DurationVar(&configuration.maximumLifetime, "maximum-lifetime", 24*time.Hour, "maximum entry and bundle lifetime")
	flags.DurationVar(&configuration.maximumObservationWindow, "maximum-observation-window", time.Hour, "maximum observation window")
}

func fatal(err error) {
	fmt.Fprintln(os.Stderr, err)
	os.Exit(1)
}

func activate(now time.Time, configuration activationConfig, bundlePath string, persist bool, source string) error {
	if err := writeActivePolicyState(configuration.stateDirectory, now, "transition", nil); err != nil {
		return err
	}
	validated, err := readBundle(now, configuration, bundlePath)
	if err != nil {
		return err
	}
	temporaryDirectory, err := os.MkdirTemp(configuration.stateDirectory, ".activate-")
	if err != nil {
		return err
	}
	defer func() { _ = os.RemoveAll(temporaryDirectory) }()
	rulesetPath := filepath.Join(temporaryDirectory, "smtp-deny.nft")
	ruleset, err := os.OpenFile(rulesetPath, os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0600)
	if err != nil {
		return err
	}
	if err := render(ruleset, now, validated, configuration.publicInterface); err != nil {
		_ = ruleset.Close()
		return err
	}
	if err := ruleset.Close(); err != nil {
		return err
	}
	if err := ensureTable(configuration.nft); err != nil {
		return err
	}
	if output, err := exec.Command(configuration.nft, "--check", "--file", rulesetPath).CombinedOutput(); err != nil {
		return fmt.Errorf("validate rendered ruleset: %w: %s", err, strings.TrimSpace(string(output)))
	}
	if output, err := exec.Command(configuration.nft, "--file", rulesetPath).CombinedOutput(); err != nil {
		return fmt.Errorf("load rendered ruleset: %w: %s", err, strings.TrimSpace(string(output)))
	}
	if persist {
		if err := replaceFile(bundlePath, filepath.Join(configuration.stateDirectory, "last-good.json.gz")); err != nil {
			return err
		}
	}
	if err := writeActivePolicyState(configuration.stateDirectory, now, source, &validated); err != nil {
		return err
	}
	fmt.Printf("denied_intervals=%d denied_addresses=%d expires_at=%s\n", len(validated.entries), validated.DeniedAddressCount, validated.ExpiresAt.Format(time.RFC3339))
	return nil
}

func restore(now time.Time, configuration activationConfig) error {
	if configuration.publicInterface == "" {
		return errors.New("public-interface is required")
	}
	return activate(now, configuration, filepath.Join(configuration.stateDirectory, "last-good.json.gz"), false, "cache")
}

func failOpen(now time.Time, configuration activationConfig) error {
	if err := writeActivePolicyState(configuration.stateDirectory, now, "transition", nil); err != nil {
		return err
	}
	if err := ensureTable(configuration.nft); err != nil {
		return err
	}
	ruleset, err := os.CreateTemp("", "oci-l4-relay-smtp-deny-fail-open-*.nft")
	if err != nil {
		return err
	}
	path := ruleset.Name()
	defer func() { _ = os.Remove(path) }()
	if err := renderFailOpen(ruleset, configuration.publicInterface); err != nil {
		_ = ruleset.Close()
		return err
	}
	if err := ruleset.Close(); err != nil {
		return err
	}
	if output, err := exec.Command(configuration.nft, "--check", "--file", path).CombinedOutput(); err != nil {
		return fmt.Errorf("validate fail-open ruleset: %w: %s", err, strings.TrimSpace(string(output)))
	}
	if output, err := exec.Command(configuration.nft, "--file", path).CombinedOutput(); err != nil {
		return fmt.Errorf("load fail-open ruleset: %w: %s", err, strings.TrimSpace(string(output)))
	}
	if err := writeActivePolicyState(configuration.stateDirectory, now, "fail-open", nil); err != nil {
		return err
	}
	fmt.Println("smtp_deny_fail_open=true")
	return nil
}

func writeActivePolicyState(directory string, now time.Time, source string, active *validatedBundle) error {
	state := activePolicyState{ActivatedAt: now, Source: source}
	path := filepath.Join(directory, activePolicyStateFile)
	if previous, err := readActivePolicyState(path); err == nil {
		state.LastRemoteCurrentAt = previous.LastRemoteCurrentAt
	}
	if active != nil {
		state.BundleGeneratedAt = active.GeneratedAt
		state.BundleExpiresAt = active.ExpiresAt
	}
	if source == "current" {
		state.LastRemoteCurrentAt = now
	}
	contents, err := json.Marshal(state)
	if err != nil {
		return err
	}
	if err := os.MkdirAll(directory, 0700); err != nil {
		return err
	}
	temporary, err := os.CreateTemp(directory, ".active-state-")
	if err != nil {
		return err
	}
	temporaryPath := temporary.Name()
	defer func() { _ = os.Remove(temporaryPath) }()
	if err := temporary.Chmod(0600); err != nil {
		_ = temporary.Close()
		return err
	}
	if _, err := temporary.Write(append(contents, '\n')); err != nil {
		_ = temporary.Close()
		return err
	}
	if err := temporary.Sync(); err != nil {
		_ = temporary.Close()
		return err
	}
	if err := temporary.Close(); err != nil {
		return err
	}
	return os.Rename(temporaryPath, path)
}

func readActivePolicyState(path string) (activePolicyState, error) {
	contents, err := os.ReadFile(path)
	if err != nil {
		return activePolicyState{}, err
	}
	decoder := json.NewDecoder(strings.NewReader(string(contents)))
	decoder.DisallowUnknownFields()
	state := activePolicyState{}
	if err := decoder.Decode(&state); err != nil {
		return activePolicyState{}, err
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		return activePolicyState{}, errors.New("active policy state contains trailing data")
	}
	return state, nil
}

func ensureTable(nft string) error {
	if err := exec.Command(nft, "list", "table", "inet", tableName).Run(); err == nil {
		return nil
	}
	if output, err := exec.Command(nft, "add", "table", "inet", tableName).CombinedOutput(); err != nil {
		return fmt.Errorf("create SMTP deny table: %w: %s", err, strings.TrimSpace(string(output)))
	}
	return nil
}

func replaceFile(source, destination string) error {
	contents, err := os.ReadFile(source)
	if err != nil {
		return err
	}
	temporary, err := os.CreateTemp(filepath.Dir(destination), ".last-good-")
	if err != nil {
		return err
	}
	name := temporary.Name()
	defer func() { _ = os.Remove(name) }()
	if err := temporary.Chmod(0600); err != nil {
		_ = temporary.Close()
		return err
	}
	if _, err := temporary.Write(contents); err != nil {
		_ = temporary.Close()
		return err
	}
	if err := temporary.Sync(); err != nil {
		_ = temporary.Close()
		return err
	}
	if err := temporary.Close(); err != nil {
		return err
	}
	return os.Rename(name, destination)
}

func readBundle(now time.Time, configuration activationConfig, path string) (validatedBundle, error) {
	file, err := os.Open(path)
	if err != nil {
		return validatedBundle{}, err
	}
	defer func() { _ = file.Close() }()
	reader := io.LimitReader(file, maximumBundleBytes+1)
	if strings.HasSuffix(path, ".gz") {
		compressed, err := gzip.NewReader(reader)
		if err != nil {
			return validatedBundle{}, fmt.Errorf("open bundle gzip: %w", err)
		}
		defer func() { _ = compressed.Close() }()
		reader = io.LimitReader(compressed, maximumBundleBytes+1)
	}
	decoder := json.NewDecoder(reader)
	decoder.DisallowUnknownFields()
	document := bundle{}
	if err := decoder.Decode(&document); err != nil {
		return validatedBundle{}, fmt.Errorf("decode bundle: %w", err)
	}
	if err := decoder.Decode(&struct{}{}); err != io.EOF {
		return validatedBundle{}, errors.New("bundle contains trailing data")
	}
	return validateBundle(now, configuration, document)
}

func validateBundle(now time.Time, configuration activationConfig, document bundle) (validatedBundle, error) {
	if configuration.minimumObservations == 0 ||
		configuration.maximumAddresses == 0 ||
		configuration.maximumIntervals < 0 ||
		configuration.maximumIntervals > absoluteMaxEntries ||
		configuration.maximumLifetime <= 0 ||
		configuration.maximumLifetime > absoluteMaxLife ||
		configuration.maximumObservationWindow <= 0 ||
		configuration.maximumObservationWindow > configuration.maximumLifetime {
		return validatedBundle{}, errors.New("invalid policy bounds")
	}
	if document.SchemaVersion != 1 {
		return validatedBundle{}, fmt.Errorf("unsupported schema version %d", document.SchemaVersion)
	}
	if document.GeneratedAt.After(now.Add(5 * time.Minute)) {
		return validatedBundle{}, errors.New("bundle generation time is in the future")
	}
	if !document.ExpiresAt.After(now) ||
		!document.ExpiresAt.After(document.GeneratedAt) ||
		document.ExpiresAt.Sub(document.GeneratedAt) > configuration.maximumLifetime ||
		document.GeneratedAt.Before(now.Add(-configuration.maximumLifetime)) {
		return validatedBundle{}, errors.New("bundle lifetime or generation age exceeds bounds")
	}
	if len(document.DeniedIPv4) > configuration.maximumIntervals {
		return validatedBundle{}, errors.New("denied interval count exceeds configured bound")
	}
	entries := make([]parsedEntry, 0, len(document.DeniedIPv4))
	var addressCount uint64
	var previousEnd uint32
	for index, entry := range document.DeniedIPv4 {
		start, err := parseIPv4(entry.Start)
		if err != nil {
			return validatedBundle{}, fmt.Errorf("entry %d start: %w", index, err)
		}
		end, err := parseIPv4(entry.End)
		if err != nil {
			return validatedBundle{}, fmt.Errorf("entry %d end: %w", index, err)
		}
		if start > end {
			return validatedBundle{}, fmt.Errorf("entry %d is reversed", index)
		}
		if index > 0 && start <= previousEnd {
			return validatedBundle{}, fmt.Errorf("entry %d overlaps or is unsorted", index)
		}
		if _, exists := transportReasons[entry.Reason]; !exists {
			return validatedBundle{}, fmt.Errorf("entry %d has invalid transport reason", index)
		}
		if entry.ObservationCount < configuration.minimumObservations {
			return validatedBundle{}, fmt.Errorf("entry %d is below the observation threshold", index)
		}
		if entry.ObservationStarted.After(entry.ObservationEnded) ||
			entry.ObservationEnded.After(document.GeneratedAt.Add(5*time.Minute)) ||
			entry.ObservationEnded.Sub(entry.ObservationStarted) > configuration.maximumObservationWindow {
			return validatedBundle{}, fmt.Errorf("entry %d has invalid observation window", index)
		}
		if !entry.ExpiresAt.After(entry.ObservationEnded) ||
			entry.ExpiresAt.After(document.ExpiresAt) ||
			entry.ExpiresAt.Sub(entry.ObservationEnded) > configuration.maximumLifetime {
			return validatedBundle{}, fmt.Errorf("entry %d has invalid expiry", index)
		}
		addressCount += uint64(end) - uint64(start) + 1
		if addressCount > configuration.maximumAddresses {
			return validatedBundle{}, errors.New("denied address count exceeds configured bound")
		}
		entries = append(entries, parsedEntry{start: start, end: end, expiresAt: entry.ExpiresAt})
		previousEnd = end
	}
	if document.DeniedIntervalCount != len(entries) || document.DeniedAddressCount != addressCount {
		return validatedBundle{}, errors.New("denied statistics do not match payload")
	}
	return validatedBundle{bundle: document, entries: entries}, nil
}

func parseIPv4(value string) (uint32, error) {
	address, err := netip.ParseAddr(value)
	if err != nil || !address.Is4() || address.String() != value {
		return 0, fmt.Errorf("invalid canonical IPv4 address %q", value)
	}
	bytes := address.As4()
	return uint32(bytes[0])<<24 | uint32(bytes[1])<<16 | uint32(bytes[2])<<8 | uint32(bytes[3]), nil
}

func ipv4String(value uint32) string {
	return netip.AddrFrom4([4]byte{byte(value >> 24), byte(value >> 16), byte(value >> 8), byte(value)}).String()
}

func render(output io.Writer, now time.Time, document validatedBundle, publicInterface string) error {
	if err := validateInterface(publicInterface); err != nil {
		return err
	}
	writer := bufio.NewWriter(output)
	if _, err := fmt.Fprintf(writer, "delete table inet %s\ntable inet %s {\n", tableName, tableName); err != nil {
		return err
	}
	if _, err := fmt.Fprintln(writer, "    set denied_v4 {\n        type ipv4_addr\n        flags interval, timeout"); err != nil {
		return err
	}
	active := make([]parsedEntry, 0, len(document.entries))
	for _, entry := range document.entries {
		if entry.expiresAt.Sub(now) >= time.Second {
			active = append(active, entry)
		}
	}
	if len(active) > 0 {
		if _, err := fmt.Fprintln(writer, "        elements = {"); err != nil {
			return err
		}
		for index, entry := range active {
			address := ipv4String(entry.start)
			if entry.start != entry.end {
				address += "-" + ipv4String(entry.end)
			}
			separator := ","
			if index == len(active)-1 {
				separator = ""
			}
			if _, err := fmt.Fprintf(writer, "            %s timeout %ds%s\n", address, int64(entry.expiresAt.Sub(now)/time.Second), separator); err != nil {
				return err
			}
		}
		if _, err := fmt.Fprintln(writer, "        }"); err != nil {
			return err
		}
	}
	if _, err := fmt.Fprintf(writer, "    }\n    counter smtp_denied {}\n    chain ingress {\n        type filter hook prerouting priority raw - 19; policy accept;\n        iifname %s tcp dport 25 ip saddr @denied_v4 counter name smtp_denied drop\n    }\n}\n", strconv.Quote(publicInterface)); err != nil {
		return err
	}
	return writer.Flush()
}

func renderFailOpen(output io.Writer, publicInterface string) error {
	if err := validateInterface(publicInterface); err != nil {
		return err
	}
	_, err := fmt.Fprintf(output, "delete table inet %s\ntable inet %s {\n    set denied_v4 {\n        type ipv4_addr\n        flags interval, timeout\n    }\n    counter smtp_denied {}\n    chain ingress {\n        type filter hook prerouting priority raw - 19; policy accept;\n        iifname %s tcp dport 25 ip saddr @denied_v4 counter name smtp_denied drop\n    }\n}\n", tableName, tableName, strconv.Quote(publicInterface))
	return err
}

func validateInterface(value string) error {
	if value == "" || strings.ContainsAny(value, "\"\\\n\r\t ") {
		return errors.New("invalid public interface")
	}
	return nil
}