From 7b487db1b2abc99d702c3444bc4d4b5284cd6abd Mon Sep 17 00:00:00 2001
From: Darryl Nixon <git@nixon.mozmail.com>
Date: Sun, 2 Jul 2023 14:00:34 -0700
Subject: [PATCH] Move to pre-routing

---
 drawbridge/drawbridge.py   |  3 +--
 drawbridge/net_queue.py    | 36 ++++++++++++++++++------------------
 drawbridge/utils/lookup.py |  5 ++++-
 3 files changed, 23 insertions(+), 21 deletions(-)

diff --git a/drawbridge/drawbridge.py b/drawbridge/drawbridge.py
index 3e468fb..cedb66c 100644
--- a/drawbridge/drawbridge.py
+++ b/drawbridge/drawbridge.py
@@ -46,7 +46,6 @@ class DrawBridge:
             if packet.payload != original:
                 packet.mangle()
 
-    @staticmethod
     def _delete_rules(self):
         for queue in self.net_queues:
             try:
@@ -60,6 +59,6 @@ class DrawBridge:
             connection = fnfqueue.Connection()
             listener = connection.bind(queue.queue)
             listener.set_mode(65535, fnfqueue.COPY_PACKET)
-            task = asyncio.create_task(self._listen(listener, queue.callback))
+            task = asyncio.create_task(self._listen(connection, queue.callback))
             tasks.append(task)
         await asyncio.gather(*tasks)
diff --git a/drawbridge/net_queue.py b/drawbridge/net_queue.py
index 427c948..91cc90f 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
+from .utils.lookup import PROTOCOLS, TABLES
 
 
 class NetQueue:
     def __init__(
         self,
-        callback: Callable,
         queue: int,
+        callback: Callable,
         src_ip: Optional[str] = None,
         dst_ip: Optional[str] = None,
         src_port: Optional[int] = None,
@@ -44,13 +44,13 @@ class NetQueue:
         return rule
 
     def write_rule(self):
-        table = iptc.Table(iptc.Table.MANGLE)
-        chain = iptc.Chain(table, "INPUT")
+        table = iptc.Table(iptc.Table.NAT)
+        chain = iptc.Chain(table, "PREROUTING")
         chain.insert_rule(self.rule)
 
     def delete_rule(self):
-        table = iptc.Table(iptc.Table.MANGLE)
-        chain = iptc.Chain(table, "INPUT")
+        table = iptc.Table(iptc.Table.NAT)
+        chain = iptc.Chain(table, "PREROUTING")
         try:
             chain.delete_rule(self.rule)
         except iptc.ip4tc.IPTCError:
@@ -83,22 +83,22 @@ class NetQueue:
     def validate_protocol(protocol: Optional[str]) -> Union[str, None]:
         if protocol:
             try:
-                Protocols[protocol]
+                PROTOCOLS[protocol]
             except KeyError:
                 raise KeyError(f"Invalid protocol: {protocol}")
         return protocol
 
     @staticmethod
     def _is_queue_taken(queue: int, override: bool) -> bool:
-        table = iptc.Table(iptc.Table.FILTER)
-        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):
-                    if override:
-                        logger.warning(f"Queue {queue} is already taken, clearing it")
-                        chain.delete_rule(rule)
-                        return False
-                    return True
+        for table in 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):
+                        if override:
+                            logger.warning(f"Queue {queue} is already taken, clearing it")
+                            chain.delete_rule(rule)
+                            return False
+                        return True
         return False
 
     @staticmethod
@@ -116,13 +116,13 @@ class NetQueue:
         return (
             f"NetQueueFilter("
             f"queue={self.queue}, "
+            f"callback={self.callback}, "
             f"src_ip={self.src_ip}, "
             f"dst_ip={self.dst_ip}, "
             f"src_port={self.src_port}, "
             f"dst_port={self.dst_port}, "
             f"protocol={self.protocol}, "
-            f"callback={self.callback}, "
-            f"async_callback={self.async_callback}"
+            f"rule={self.rule}, "
             f")"
         )
 
diff --git a/drawbridge/utils/lookup.py b/drawbridge/utils/lookup.py
index 77c4765..b84bc36 100644
--- a/drawbridge/utils/lookup.py
+++ b/drawbridge/utils/lookup.py
@@ -1,6 +1,7 @@
 import socket
+import iptc
 
-Protocols = {
+PROTOCOLS = {
     "ah": socket.IPPROTO_AH,
     "dstopts": socket.IPPROTO_DSTOPTS,
     "egp": socket.IPPROTO_EGP,
@@ -26,3 +27,5 @@ Protocols = {
     "tp": socket.IPPROTO_TP,
     "udp": socket.IPPROTO_UDP,
 }
+
+TABLES = [iptc.Table(t) for t in iptc.Table.ALL]
\ No newline at end of file