package main

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

const (
	tableName       = "oci_l4_relay_web_allow"
	maxBundleBytes  = 64 << 20
	maxIntervals    = 500000
	maxExceptions   = 10000
	minimumLifetime = time.Hour
	maximumLifetime = 7 * 24 * time.Hour
)

type addressRange struct {
	Start string `json:"start"`
	End   string `json:"end"`
}

type bundle struct {
	SchemaVersion          int            `json:"schema_version"`
	GeneratedAt            time.Time      `json:"generated_at"`
	ExpiresAt              time.Time      `json:"expires_at"`
	AdmittedIPv4           []addressRange `json:"admitted_ipv4"`
	ExceptionIPv4          []addressRange `json:"exception_ipv4"`
	AdmittedIntervalCount  int            `json:"admitted_interval_count"`
	AdmittedAddressCount   uint64         `json:"admitted_address_count"`
	ExceptionIntervalCount int            `json:"exception_interval_count"`
	ExceptionAddressCount  uint64         `json:"exception_address_count"`
}

type parsedRange struct {
	start uint32
	end   uint32
}

type validatedBundle struct {
	bundle
	admitted   []parsedRange
	exceptions []parsedRange
}

type updateConfig struct {
	bucket                    string
	object                    string
	publicInterface           string
	stateDirectory            string
	ociCLI                    string
	nft                       string
	minimumAddresses          uint64
	maximumAddresses          uint64
	minimumIntervals          int
	maximumExceptionAddresses uint64
	maximumExceptionIntervals int
	staticExceptions          []string
}

type stringValues []string

func (values *stringValues) String() string {
	return strings.Join(*values, ",")
}

func (values *stringValues) Set(value string) error {
	*values = append(*values, value)
	return nil
}

