package main

import (
	"context"
	"encoding/base64"
	"encoding/json"
	"os"
	"os/exec"
	"path/filepath"
	"strings"
	"testing"
)

func TestDeriveWireGuardKey(t *testing.T) {
	privateKey, publicKey, err := deriveWireGuardKey([]byte("stable seed"))
	if err != nil {
		t.Fatal(err)
	}
	if got := base64.StdEncoding.EncodeToString(privateKey); got != "WBlFzkJLD63HWvvdo3EpeSDetKKqWSlTFpSi/I8K93o=" {
		t.Fatalf("private key = %q", got)
	}
	if got := base64.StdEncoding.EncodeToString(publicKey); got != "20/HiRcrrKinQNH2WEJk8QcuB+3fv9uOW8Ptofn3CCA=" {
		t.Fatalf("public key = %q", got)
	}
}

func TestAdmissionRefreshCommandIsSupported(t *testing.T) {
	if err := runAdmissionOperation(context.TODO(), validConfig(), "invalid"); err == nil {
		t.Fatal("unknown admission operation was accepted")
	}
}

func TestRenderRuleset(t *testing.T) {
	configuration := validConfig()
	ruleset := renderRuleset(configuration, "ens3")
	for _, expected := range []string{
		`iifname "ens3" tcp dport { 80, 443, 465, 993 } counter name web_allow_denied drop`,
		`iifname "ens3" tcp flags syn tcp dport { 80, 443 } ip saddr @syn_quarantine_v4`,
		"set syn_quarantine_v4 {\n    type ipv4_addr\n    flags dynamic,timeout\n    counter\n    timeout 1m",
		`iifname "wg-oci0" tcp dport 19100 accept`,
		`limit rate over 1000/second burst 1000 packets counter name web_syn_aggregate_dropped`,
		`limit rate over 20000/second burst 20000 packets counter name web_packet_aggregate_dropped`,
		`dnat ip to 10.104.40.127`,
	} {
		if !strings.Contains(ruleset, expected) {
			t.Errorf("ruleset does not contain %q", expected)
		}
	}
	if strings.Contains(ruleset, "587") {
		t.Fatal("ruleset exposes cluster-local submission")
	}
	if strings.Index(ruleset, "ip saddr @syn_quarantine_v4") > strings.Index(ruleset, "limit rate over 1000/second") {
		t.Error("per-source quarantine must precede the aggregate SYN ceiling")
	}
	if strings.Contains(ruleset, "10.129.13.87") || strings.Contains(ruleset, "dport { 25, 465, 993 }") {
		t.Fatal("mail forwarding is present while mail configuration is inactive")
	}
}

func TestRenderRulesetMapsOnlyDeclaredMailPorts(t *testing.T) {
	configuration := validConfig()
	configuration.MailRelayAddress = "10.129.13.87"
	ruleset := renderRuleset(configuration, "ens3")
	for _, expected := range []string{
		`iifname "ens3" oifname "wg-oci0" ip daddr 10.104.40.127 tcp dport { 80, 443 } accept`,
		`iifname "ens3" oifname "wg-oci0" ip daddr 10.129.13.87 tcp dport { 25, 465, 993 } accept`,
		`iifname "ens3" tcp dport { 80, 443 } dnat ip to 10.104.40.127`,
		`iifname "ens3" tcp dport { 25, 465, 993 } dnat ip to 10.129.13.87`,
		`tcp dport 25 meter smtp_syn_by_source`,
		"set smtp_syn_quarantine_v4 {\n    type ipv4_addr\n    flags dynamic,timeout\n    counter\n    timeout 600s\n    size 65536",
		`tcp dport 25 limit rate over 500/second burst 500 packets counter name smtp_syn_aggregate_dropped`,
		`tcp dport { 465, 993 } meter auth_syn_by_source`,
		`tcp dport { 465, 993 } limit rate over 200/second burst 200 packets counter name auth_syn_aggregate_dropped`,
	} {
		if !strings.Contains(ruleset, expected) {
			t.Errorf("ruleset does not contain %q", expected)
		}
	}
	for _, excluded := range []string{"snat", "masquerade", "dport 80 dnat ip to 10.129.13.87", "dport 25 dnat ip to 10.104.40.127"} {
		if strings.Contains(strings.ToLower(ruleset), excluded) {
			t.Errorf("ruleset contains %q", excluded)
		}
	}
}

