firewall/nft.py

from dataclasses import dataclass, field
from itertools import chain
from typing import Any, TypeVar

T = TypeVar("T")
JsonNftables = dict[str, Any]


def flatten(l: list[list[T]]) -> list[T]:
    return list(chain.from_iterable(l))


@dataclass
class Set:
    name: str

    flags: list[str] | None = None
    type: str | list[str] | None = None

    def to_nft(self, family: str, table: str) -> JsonNftables:
        set: JsonNftables = {
            "name": self.name,
            "family": family,
            "table": table,
        }

        if self.flags is not None:
            set["flags"] = self.flags

        if self.type is not None:
            set["type"] = self.type

        return {"add": {"set": set}}


@dataclass
class Immediate:
    value: str

    def to_nft(self) -> str:
        return self.value


@dataclass
class Payload:
    protocol: str
    field: str

    def to_nft(self) -> JsonNftables:
        return {"payload": {"protocol": self.protocol, "field": self.field}}


Expression = Immediate | Payload


@dataclass
class Verdict:
    verdict: str

    target: str | None = None

    def to_nft(self) -> JsonNftables:
        return {self.verdict: self.target}


@dataclass
class Match:
    op: str
    left: Expression
    right: Expression

    def to_nft(self) -> JsonNftables:
        return {
            "match": {
                "op": self.op,
                "left": self.left.to_nft(),
                "right": self.right.to_nft(),
            }
        }


Statement = Verdict | Match


@dataclass
class Rule:
    stmts: list[Statement]

    def to_nft(self, family: str, table: str, chain: str) -> JsonNftables:
        return {
            "add": {
                "rule": {
                    "family": family,
                    "table": table,
                    "chain": chain,
                    "expr": [stmt.to_nft() for stmt in self.stmts],
                }
            }
        }


@dataclass
class Chain:
    name: str

    type: str | None = None
    hook: str | None = None
    priority: int | None = None
    policy: str | None = None

    rules: list[Rule] = field(default_factory=list)

    def to_nft(self, family: str, table: str) -> list[JsonNftables]:
        chain: JsonNftables = {
            "name": self.name,
            "family": family,
            "table": table,
        }

        if self.type is not None:
            chain["type"] = self.type

        if self.hook is not None:
            chain["hook"] = self.hook

        if self.priority is not None:
            chain["prio"] = self.priority

        if self.policy is not None:
            chain["policy"] = self.policy

        commands = [{"add": {"chain": chain}}]

        for rule in self.rules:
            commands.append(rule.to_nft(family, table, self.name))

        return commands


@dataclass
class Table:
    family: str
    name: str

    chains: list[Chain] = field(default_factory=list)
    sets: list[Set] = field(default_factory=list)

    def to_nft(self) -> list[JsonNftables]:
        commands = [
            {"add": {"table": {"family": self.family, "name": self.name}}}
        ]

        for set in self.sets:
            commands.append(set.to_nft(self.family, self.name))

        for chain in self.chains:
            commands.extend(chain.to_nft(self.family, self.name))

        return commands


@dataclass
class Ruleset:
    flush: bool

    tables: list[Table] = field(default_factory=list)

    def to_nft(self) -> JsonNftables:
        ruleset = flatten([table.to_nft() for table in self.tables])

        if self.flush:
            ruleset.insert(0, {"flush": {"ruleset": None}})

        return {"nftables": ruleset}