func addUpdateFlags(flags *flag.FlagSet, configuration *updateConfig) {
	flags.StringVar(&configuration.publicInterface, "public-interface", "", "public network interface")
	flags.StringVar(&configuration.stateDirectory, "state-directory", "/var/lib/oci-l4-relay/web-allow", "state directory")
	flags.StringVar(&configuration.nft, "nft", "/usr/sbin/nft", "nft path")
	flags.Uint64Var(&configuration.minimumAddresses, "minimum-admitted-addresses", 1, "minimum admitted IPv4 address count")
	flags.Uint64Var(&configuration.maximumAddresses, "maximum-admitted-addresses", 1<<32, "maximum admitted IPv4 address count")
	flags.IntVar(&configuration.minimumIntervals, "minimum-admitted-intervals", 1, "minimum admitted interval count")
	flags.Uint64Var(&configuration.maximumExceptionAddresses, "maximum-exception-addresses", 1<<32, "maximum exception IPv4 address count")
	flags.IntVar(&configuration.maximumExceptionIntervals, "maximum-exception-intervals", maxExceptions, "maximum exception interval count")
	flags.Var((*stringValues)(&configuration.staticExceptions), "static-exception-ipv4", "static IPv4 exception")
}

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 main() {
	if len(os.Args) < 2 {
		fatal(errors.New("expected apply, update, restore, validate, or render"))
	}
	switch os.Args[1] {
	case "apply":
		flags := flag.NewFlagSet("apply", flag.ExitOnError)
		configuration := updateConfig{}
		input := flags.String("input", "", "downloaded bundle path")
		addUpdateFlags(flags, &configuration)
		_ = flags.Parse(os.Args[2:])
		now := time.Now().UTC()
		if *input == "" {
			fatal(errors.New("input is required"))
		}
		if err := activateBundle(now, configuration, *input, true); err != nil {
			fmt.Fprintln(os.Stderr, err)
			if restoreErr := restore(now, configuration); restoreErr != nil {
				if failClosedErr := failClosed(configuration); failClosedErr != nil {
					fatal(errors.Join(restoreErr, failClosedErr))
				}
				fatal(restoreErr)
			}
			fmt.Println("local_apply=false cache_restored=true")
		}
	case "update":
		flags := flag.NewFlagSet("update", flag.ExitOnError)
		configuration := updateConfig{}
		flags.StringVar(&configuration.bucket, "bucket", "", "OCI Object Storage bucket")
		flags.StringVar(&configuration.object, "object", "", "OCI Object Storage object")
		flags.StringVar(&configuration.ociCLI, "oci-cli", "/opt/oci-cli/bin/oci", "OCI CLI path")
		addUpdateFlags(flags, &configuration)
		_ = flags.Parse(os.Args[2:])
		now := time.Now().UTC()
		if err := update(now, configuration); err != nil {
			fmt.Fprintln(os.Stderr, err)
			if restoreErr := restore(now, configuration); restoreErr != nil {
				if failClosedErr := failClosed(configuration); failClosedErr != nil {
					fatal(errors.Join(restoreErr, failClosedErr))
				}
				fatal(restoreErr)
			}
			fmt.Println("remote_update=false cache_restored=true")
		}
	case "restore":
		flags := flag.NewFlagSet("restore", flag.ExitOnError)
		configuration := updateConfig{}
		addUpdateFlags(flags, &configuration)
		_ = flags.Parse(os.Args[2:])
		if err := restore(time.Now().UTC(), configuration); err != nil {
			if failClosedErr := failClosed(configuration); failClosedErr != nil {
				fatal(errors.Join(err, failClosedErr))
			}
			fatal(err)
		}
	case "validate", "render":
		flags := flag.NewFlagSet(os.Args[1], flag.ExitOnError)
		input := flags.String("input", "", "bundle path")
		publicInterface := flags.String("public-interface", "", "public network interface")
		minimumAddresses := flags.Uint64("minimum-admitted-addresses", 1, "minimum admitted IPv4 address count")
		maximumAddresses := flags.Uint64("maximum-admitted-addresses", 1<<32, "maximum admitted IPv4 address count")
		minimumIntervals := flags.Int("minimum-admitted-intervals", 1, "minimum admitted interval count")
		maximumExceptionAddresses := flags.Uint64("maximum-exception-addresses", 1<<32, "maximum exception IPv4 address count")
		maximumExceptionIntervals := flags.Int("maximum-exception-intervals", maxExceptions, "maximum exception interval count")
		summaryOutput := flags.String("summary-output", "", "write validated bundle summary")
		_ = flags.Parse(os.Args[2:])
		validated, err := readBundle(time.Now().UTC(), *input)
		if err != nil {
			fatal(err)
		}
		if err := validatePolicy(updateConfig{
			minimumAddresses:          *minimumAddresses,
			maximumAddresses:          *maximumAddresses,
			minimumIntervals:          *minimumIntervals,
			maximumExceptionAddresses: *maximumExceptionAddresses,
			maximumExceptionIntervals: *maximumExceptionIntervals,
		}, validated); err != nil {
			fatal(err)
		}
		if *summaryOutput != "" {
			summary := struct {
				GeneratedAt            time.Time `json:"generated_at"`
				ExpiresAt              time.Time `json:"expires_at"`
				AdmittedIntervalCount  int       `json:"admitted_interval_count"`
				AdmittedAddressCount   uint64    `json:"admitted_address_count"`
				ExceptionIntervalCount int       `json:"exception_interval_count"`
				ExceptionAddressCount  uint64    `json:"exception_address_count"`
			}{
				GeneratedAt:            validated.GeneratedAt,
				ExpiresAt:              validated.ExpiresAt,
				AdmittedIntervalCount:  validated.AdmittedIntervalCount,
				AdmittedAddressCount:   validated.AdmittedAddressCount,
				ExceptionIntervalCount: validated.ExceptionIntervalCount,
				ExceptionAddressCount:  validated.ExceptionAddressCount,
			}
			if err := writeJSONFile(*summaryOutput, summary); err != nil {
				fatal(err)
			}
		}
		if os.Args[1] == "render" {
			if err := render(os.Stdout, validated, *publicInterface); err != nil {
				fatal(err)
			}
		}
	default:
		fatal(fmt.Errorf("unknown command %q", os.Args[1]))
	}
}

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

func update(now time.Time, configuration updateConfig) error {
	if configuration.bucket == "" || configuration.object == "" || configuration.publicInterface == "" {
		return errors.New("bucket, object, and public-interface are required")
	}
	if err := os.MkdirAll(configuration.stateDirectory, 0700); err != nil {
		return err
	}
	temporaryDirectory, err := os.MkdirTemp(configuration.stateDirectory, ".update-")
	if err != nil {
		return err
	}
	defer func() { _ = os.RemoveAll(temporaryDirectory) }()
	bundlePath := filepath.Join(temporaryDirectory, "bundle.json.gz")
	command := exec.Command(configuration.ociCLI, "os", "object", "get", "--auth", "instance_principal", "--bucket-name", configuration.bucket, "--name", configuration.object, "--file", bundlePath)
	command.Stdout = os.Stdout
	command.Stderr = os.Stderr
	if err := command.Run(); err != nil {
		return fmt.Errorf("fetch address bundle: %w", err)
	}
	return activateBundle(now, configuration, bundlePath, true)
}

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

