diff --git a/drawbridge/__init__.py b/drawbridge/__init__.py index 44badb1..e05180e 100644 --- a/drawbridge/__init__.py +++ b/drawbridge/__init__.py @@ -2,7 +2,7 @@ import os import iptc -from net_queue import NetQueue +from .net_queue import NetQueue def is_root(): diff --git a/drawbridge/net_queue.py b/drawbridge/net_queue.py index 605d7ae..9952385 100644 --- a/drawbridge/net_queue.py +++ b/drawbridge/net_queue.py @@ -7,8 +7,8 @@ from typing import Union import iptc -from utils.logger import logger -from utils.lookup import Protocols +from .utils.logger import logger +from .utils.lookup import Protocols class NetQueue: @@ -30,7 +30,7 @@ class NetQueue: self.dst_ip = self.validate_ip(dst_ip, "destination") self.protocol = self.validate_protocol(protocol) - self.queue = self.validate_queue(queue, override) + self.queue = self._validate_queue(queue, override) self.rule = self._create_rule() def _create_rule(self) -> iptc.Rule: @@ -72,14 +72,14 @@ class NetQueue: return ip @staticmethod - def validate_port(port: Optional[int], description: str) -> Union(int, None): + def validate_port(port: Optional[int], description: str) -> Union[int, None]: if port: if not 0 <= port <= 65535: raise ValueError(f"Invalid {description} port: {port}") return port @staticmethod - def validate_protocol(protocol: Optional[str]) -> Union(str, None): + def validate_protocol(protocol: Optional[str]) -> Union[str, None]: if protocol: try: Protocols(protocol)