func TestRulesetLoadsWithIndependentMailBudgets(t *testing.T) {
	if os.Getenv("RELAY_NFT_TEST") != "1" {
		t.Skip("RELAY_NFT_TEST is not enabled")
	}
	configuration := validConfig()
	configuration.MailRelayAddress = "10.129.13.87"
	path := filepath.Join(t.TempDir(), "ruleset.nft")
	if err := os.WriteFile(path, []byte(renderRuleset(configuration, "ens3")), 0600); err != nil {
		t.Fatal(err)
	}
	if output, err := exec.Command("nft", "--check", "--file", path).CombinedOutput(); err != nil {
		t.Fatalf("nft check: %v: %s", err, output)
	}
	if output, err := exec.Command("nft", "--file", path).CombinedOutput(); err != nil {
		t.Fatalf("nft apply: %v: %s", err, output)
	}
	output, err := exec.Command("nft", "list", "chain", "inet", "oci_l4_relay", "ingress_guard").CombinedOutput()
	if err != nil {
		t.Fatalf("nft list: %v: %s", err, output)
	}
	ruleset := string(output)
	for _, expected := range []string{
		"smtp_syn_by_source",
		"auth_syn_by_source",
		"web_syn_by_source",
		"smtp_packet_aggregate_dropped",
		"auth_packet_aggregate_dropped",
		"web_packet_aggregate_dropped",
	} {
		if !strings.Contains(ruleset, expected) {
			t.Fatalf("active ruleset does not contain %q", expected)
		}
	}
}

func TestTrafficControlCommandsSeparateServiceClasses(t *testing.T) {
	configuration := validConfig()
	configuration.MailRelayAddress = "10.129.13.87"
	commands := trafficControlCommands(configuration)
	joined := make([]string, 0, len(commands))
	for _, command := range commands {
		joined = append(joined, strings.Join(command, " "))
	}
	output := strings.Join(joined, "\n")
	for _, expected := range []string{
		"classid 1:10 htb rate 60mbit",
		"classid 1:20 htb rate 15mbit",
		"classid 1:30 htb rate 30mbit",
		"dst_ip 10.104.40.127 ip_proto tcp dst_port 443 classid 1:10",
		"dst_ip 10.129.13.87 ip_proto tcp dst_port 25 classid 1:20",
		"dst_ip 10.129.13.87 ip_proto tcp dst_port 993 classid 1:30",
	} {
		if !strings.Contains(output, expected) {
			t.Fatalf("traffic-control commands do not contain %q", expected)
		}
	}
}

func TestTrafficControlActivation(t *testing.T) {
	if os.Getenv("RELAY_NFT_TEST") != "1" {
		t.Skip("RELAY_NFT_TEST is not enabled")
	}
	if output, err := exec.Command("ip", "link", "add", interfaceName, "type", "dummy").CombinedOutput(); err != nil {
		t.Fatalf("create test interface: %v: %s", err, output)
	}
	configuration := validConfig()
	configuration.MailRelayAddress = "10.129.13.87"
	for attempt := 0; attempt < 2; attempt++ {
		clearTrafficControl(context.Background())
		for _, command := range trafficControlCommands(configuration) {
			if output, err := exec.Command(command[0], command[1:]...).CombinedOutput(); err != nil {
				t.Fatalf("attempt %d: %s: %v: %s", attempt+1, strings.Join(command, " "), err, output)
			}
		}
	}
	output, err := exec.Command("tc", "class", "show", "dev", interfaceName).CombinedOutput()
	if err != nil {
		t.Fatalf("list traffic-control classes: %v: %s", err, output)
	}
	for _, expected := range []string{"1:10", "1:20", "1:30", "1:40"} {
		if !strings.Contains(string(output), expected) {
			t.Fatalf("active traffic-control classes do not contain %s: %s", expected, output)
		}
	}
}

func TestLoadConfigRejectsUnknownFields(t *testing.T) {
	path := filepath.Join(t.TempDir(), "config.json")
	if err := os.WriteFile(path, []byte(`{"unknown":true}`), 0600); err != nil {
		t.Fatal(err)
	}
	if _, err := loadConfig(path); err == nil {
		t.Fatal("loadConfig accepted an unknown field")
	}
}

func TestMalformedMailConfigurationDisablesOnlyMail(t *testing.T) {
	configuration := validConfig()
	configuration.MailRelayAddress = "not-an-address"
	configuration.SMTPDeny = &smtpDenyConfig{}
	contents, err := json.Marshal(configuration)
	if err != nil {
		t.Fatal(err)
	}
	path := filepath.Join(t.TempDir(), "config.json")
	if err := os.WriteFile(path, contents, 0600); err != nil {
		t.Fatal(err)
	}
	loaded, err := loadConfig(path)
	if err != nil {
		t.Fatal(err)
	}
	if loaded.MailRelayAddress != "" || loaded.SMTPDeny != nil || loaded.mailConfigurationError == "" {
		t.Fatal("malformed mail configuration remained active")
	}
	if loaded.GatewayRelayAddress != configuration.GatewayRelayAddress {
		t.Fatal("malformed mail configuration changed the web relay")
	}
}

func TestValidMailConfiguration(t *testing.T) {
	configuration := validConfig()
	configuration.MailRelayAddress = "10.129.13.87"
	configuration.SMTPDeny = &smtpDenyConfig{
		Bucket:                   "bucket",
		MaximumAddresses:         4096,
		MaximumIntervals:         1024,
		MaximumLifetimeSeconds:   86400,
		MaximumObservationWindow: 600,
		MinimumObservations:      20,
		Namespace:                "namespace",
		Object:                   "smtp-deny/v1/current.json.gz",
	}
	configuration.SMTPDenyHelper = "/usr/local/libexec/oci-l4-relay-smtp-deny"
	if err := configuration.validateMail(); err != nil {
		t.Fatal(err)
	}
}

func validConfig() config {
	return config{
		Admission: admissionConfig{
			Bucket:           "bucket",
			MaximumAddresses: 3500000000,
			MinimumAddresses: 1000000000,
			MinimumIntervals: 10000,
			Namespace:        "namespace",
			Object:           "web-allow/v1/current.json.gz",
		},
		AdmissionHelper:               "/usr/local/libexec/oci-l4-relay-admission",
		GatewayRelayAddress:           "10.104.40.127",
		HomeWireGuardPublicKey:        "key",
		LogIngestHostname:             "relay-logs.quorate.dev",
		LogIngestSecretOCID:           "log-secret",
		MonitoringSourceIPv4:          []string{"192.0.2.1", "192.0.2.2", "192.0.2.3"},
		PublicKeySecretOCID:           "public-secret",
		RateLimitMbit:                 150,
		RelayHTTPQuarantineEnabled:    true,
		RelayHTTPQuarantineMinutes:    30,
		RelayHTTPSRecentMinutes:       10,
		RelayMetricsPort:              19100,
		RelayAuthPacketAggregateBurst: 5000,
		RelayAuthPacketAggregateRate:  5000,
		RelayAuthSYNAggregateBurst:    200,
		RelayAuthSYNAggregateRate:     200,
		RelayAuthSYNBurst:             15,
		RelayAuthSYNRate:              30,
		RelaySMTPPacketAggregateBurst: 10000,
		RelaySMTPPacketAggregateRate:  10000,
		RelaySMTPSYNAggregateBurst:    500,
		RelaySMTPSYNAggregateRate:     500,
		RelaySMTPSYNBurst:             20,
		RelaySMTPSYNRate:              60,
		RelayTrackerMaxSources:        65536,
		RelayTunnelAddress:            "10.207.54.68/30",
		RelayTunnelMTU:                1420,
		RelayUnmatchedHTTPBurst:       5,
		RelayUnmatchedHTTPRate:        30,
		RelayWebPacketAggregateBurst:  20000,
		RelayWebPacketAggregateRate:   20000,
		RelayWebSYNAggregateBurst:     1000,
		RelayWebSYNAggregateRate:      1000,
		RelayWebSYNBurst:              100,
		RelayWebSYNRate:               120,
		WireGuardListenPort:           51820,
		WireGuardSeedSecretOCID:       "seed-secret",
	}
}

func TestValidLogIngestToken(t *testing.T) {
	for _, test := range []struct {
		name  string
		token string
		valid bool
	}{
		{name: "valid", token: strings.Repeat("a", 48), valid: true},
		{name: "short", token: strings.Repeat("a", 31)},
		{name: "newline", token: strings.Repeat("a", 47) + "\n"},
		{name: "space", token: strings.Repeat("a", 47) + " "},
	} {
		t.Run(test.name, func(t *testing.T) {
			if got := validLogIngestToken([]byte(test.token)); got != test.valid {
				t.Fatalf("validLogIngestToken() = %t", got)
			}
		})
	}
}