func activateBundle(now time.Time, configuration updateConfig, bundlePath string, persist bool) error {
	validated, err := readBundle(now, bundlePath)
	if err != nil {
		return err
	}
	if err := validatePolicy(configuration, validated); err != nil {
		return err
	}
	validated, err = addStaticExceptions(validated, configuration.staticExceptions)
	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, "admission.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, 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 := replaceFile(rulesetPath, filepath.Join(configuration.stateDirectory, "last-good.nft")); err != nil {
			return err
		}
	}
	fmt.Printf("admitted_intervals=%d admitted_addresses=%d exception_intervals=%d exception_addresses=%d expires_at=%s\n", len(validated.admitted), validated.AdmittedAddressCount, len(validated.exceptions), validated.ExceptionAddressCount, validated.ExpiresAt.Format(time.RFC3339))
	return nil
}

func addStaticExceptions(document validatedBundle, values []string) (validatedBundle, error) {
	exceptions := append([]parsedRange(nil), document.exceptions...)
	for _, value := range values {
		address, err := parseIPv4(value)
		if err != nil {
			return validatedBundle{}, fmt.Errorf("static exception: %w", err)
		}
		exceptions = append(exceptions, parsedRange{start: address, end: address})
	}
	sort.Slice(exceptions, func(left, right int) bool {
		return exceptions[left].start < exceptions[right].start
	})
	merged := make([]parsedRange, 0, len(exceptions))
	for _, candidate := range exceptions {
		if len(merged) == 0 || uint64(candidate.start) > uint64(merged[len(merged)-1].end)+1 {
			merged = append(merged, candidate)
			continue
		}
		if candidate.end > merged[len(merged)-1].end {
			merged[len(merged)-1].end = candidate.end
		}
	}
	if len(merged) > maxExceptions {
		return validatedBundle{}, errors.New("effective exception interval count exceeds limit")
	}
	document.exceptions = merged
	return document, nil
}

func validatePolicy(configuration updateConfig, validated validatedBundle) error {
	if validated.AdmittedAddressCount < configuration.minimumAddresses || validated.AdmittedAddressCount > configuration.maximumAddresses {
		return errors.New("admitted address count is outside configured policy bounds")
	}
	if len(validated.admitted) < configuration.minimumIntervals {
		return errors.New("admitted interval count is below the configured policy bound")
	}
	if validated.ExceptionAddressCount > configuration.maximumExceptionAddresses {
		return errors.New("exception address count exceeds the configured policy bound")
	}
	if len(validated.exceptions) > configuration.maximumExceptionIntervals {
		return errors.New("exception interval count exceeds the configured policy bound")
	}
	return nil
}

func failClosed(configuration updateConfig) error {
	if configuration.publicInterface == "" {
		return errors.New("public-interface is required")
	}
	if err := ensureTable(configuration.nft); err != nil {
		return err
	}
	ruleset, err := os.CreateTemp("", "oci-l4-relay-fail-closed-*.nft")
	if err != nil {
		return err
	}
	path := ruleset.Name()
	defer func() { _ = os.Remove(path) }()
	if err := renderFailClosed(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-closed 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-closed ruleset: %w: %s", err, strings.TrimSpace(string(output)))
	}
	fmt.Println("web_allow_fail_closed=true")
	return 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 admission 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, path string) (validatedBundle, error) {
	file, err := os.Open(path)
	if err != nil {
		return validatedBundle{}, err
	}
	defer func() { _ = file.Close() }()
	reader := io.LimitReader(file, maxBundleBytes+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, maxBundleBytes+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, document)
}

