firewall/nftables.py

#!/usr/bin/env python3

from argparse import ArgumentParser, FileType
from dataclasses import dataclass
from enum import Enum
from graphlib import TopologicalSorter
from pydantic import (
    BaseModel,
    Extra,
    FilePath,
    IPvAnyAddress,
    IPvAnyNetwork,
    conint,
    parse_obj_as,
    validator,
    root_validator,
)
from yaml import safe_load


class RestrictiveBaseModel(BaseModel, extra=Extra.forbid):
    pass


# Ports
Port = conint(ge=0, le=2**16)


class PortRange(str):
    @classmethod
    def __get_validators__(cls):
        yield cls.validate

    @classmethod
    def validate(cls, v):
        try:
            start, end = v.split("..")
        except ValueError:
            raise ValueError("invalid port range: must be in the form start..end")

        start, end = parse_obj_as(Port, start), parse_obj_as(Port, end)
        if start > end:
            raise ValueError("invalid port range: start must be less than end")

        return range(start, end)


# ===== First pass: Zones =====


# Zones
class ZoneName(str):
    pass


@dataclass
class Zone:
    addrs: set[IPvAnyNetwork]
    negate: bool


# Zones: Parsing YAML
class ZoneYAML(RestrictiveBaseModel):
    addrs: set[IPvAnyNetwork] = set()
    files: set[FilePath] = set()
    negate: bool = False
    zones: set[ZoneName] = set()


# Zones: Graph resolver
def convert_to_zone_and_deps(zone_yaml: ZoneYAML) -> tuple[Zone, list[ZoneName]]:
    return (Zone(addrs=zone_yaml.addrs, negate=zone_yaml.negate), zone_yaml.zones)


def resolve_zones(zones):
    zones = { name: convert_to_zone_and_deps(ZoneYAML(**zone)) for (name, zone) in zones.items() }
    zone_name = { name: set(zones) for (name, (_, zones)) in zones.items() }

    print(zones)

    for name in TopologicalSorter(zone_name).static_order():
        print(name)

        # TODO: Check negation inclusion


# Blacklist
class BlackList(RestrictiveBaseModel):
    enabled: bool = False
    addr: list[IPvAnyAddress] = []


# Reverse Path Filter
class ReversePathFilter(RestrictiveBaseModel):
    enabled: bool = False


# Filters
class Verdict(str, Enum):
    accept = "accept"
    drop = "drop"
    reject = "reject"


class TcpProtocol(RestrictiveBaseModel):
    dport: list[Port | PortRange] | None
    sport: list[Port | PortRange] | None


class UdpProtocol(RestrictiveBaseModel):
    dport: list[Port | PortRange] | None
    sport: list[Port | PortRange] | None


class Protocols(RestrictiveBaseModel):
    icmp: bool = False
    ospf: bool = False
    tcp: TcpProtocol | None
    udp: UdpProtocol | None
    vrrp: bool = False


class Rule(RestrictiveBaseModel):
    iif: str | None
    oif: str | None
    protocols: Protocols = Protocols()
    src: ZoneName | list[IPvAnyNetwork | ZoneName | FilePath] | None
    verdict: Verdict = Verdict.accept


class ForwardRule(Rule):
    # dest: ZoneEntries | None
    dest: None


class Filter(RestrictiveBaseModel):
    input: list[Rule] = []
    output: list[Rule] = []
    forward: list[ForwardRule] = []


# Nat
class SNat(RestrictiveBaseModel):
    addr: IPvAnyAddress
    persistent: bool = True


class Nat(RestrictiveBaseModel):
    # src: ZoneEntries | None
    src: None
    snat: SNat


# Root model
class Firewall(RestrictiveBaseModel):
    blacklist: BlackList | None
    reverse_path_filter: ReversePathFilter | None
    filter: Filter | None
    nat: list[Nat] = []


def main():
    parser = ArgumentParser()
    parser.add_argument("file", type=FileType("r"), help="YAML rule file")

    args = parser.parse_args()
    contents = safe_load(args.file)

    zones = resolve_zones(contents.pop("zones"))
    print(zones)

    exit(0)

    rules = Firewall(**contents)
    print(rules)

    return 0


if __name__ == "__main__":
    main()