firewall/nftables.py

#!/usr/bin/env python3

from __future__ import annotations
from argparse import ArgumentParser, FileType
from enum import Enum
from pydantic import (
    BaseModel,
    Extra,
    FilePath,
    IPvAnyAddress,
    IPvAnyNetwork,
    validator,
    root_validator,
)
from yaml import safe_load


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


def parse_range_string(s):
    parts = s.split(",")
    values = []

    for part in parts:
        if ".." in part:
            start, end = part.split("..")
            values.append(range(int(start), int(end) + 1))
        else:
            values.append(int(part))

    return values


# Zones
class ZoneName(str):
    pass


class Zone(RestrictiveBaseModel):
    name: ZoneName
    exclude: list[IPvAnyNetwork | ZoneName | FilePath] | None
    include: list[IPvAnyNetwork | ZoneName | FilePath] | None

    @root_validator()
    def validate_mutually_exactly_one(cls, values):
        if values.get("exclude") and values.get("include"):
            raise ValueError("exclude and include are mutually exclusive")

        if values.get("exclude") is None and values.get("include") is None:
            raise ValueError("exactly one of exclude and include must be set")

        return values


# 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: str | None
    sport: str | None

    @validator("dport", "sport")
    def parse_range(cls, v):
        return parse_range_string(v)


class UdpProtocol(RestrictiveBaseModel):
    dport: str | None
    sport: str | None

    @validator("dport", "sport")
    def parse_range(cls, v):
        return parse_range_string(v)


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: ZoneName | list[IPvAnyNetwork | ZoneName | FilePath] | 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: ZoneName | list[IPvAnyNetwork | ZoneName | FilePath] | None
    snat: SNat


# Root model
class Firewall(RestrictiveBaseModel):
    zones: list[Zone] = []
    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()

    rules = Firewall(**safe_load(args.file))

    print(rules)

    return 0


if __name__ == "__main__":
    main()