refactored nftable managment, and fixed stop of the container

Dockerfile

@@ -38,6 +38,7 @@ RUN g++ binsrc/proxy.cpp -o modules/proxy -O3 -pthread -lboost_system -lboost_th
 
 COPY ./backend/ /execute/
 COPY --from=frontend /app/build/ ./frontend/
-ENTRYPOINT ["/bin/sh", "/execute/docker-entrypoint.sh"]
+
+CMD ["/bin/sh", "/execute/docker-entrypoint.sh"]
 
 

backend/docker-entrypoint.sh

@@ -2,8 +2,7 @@
 
 chown nobody:nobody -R /execute/
 
-capsh --caps="cap_net_admin+eip cap_setpcap,cap_setuid,cap_setgid+ep" \
-    --keep=1 --user=nobody --addamb=cap_net_admin -- \
-    -c "python3 /execute/app.py DOCKER"
+exec capsh --caps="cap_net_admin+eip cap_setpcap,cap_setuid,cap_setgid+ep" \
+    --keep=1 --user=nobody --addamb=cap_net_admin -- -c "python3 /execute/app.py DOCKER"
 
 

backend/modules/nfregex/firegex.py

@@ -1,10 +1,12 @@
 from typing import Dict, List, Set
-from utils.firegextables import FiregexFilter, FiregexTables
-from utils import ip_parse, ip_family, run_func
+from modules.nfregex.nftables import FiregexFilter, FiregexTables
+from utils import ip_parse, run_func
 from modules.nfregex.models import Service, Regex
 import re, os, asyncio
 import traceback
 
