"""Tests for abstract nftables ruleset model."""

from gaol.networking.nftables import (
    NftAction,
    NftChain,
    NftChainType,
    NftFamily,
    NftHook,
    NftRule,
    NftRuleset,
    chain_forwarding_ruleset,
    kill_switch_rules,
    kill_switch_ruleset,
)


class TestNftRule:
    def test_simple_accept(self):
        rule = NftRule(oif="lo", action=NftAction.ACCEPT)
        assert rule.render() == 'oif "lo" accept'

    def test_full_rule(self):
        rule = NftRule(
            iif="ve-research",
            protocol="tcp",
            dst_addr="10.200.1.0/30",
            dst_port=443,
            action=NftAction.ACCEPT,
            comment="allow HTTPS",
        )
        rendered = rule.render()
        assert 'iif "ve-research"' in rendered
        assert "tcp" in rendered
        assert "ip daddr 10.200.1.0/30" in rendered
        assert "dport 443" in rendered
        assert "accept" in rendered
        assert 'comment "allow HTTPS"' in rendered

    def test_masquerade(self):
        rule = NftRule(src_addr="10.200.1.0/30", action=NftAction.MASQUERADE)
        assert rule.render() == "ip saddr 10.200.1.0/30 masquerade"

    def test_ct_state(self):
        rule = NftRule(ct_state="established,related", action=NftAction.ACCEPT)
        assert rule.render() == "ct state established,related accept"

    def test_dnat(self):
        rule = NftRule(
            protocol="tcp",
            dst_port=8080,
            action=NftAction.DNAT,
            dnat_target="10.0.0.2:80",
        )
        assert "dnat to 10.0.0.2:80" in rule.render()

    def test_log_prefix(self):
        rule = NftRule(action=NftAction.LOG, log_prefix="gaol: ")
        assert rule.render() == 'log prefix "gaol: "'


class TestNftChain:
    def test_render_filter_chain(self):
        chain = NftChain(
            name="output",
            chain_type=NftChainType.FILTER,
            hook=NftHook.OUTPUT,
            priority=0,
            policy="drop",
            rules=[
                NftRule(oif="lo", action=NftAction.ACCEPT, priority=10),
                NftRule(ct_state="established,related", action=NftAction.ACCEPT, priority=20),
            ],
        )
        rendered = chain.render()
        assert "chain output {" in rendered
        assert "type filter hook output priority 0; policy drop;" in rendered
        assert 'oif "lo" accept' in rendered
        assert "ct state established,related accept" in rendered

    def test_rules_sorted_by_priority(self):
        chain = NftChain(
            name="forward",
            chain_type=NftChainType.FILTER,
            hook=NftHook.FORWARD,
            rules=[
                NftRule(oif="lo", action=NftAction.DROP, priority=100),
                NftRule(oif="lo", action=NftAction.ACCEPT, priority=10),
            ],
        )
        rendered = chain.render()
        lines = rendered.splitlines()
        accept_idx = next(i for i, l in enumerate(lines) if "accept" in l)
        drop_idx = next(i for i, l in enumerate(lines) if "drop" in l)
        assert accept_idx < drop_idx


class TestNftRuleset:
    def test_render_complete_ruleset(self):
        ruleset = NftRuleset(table_name="gaol_test", family=NftFamily.IP)
        ruleset.add_chain(
            NftChain(
                name="output",
                chain_type=NftChainType.FILTER,
                hook=NftHook.OUTPUT,
                policy="drop",
                rules=[
                    NftRule(oif="lo", action=NftAction.ACCEPT),
                ],
            )
        )
        rendered = ruleset.render()
        assert "table ip gaol_test {" in rendered
        assert "chain output {" in rendered
        assert 'oif "lo" accept' in rendered

    def test_render_starts_with_comment(self):
        ruleset = NftRuleset(table_name="gaol_test")
        rendered = ruleset.render()
        assert rendered.startswith("# Auto-generated by gaol")

    def test_render_includes_table_reference(self):
        ruleset = NftRuleset(table_name="gaol_research", family=NftFamily.IP)
        rendered = ruleset.render()
        assert "# Table: ip gaol_research" in rendered


class TestArchetypeGenerators:
    def test_kill_switch_rules(self):
        rules = kill_switch_rules("10.200.1.1", allowed_ips=["10.200.2.2"])
        assert len(rules) == 4  # loopback, gateway, established, allowed IP
        actions = [r.action for r in rules]
        assert all(a == NftAction.ACCEPT for a in actions)
        assert rules[0].oif == "lo"
        assert rules[1].dst_addr == "10.200.1.1"
        assert rules[3].dst_addr == "10.200.2.2"

    def test_kill_switch_ruleset(self):
        ruleset = kill_switch_ruleset("gaol_test_ks", "10.200.1.1")
        rendered = ruleset.render()
        assert "policy drop" in rendered
        assert 'oif "lo" accept' in rendered
        assert "ip daddr 10.200.1.1 accept" in rendered

    def test_chain_forwarding_ruleset(self):
        ruleset = chain_forwarding_ruleset(
            "gaol_test_chain", "ve-research", "10.200.1.0/30"
        )
        rendered = ruleset.render()
        assert "chain forward {" in rendered
        assert "chain postrouting {" in rendered
        assert 'iif "ve-research" accept' in rendered
        assert "ip saddr 10.200.1.0/30 masquerade" in rendered
        assert "policy drop" in rendered  # forward chain should default drop

    def test_chain_forwarding_has_established_return(self):
        ruleset = chain_forwarding_ruleset(
            "gaol_test", "ve-test", "10.200.7.0/30"
        )
        rendered = ruleset.render()
        assert "ct state established,related accept" in rendered
