adding firewall function to firegex!

backend/modules/porthijack/firewall.py

@@ -1,47 +1,10 @@
 import asyncio
-from typing import Dict
 from modules.porthijack.nftables import FiregexTables
 from modules.porthijack.models import Service
 from utils.sqlite import SQLite
 
 nft = FiregexTables()
 
-class FirewallManager:
-    def __init__(self, db:SQLite):
-        self.db = db
-        self.service_table: Dict[str, ServiceManager] = {}
-        self.lock = asyncio.Lock()
-
-    async def close(self):
-        for key in list(self.service_table.keys()):
-            await self.remove(key)
-
-    async def remove(self,srv_id):
-        async with self.lock: 
-            if srv_id in self.service_table:
-                await self.service_table[srv_id].disable()
-                del self.service_table[srv_id]
-    
-    async def init(self):
-        FiregexTables().init()
-        await self.reload()
-
-    async def reload(self):
-        async with self.lock: 
-            for srv in self.db.query('SELECT * FROM services;'):
-                srv = Service.from_dict(srv)
-                if srv.service_id in self.service_table:
-                    continue
-                self.service_table[srv.service_id] = ServiceManager(srv, self.db)
-                if srv.active:
-                    await self.service_table[srv.service_id].enable()
-
-    def get(self,srv_id):
-        if srv_id in self.service_table:
-            return self.service_table[srv_id]
-        else:
-            raise ServiceNotFoundException()
-        
 class ServiceNotFoundException(Exception): pass
 
 class ServiceManager:
@@ -74,4 +37,41 @@ class ServiceManager:
     
     async def restart(self):
         await self.disable()
-        await self.enable()
+        await self.enable()
+
+class FirewallManager:
+    def __init__(self, db:SQLite):
+        self.db = db
+        self.service_table: dict[str, ServiceManager] = {}
+        self.lock = asyncio.Lock()
+
+    async def close(self):
+        for key in list(self.service_table.keys()):
+            await self.remove(key)
+
+    async def remove(self,srv_id):
+        async with self.lock: 
+            if srv_id in self.service_table:
+                await self.service_table[srv_id].disable()
+                del self.service_table[srv_id]
+    
+    async def init(self):
+        FiregexTables().init()
+        await self.reload()
+
+    async def reload(self):
+        async with self.lock: 
+            for srv in self.db.query('SELECT * FROM services;'):
+                srv = Service.from_dict(srv)
+                if srv.service_id in self.service_table:
+                    continue
+                self.service_table[srv.service_id] = ServiceManager(srv, self.db)
+                if srv.active:
+                    await self.service_table[srv.service_id].enable()
+
+    def get(self,srv_id) -> ServiceManager:
+        if srv_id in self.service_table:
+            return self.service_table[srv_id]
+        else:
+            raise ServiceNotFoundException()
+        

backend/modules/porthijack/nftables.py

@@ -1,4 +1,3 @@
-from typing import List
 from modules.porthijack.models import Service
 from utils import addr_parse, ip_parse, ip_family, NFTableManager, nftables_json_to_int
 
@@ -13,9 +12,7 @@ class FiregexHijackRule():
         self.ip_dst = str(ip_dst)
 
     def __eq__(self, o: object) -> bool:
-        if isinstance(o, FiregexHijackRule):
-            return self.public_port == o.public_port and self.proto == o.proto and ip_parse(self.ip_src) == ip_parse(o.ip_src)
-        elif isinstance(o, Service):
+        if isinstance(o, FiregexHijackRule) or isinstance(o, Service):
             return self.public_port == o.public_port and self.proto == o.proto and ip_parse(self.ip_src) == ip_parse(o.ip_src)
         return False
 
@@ -79,10 +76,9 @@ class FiregexTables(NFTableManager):
         }}})
 
 
-    def get(self) -> List[FiregexHijackRule]:
+    def get(self) -> list[FiregexHijackRule]:
         res = []
         for filter in self.list_rules(tables=[self.table_name], chains=[self.prerouting_porthijack,self.postrouting_porthijack]):
-            filter["expr"][0]["match"]["right"]
             res.append(FiregexHijackRule(
                 target=filter["chain"],
                 id=int(filter["handle"]),