Default queue #0 and remove some redundancy

drawbridge/drawbridge.py

@@ -16,8 +16,8 @@ class DrawBridge:
 
     def add_queue(
         self,
-        queue: int,
         callback: Callable,
+        queue: int = 0,
         src_ip: Optional[str] = None,
         dst_ip: Optional[str] = None,
         src_port: Optional[int] = None,
@@ -26,7 +26,7 @@ class DrawBridge:
         override: bool = False,
     ):
         try:
-            new_queue = NetQueue(queue, callback, src_ip, dst_ip, src_port, dst_port, protocol, override)
+            new_queue = NetQueue(callback, queue, src_ip, dst_ip, src_port, dst_port, protocol, override)
             new_queue.write_rule()
         except Exception as e:
             logger.error(f"Failed to initialize NetQueue: {e}")

drawbridge/net_queue.py

@@ -8,14 +8,14 @@ from typing import Union
 import iptc
 
 from .utils.logger import logger
-from .utils.lookup import PROTOCOLS, TABLES
+from .utils.lookup import PROTOCOLS, ALL_TABLES, PREROUTING_MANGLE
 
 
 class NetQueue:
     def __init__(
         self,
-        queue: int,
         callback: Callable,
+        queue: int,
         src_ip: Optional[str] = None,
         dst_ip: Optional[str] = None,
         src_port: Optional[int] = None,
@@ -44,15 +44,11 @@ class NetQueue:
         return rule
 
     def write_rule(self):
-        table = iptc.Table(iptc.Table.MANGLE)
-        chain = iptc.Chain(table, "PREROUTING")
-        chain.insert_rule(self.rule)
+        PREROUTING_MANGLE.insert_rule(self.rule)
 
     def delete_rule(self):
-        table = iptc.Table(iptc.Table.MANGLE)
-        chain = iptc.Chain(table, "PREROUTING")
         try:
-            chain.delete_rule(self.rule)
+            PREROUTING_MANGLE.delete_rule(self.rule)
         except iptc.ip4tc.IPTCError:
             logger.warning("Failed to delete rule, it may have already been deleted")
 
@@ -90,7 +86,7 @@ class NetQueue:
 
     @staticmethod
     def _is_queue_taken(queue: int, override: bool) -> bool:
-        for table in TABLES:
+        for table in ALL_TABLES:
             for chain in table.chains:
                 for rule in chain.rules:
                     if rule.target.name == "NFQUEUE" and rule.target.get_all_parameters()["queue-num"] == str(queue):

drawbridge/utils/lookup.py

@@ -28,4 +28,5 @@ PROTOCOLS = {
     "udp": socket.IPPROTO_UDP,
 }
 
-TABLES = [iptc.Table(t) for t in iptc.Table.ALL]
+ALL_TABLES = [iptc.Table(t) for t in iptc.Table.ALL]
+PREROUTING_MANGLE = iptc.Chain(iptc.Table(iptc.Table.MANGLE), "PREROUTING")