pdf/drawbridge
1
0
Fork 0
mirror of https://github.com/DarrylNixon/drawbridge synced 2024-04-22 12:17:07 -07:00
Code Issues Projects Releases Packages Wiki Activity

Default queue #0 and remove some redundancy

Browse source
This commit is contained in:
Darryl Nixon 2023-07-02 14:39:10 -07:00
parent 07decd8707
commit 1e9e4fdc4c
3 changed files with 9 additions and 12 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

4
drawbridge/drawbridge.py
View file

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

14
drawbridge/net_queue.py
View file

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

3
drawbridge/utils/lookup.py
View file

@ -28,4 +28,5 @@ PROTOCOLS = {
"udp": socket.IPPROTO_UDP, "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")