+nft = FiregexTables()
+
 class RegexFilter:
     def __init__(
         self, regex,
@@ -68,9 +70,9 @@ class FiregexInterceptor:
         self.update_config_lock = asyncio.Lock()
         input_range, output_range = await self._start_binary()
         self.update_task = asyncio.create_task(self.update_blocked())
-        if not filter in FiregexTables().get():
-            FiregexTables().add_input(queue_range=input_range, proto=self.filter.proto, port=self.filter.port, ip_int=self.filter.ip_int)
-            FiregexTables().add_output(queue_range=output_range, proto=self.filter.proto, port=self.filter.port, ip_int=self.filter.ip_int)
+        if not filter in nft.get():
+            nft.add_input(queue_range=input_range, proto=self.filter.proto, port=self.filter.port, ip_int=self.filter.ip_int)
+            nft.add_output(queue_range=output_range, proto=self.filter.proto, port=self.filter.port, ip_int=self.filter.ip_int)
         return self
     
     async def _start_binary(self):
@@ -140,7 +142,6 @@ class FiregexInterceptor:
         return res
 
 def delete_by_srv(srv:Service):
-    nft = FiregexTables()
     for filter in nft.get():
         if filter.port == srv.port and filter.proto == srv.proto and ip_parse(filter.ip_int) == ip_parse(srv.ip_int):
             nft.cmd({"delete":{"rule": {"handle": filter.id, "table": nft.table_name, "chain": filter.target, "family": "inet"}}})

backend/modules/nfregex/firewall.py

@@ -1,6 +1,7 @@
 import asyncio
 from typing import Dict
-from modules.nfregex.firegex import FiregexFilter, FiregexInterceptor, FiregexTables, RegexFilter, delete_by_srv
+from modules.nfregex.firegex import FiregexInterceptor, RegexFilter, delete_by_srv
+from modules.nfregex.nftables import FiregexTables, FiregexFilter
 from modules.nfregex.models import Regex, Service
 from utils.sqlite import SQLite
 

backend/modules/nfregex/nftables.py

@@ -1,10 +1,8 @@
 from typing import List
-import nftables
-from utils import ip_parse, ip_family
+from utils import ip_parse, ip_family, NFTableManager
 
 class FiregexFilter():
     def __init__(self, proto:str, port:int, ip_int:str, queue=None, target:str=None, id=None):
-        self.nftables = nftables.Nftables()
         self.id = int(id) if id else None
         self.queue = queue
         self.target = target
@@ -17,55 +15,36 @@ class FiregexFilter():
             return self.port == o.port and self.proto == o.proto and ip_parse(self.ip_int) == ip_parse(o.ip_int)
         return False
 
-class FiregexTables:
-
+class FiregexTables(NFTableManager):
+    input_chain = "nfregex_input"
+    output_chain = "nfregex_output"
+    
     def __init__(self):
-        self.table_name = "firegex"
-        self.nft = nftables.Nftables()
-    
-    def raw_cmd(self, *cmds):
-        return self.nft.json_cmd({"nftables": list(cmds)})
-
-    def cmd(self, *cmds):
-        code, out, err = self.raw_cmd(*cmds)
-
-        if code == 0: return out
-        else: raise Exception(err)
-    
-    def init(self):
-        self.reset()
-        code, out, err = self.raw_cmd({"create":{"table":{"name":self.table_name,"family":"inet"}}})
-        if code == 0:
-            self.cmd(
-                {"create":{"chain":{
-                    "family":"inet",
-                    "table":self.table_name,
-                    "name":"input",
-                    "type":"filter",
-                    "hook":"prerouting",
-                    "prio":-150,
-                    "policy":"accept"
-                }}},
-                {"create":{"chain":{
-                    "family":"inet",
-                    "table":self.table_name,
-                    "name":"output",
-                    "type":"filter",
-                    "hook":"postrouting",
-                    "prio":-150,
-                    "policy":"accept"
-                }}}
-            )
-        
-            
-    def reset(self):
-        self.raw_cmd(
-            {"flush":{"table":{"name":"firegex","family":"inet"}}},
-            {"delete":{"table":{"name":"firegex","family":"inet"}}},
-        )
-
-    def list(self):
-        return self.cmd({"list": {"ruleset": None}})["nftables"]
+        super().__init__([
+            {"create":{"chain":{
+                "family":"inet",
+                "table":self.table_name,
+                "name":self.input_chain,
+                "type":"filter",
+                "hook":"prerouting",
+                "prio":-150,
+                "policy":"accept"
+            }}},
+            {"create":{"chain":{
+                "family":"inet",
+                "table":self.table_name,
+                "name":self.output_chain,
+                "type":"filter",
+                "hook":"postrouting",
+                "prio":-150,
+                "policy":"accept"
+            }}}
+        ],[
+            {"flush":{"chain":{"table":self.table_name,"family":"inet", "name":self.input_chain}}},
+            {"delete":{"chain":{"table":self.table_name,"family":"inet", "name":self.input_chain}}},
+            {"flush":{"chain":{"table":self.table_name,"family":"inet", "name":self.output_chain}}},
+            {"delete":{"chain":{"table":self.table_name,"family":"inet", "name":self.output_chain}}},
+        ])
 
     def add_output(self, queue_range, proto, port, ip_int):
         init, end = queue_range
@@ -76,7 +55,7 @@ class FiregexTables:
         self.cmd({ "insert":{ "rule": {
             "family": "inet",
             "table": self.table_name,
-            "chain": "output",
+            "chain": self.output_chain,
             "expr": [
                     {'match': {'left': {'payload': {'protocol': ip_family(ip_int), 'field': 'saddr'}}, 'op': '==', 'right': {"prefix": {"addr": ip_addr, "len": ip_addr_cidr}}}},
                     {'match': {"left": { "payload": {"protocol": str(proto), "field": "sport"}}, "op": "==", "right": int(port)}},
@@ -93,7 +72,7 @@ class FiregexTables:
         self.cmd({"insert":{"rule":{
             "family": "inet",
             "table": self.table_name,
-            "chain": "input",
+            "chain": self.input_chain,
             "expr": [
                     {'match': {'left': {'payload': {'protocol': ip_family(ip_int), 'field': 'daddr'}}, 'op': '==', 'right': {"prefix": {"addr": ip_addr, "len": ip_addr_cidr}}}},
                     {'match': {"left": { "payload": {"protocol": str(proto), "field": "dport"}}, "op": "==", "right": int(port)}},

backend/routers/nfregex.py

@@ -5,7 +5,7 @@ import sqlite3
 from typing import List, Union
 from fastapi import APIRouter, HTTPException
 from pydantic import BaseModel
-from modules.nfregex.firegex import FiregexTables
+from modules.nfregex.nftables import FiregexTables
 from modules.nfregex.firewall import STATUS, FirewallManager
 from utils.sqlite import SQLite
 from utils import ip_parse, refactor_name, refresh_frontend

backend/utils/__init__.py

@@ -1,7 +1,6 @@
 import asyncio
 from ipaddress import ip_interface
-import os, socket, psutil
-import sys
+import os, socket, psutil, sys, nftables
 from fastapi_socketio import SocketManager
 
 LOCALHOST_IP = socket.gethostbyname(os.getenv("LOCALHOST_IP","127.0.0.1"))
@@ -47,4 +46,41 @@ def get_interfaces():
             for interf in interfs:
                 if interf.family in [socket.AF_INET, socket.AF_INET6]:
                     yield {"name": int_name, "addr":interf.address}
-    return list(_get_interfaces())
+    return list(_get_interfaces())
+
+class Singleton(object):
+    __instance = None
+    def __new__(class_, *args, **kwargs):
+        if not isinstance(class_.__instance, class_):
+            class_.__instance = object.__new__(class_, *args, **kwargs)
+        return class_.__instance
+
+class NFTableManager(Singleton):
+    
+    table_name = "firegex"
+    
+    def __init__(self, init_cmd, reset_cmd):
+        self.__init_cmds = init_cmd
+        self.__reset_cmds = reset_cmd
+        self.nft = nftables.Nftables()
+    
+    def raw_cmd(self, *cmds):
+        return self.nft.json_cmd({"nftables": list(cmds)})
+
+    def cmd(self, *cmds):
+        code, out, err = self.raw_cmd(*cmds)
+
+        if code == 0: return out
+        else: raise Exception(err)
+    
+    def init(self):
+        self.reset()
+        self.raw_cmd({"create":{"table":{"name":self.table_name,"family":"inet"}}})
+        self.cmd(*self.__init_cmds)
+            
+    def reset(self):
+        print(self.raw_cmd(*self.__reset_cmds))
+
+    def list(self):
+        return self.cmd({"list": {"ruleset": None}})["nftables"]
+            

start.py

@@ -27,12 +27,16 @@ parser.add_argument('--no-autostart', "-n", required=False, action="store_true",
 parser.add_argument('--keep','-k', required=False, action="store_true", help='Keep the docker-compose file generated', default=False)
 parser.add_argument('--build', "-b", required=False, action="store_true", help='Build the container locally', default=False)
 parser.add_argument('--stop', '-s', required=False, action="store_true", help='Stop firegex execution', default=False)
+parser.add_argument('--restart', '-r', required=False, action="store_true", help='Restart firegex', default=False)
 parser.add_argument('--psw-no-interactive',type=str, required=False, help='Password for no-interactive mode', default=None)
-parser.add_argument('--startup-psw', required=False, action="store_true", help='Insert password in the startup screen of firegex', default=False)
+parser.add_argument('--startup-psw','-P', required=False, action="store_true", help='Insert password in the startup screen of firegex', default=False)
+
 
 args = parser.parse_args()
 os.chdir(os.path.dirname(os.path.realpath(__file__)))
 
+start_operation = not (args.stop or args.restart)
+
 if args.build and not os.path.isfile("./Dockerfile"):
     puts("This is not a clone of firegex, to build firegex the clone of the repository is needed!", color=colors.red)
     exit()
@@ -40,14 +44,14 @@ if args.build and not os.path.isfile("./Dockerfile"):
 if args.threads < 1: 
     args.threads = multiprocessing.cpu_count()
 
-if not args.stop:
+if start_operation:
     sep()
     puts(f"Firegex", color=colors.yellow, end="")
     puts(" will start on port ", end="")
     puts(f"{args.port}", color=colors.cyan)
 
 psw_set = None
-if not args.stop:
+if start_operation:
     if args.psw_no_interactive:
         psw_set = args.psw_no_interactive
     elif not args.startup_psw:
@@ -100,7 +104,10 @@ services:
 sep()
 if not args.no_autostart:
     try:
-        if args.stop:
+        if args.restart:
+            puts("Running 'docker-compose restart'\n", color=colors.green)
+            os.system("docker-compose -p firegex restart")
+        elif args.stop:
             puts("Running 'docker-compose down'\n", color=colors.green)
             os.system("docker-compose -p firegex down")
         else: