From 1ce8b76555db5076398e30bd573dcb0ca53f8cee Mon Sep 17 00:00:00 2001
From: Vincent Lafeychine <vincent.lafeychine@proton.me>
Date: Sun, 13 Aug 2023 18:40:29 +0200
Subject: [PATCH] feat(zones): Early zone management

---
 example.yaml | 23 ++++++++++-----------
 nftables.py  | 58 +++++++++++++++++++++++++++++++++++-----------------
 2 files changed, 50 insertions(+), 31 deletions(-)

diff --git a/example.yaml b/example.yaml
index 20b074f..cef6e9b 100644
--- a/example.yaml
+++ b/example.yaml
@@ -1,20 +1,19 @@
 ---
 zones:
-  - name: users-internet-allowed
-    include:
-      files: [example.yaml]
+  users-internet-allowed:
+    files: [example.yaml]
 
-  - name: mgmt
-    include:
-      addrs: [10.203.0.0/16]
+  mgmt:
+    addrs: [10.203.0.0/16]
 
-  - name: adm
-    include:
-      addrs: [2a09:6840::/29, 10.128.0.0/16]
+  adm:
+    addrs: [2a09:6840::/29, 10.128.0.0/16]
 
-  - name: internet
-    exclude:
-      zones: [adm, mgmt]
+  internet:
+    negate: true
+    zones: [adm, mgmt]
+
+# interne: negate KO
 
 blacklist:
   enabled: true
diff --git a/nftables.py b/nftables.py
index c7f993d..b9d2daf 100755
--- a/nftables.py
+++ b/nftables.py
@@ -1,7 +1,9 @@
 #!/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,
@@ -43,31 +45,43 @@ class PortRange(str):
         return range(start, end)
 
 
+# ===== First pass: Zones =====
+
+
 # Zones
 class ZoneName(str):
     pass
 
 
-class ZoneEntries(RestrictiveBaseModel):
-    addrs: list[IPvAnyNetwork] | None
-    files: list[FilePath] | None
-    zones: list[ZoneName] | None
+@dataclass
+class Zone:
+    addrs: set[IPvAnyNetwork]
+    negate: bool
 
 
-class Zone(RestrictiveBaseModel):
-    name: ZoneName
-    exclude: ZoneEntries | None
-    include: ZoneEntries | None
+# Zones: Parsing YAML
+class ZoneYAML(RestrictiveBaseModel):
+    addrs: set[IPvAnyNetwork] = set()
+    files: set[FilePath] = set()
+    negate: bool = False
+    zones: set[ZoneName] = set()
 
-    @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")
+# 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)
 
-        return values
+
+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
@@ -115,7 +129,8 @@ class Rule(RestrictiveBaseModel):
 
 
 class ForwardRule(Rule):
-    dest: ZoneEntries | None
+    # dest: ZoneEntries | None
+    dest: None
 
 
 class Filter(RestrictiveBaseModel):
@@ -131,13 +146,13 @@ class SNat(RestrictiveBaseModel):
 
 
 class Nat(RestrictiveBaseModel):
-    src: ZoneEntries | None
+    # src: ZoneEntries | None
+    src: None
     snat: SNat
 
 
 # Root model
 class Firewall(RestrictiveBaseModel):
-    zones: list[Zone] = []
     blacklist: BlackList | None
     reverse_path_filter: ReversePathFilter | None
     filter: Filter | None
@@ -149,9 +164,14 @@ def main():
     parser.add_argument("file", type=FileType("r"), help="YAML rule file")
 
     args = parser.parse_args()
+    contents = safe_load(args.file)
 
-    rules = Firewall(**safe_load(args.file))
+    zones = resolve_zones(contents.pop("zones"))
+    print(zones)
 
+    exit(0)
+
+    rules = Firewall(**contents)
     print(rules)
 
     return 0