func validateBundle(now time.Time, document bundle) (validatedBundle, error) {
	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.Before(now.Add(minimumLifetime)) {
		return validatedBundle{}, errors.New("bundle expires too soon")
	}
	if !document.ExpiresAt.After(document.GeneratedAt) {
		return validatedBundle{}, errors.New("bundle expiry does not follow generation")
	}
	if document.ExpiresAt.Sub(document.GeneratedAt) > maximumLifetime || document.GeneratedAt.Before(now.Add(-maximumLifetime)) {
		return validatedBundle{}, errors.New("bundle lifetime or generation age exceeds bounds")
	}
	if len(document.AdmittedIPv4) == 0 || len(document.AdmittedIPv4) > maxIntervals {
		return validatedBundle{}, errors.New("admitted interval count is outside bounds")
	}
	if len(document.ExceptionIPv4) > maxExceptions {
		return validatedBundle{}, errors.New("exception interval count is outside bounds")
	}
	admitted, admittedAddresses, err := validateRanges(document.AdmittedIPv4)
	if err != nil {
		return validatedBundle{}, fmt.Errorf("admitted_ipv4: %w", err)
	}
	exceptions, exceptionAddresses, err := validateRanges(document.ExceptionIPv4)
	if err != nil {
		return validatedBundle{}, fmt.Errorf("exception_ipv4: %w", err)
	}
	if document.AdmittedIntervalCount != len(admitted) || document.AdmittedAddressCount != admittedAddresses {
		return validatedBundle{}, errors.New("admitted statistics do not match payload")
	}
	if document.ExceptionIntervalCount != len(exceptions) || document.ExceptionAddressCount != exceptionAddresses {
		return validatedBundle{}, errors.New("exception statistics do not match payload")
	}
	return validatedBundle{bundle: document, admitted: admitted, exceptions: exceptions}, nil
}

func validateRanges(values []addressRange) ([]parsedRange, uint64, error) {
	parsed := make([]parsedRange, 0, len(values))
	var count uint64
	var previousEnd uint32
	for index, value := range values {
		start, err := parseIPv4(value.Start)
		if err != nil {
			return nil, 0, fmt.Errorf("range %d start: %w", index, err)
		}
		end, err := parseIPv4(value.End)
		if err != nil {
			return nil, 0, fmt.Errorf("range %d end: %w", index, err)
		}
		if start > end {
			return nil, 0, fmt.Errorf("range %d is reversed", index)
		}
		if index > 0 && start <= previousEnd {
			return nil, 0, fmt.Errorf("range %d overlaps or is unsorted", index)
		}
		parsed = append(parsed, parsedRange{start: start, end: end})
		count += uint64(end) - uint64(start) + 1
		previousEnd = end
	}
	return parsed, count, 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, document validatedBundle, publicInterface string) error {
	if publicInterface == "" || strings.ContainsAny(publicInterface, "\"\\\n\r\t ") {
		return errors.New("invalid public interface")
	}
	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 := renderSet(writer, "admitted_v4", document.admitted); err != nil {
		return err
	}
	if err := renderSet(writer, "exception_v4", document.exceptions); err != nil {
		return err
	}
	if _, err := fmt.Fprintf(writer, "    counter web_allow_denied {}\n    chain ingress {\n        type filter hook prerouting priority raw - 20; policy accept;\n        iifname %s tcp dport { 80, 443, 465, 993 } ip saddr @exception_v4 accept\n        iifname %s tcp dport { 80, 443, 465, 993 } ip saddr != @admitted_v4 counter name web_allow_denied drop\n    }\n}\n", strconv.Quote(publicInterface), strconv.Quote(publicInterface)); err != nil {
		return err
	}
	return writer.Flush()
}

func renderFailClosed(output io.Writer, publicInterface string) error {
	if publicInterface == "" || strings.ContainsAny(publicInterface, "\"\\\n\r\t ") {
		return errors.New("invalid public interface")
	}
	_, err := fmt.Fprintf(output, "delete table inet %s\ntable inet %s {\n    counter web_allow_denied {}\n    chain ingress {\n        type filter hook prerouting priority raw - 20; policy accept;\n        iifname %s tcp dport { 80, 443, 465, 993 } counter name web_allow_denied drop\n    }\n}\n", tableName, tableName, strconv.Quote(publicInterface))
	return err
}

func renderSet(writer *bufio.Writer, name string, ranges []parsedRange) error {
	if _, err := fmt.Fprintf(writer, "    set %s {\n        type ipv4_addr\n        flags interval\n        auto-merge\n", name); err != nil {
		return err
	}
	if len(ranges) > 0 {
		if _, err := fmt.Fprintln(writer, "        elements = {"); err != nil {
			return err
		}
		for index, item := range ranges {
			value := ipv4String(item.start)
			if item.start != item.end {
				value += "-" + ipv4String(item.end)
			}
			separator := ","
			if index == len(ranges)-1 {
				separator = ""
			}
			if _, err := fmt.Fprintf(writer, "            %s%s\n", value, separator); err != nil {
				return err
			}
		}
		if _, err := fmt.Fprintln(writer, "        }"); err != nil {
			return err
		}
	}
	_, err := fmt.Fprintln(writer, "    }")
	return err
}