From 1e9e4fdc4c971a090cbfe1fa9e88d8ccf1c6a9e8 Mon Sep 17 00:00:00 2001
From: Darryl Nixon <git@nixon.mozmail.com>
Date: Sun, 2 Jul 2023 14:39:10 -0700
Subject: [PATCH] Default queue #0 and remove some redundancy

---
 drawbridge/drawbridge.py   |  4 ++--
 drawbridge/net_queue.py    | 14 +++++---------
 drawbridge/utils/lookup.py |  3 ++-
 3 files changed, 9 insertions(+), 12 deletions(-)

diff --git a/drawbridge/drawbridge.py b/drawbridge/drawbridge.py
index c75d8a1..9aa8467 100644
--- a/drawbridge/drawbridge.py
+++ b/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}")
diff --git a/drawbridge/net_queue.py b/drawbridge/net_queue.py
index 4fa85fa..19bcc18 100644
--- a/drawbridge/net_queue.py
+++ b/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):
diff --git a/drawbridge/utils/lookup.py b/drawbridge/utils/lookup.py
index b84bc36..9c13582 100644
--- a/drawbridge/utils/lookup.py
+++ b/drawbridge/utils/lookup.py
@@ -28,4 +28,5 @@ PROTOCOLS = {
     "udp": socket.IPPROTO_UDP,
 }
 
-TABLES = [iptc.Table(t) for t in iptc.Table.ALL]
\ No newline at end of file
+ALL_TABLES = [iptc.Table(t) for t in iptc.Table.ALL]
+PREROUTING_MANGLE = iptc.Chain(iptc.Table(iptc.Table.MANGLE), "PREROUTING")
\ No newline at end of file