diff --git a/database.py b/database.py index fc2855e..5522013 100644 --- a/database.py +++ b/database.py @@ -9,7 +9,6 @@ import time import logging import coloredlogs import pickle -import enum coloredlogs.install( level='DEBUG', @@ -22,6 +21,7 @@ Level = int class Path(): + # FP add boolean here pass @@ -53,27 +53,46 @@ class Ip4Path(Path): self.prefixlen = prefixlen -Match = typing.Tuple[Timestamp, Path, Level] +class Match(): + def __init__(self) -> None: + self.updated: int = 0 + self.level: int = 0 + self.source: Path = RulePath() + # FP dupplicate args -# class AsnNode(): -# def __init__(self, asn: int) -> None: -# self.asn = asn + def set(self, + updated: int, + level: int, + source: Path, + ) -> None: + if updated > self.updated or level > self.level: + self.updated = updated + self.level = level + self.source = source + # FP dupplicate function + + def active(self) -> bool: + return self.updated > 0 + + +class AsnNode(Match): + pass class DomainTreeNode(): def __init__(self) -> None: self.children: typing.Dict[str, DomainTreeNode] = dict() - self.match_zone: typing.Optional[Match] = None - self.match_hostname: typing.Optional[Match] = None + self.match_zone = Match() + self.match_hostname = Match() class IpTreeNode(): def __init__(self) -> None: self.children: typing.List[typing.Optional[IpTreeNode]] = [None, None] - self.match: typing.Optional[Match] = None + self.match = Match() -Node = typing.Union[DomainTreeNode, IpTreeNode, Asn] +Node = typing.Union[DomainTreeNode, IpTreeNode, AsnNode] NodeCallable = typing.Callable[[Path, Node, typing.Optional[typing.Any]], @@ -112,7 +131,7 @@ class Profiler(): class Database(Profiler): - VERSION = 9 + VERSION = 10 PATH = "blocking.p" def initialize(self) -> None: @@ -120,7 +139,7 @@ class Database(Profiler): "Creating database version: %d ", Database.VERSION) self.domtree = DomainTreeNode() - self.asns: typing.Set[Asn] = set() + self.asns: typing.Dict[Asn, AsnNode] = dict() self.ip4tree = IpTreeNode() def load(self) -> None: @@ -133,12 +152,12 @@ class Database(Profiler): return self.log.warning( "Outdated database version found: %d, " - "will be rebuilt.", + "it will be rebuilt.", version) except (TypeError, AttributeError, EOFError): self.log.error( - "Corrupt database found, " - "will be rebuilt.") + "Corrupt (or heavily outdated) database found, " + "it will be rebuilt.") except FileNotFoundError: pass self.initialize() @@ -306,7 +325,7 @@ class Database(Profiler): dic = self.domtree depth = 0 for part in domain.path: - if dic.match_zone: + if dic.match_zone.active(): self.enter_step('get_domain_yield') yield ZonePath(domain.path[:depth]) self.enter_step('get_domain_brws') @@ -314,10 +333,10 @@ class Database(Profiler): return dic = dic.children[part] depth += 1 - if dic.match_zone: + if dic.match_zone.active(): self.enter_step('get_domain_yield') yield ZonePath(domain.path) - if dic.match_hostname: + if dic.match_hostname.active(): self.enter_step('get_domain_yield') yield HostnamePath(domain.path) @@ -328,7 +347,7 @@ class Database(Profiler): dic = self.ip4tree for i in reversed(range(ip4.prefixlen)): part = (ip4.value >> i) & 0b1 - if dic.match: + if dic.match.active(): self.enter_step('get_ip4_yield') yield Ip4Path(ip4.value, 32-i) self.enter_step('get_ip4_brws') @@ -336,7 +355,7 @@ class Database(Profiler): if next_dic is None: return dic = next_dic - if dic.match: + if dic.match.active(): self.enter_step('get_ip4_yield') yield ip4 @@ -344,58 +363,61 @@ class Database(Profiler): for asn in self.asns: yield AsnPath(asn) - def set_hostname(self, - hostname_str: str, - updated: int, - is_first_party: bool = None, - source: Path = None) -> None: - self.enter_step('set_hostname_pack') + def _set_domain(self, + hostname: bool, + domain_str: str, + updated: int, + is_first_party: bool = None, + source: Path = None) -> None: + self.enter_step('set_domain_pack') if is_first_party: raise NotImplementedError - self.enter_step('set_hostname_brws') - hostname = self.pack_domain(hostname_str) + domain = self.pack_domain(domain_str) + self.enter_step('set_domain_brws') dic = self.domtree - for part in hostname.path: - if dic.match_zone: - # Refuse to add hostname whose zone is already matching + for part in domain.path: + if dic.match_zone.active(): + # Refuse to add domain whose zone is already matching return if part not in dic.children: dic.children[part] = DomainTreeNode() dic = dic.children[part] - dic.match_hostname = (updated, source or RulePath(), 0) + if hostname: + match = dic.match_hostname + else: + match = dic.match_zone + match.set( + updated, + 0, # TODO Level + source or RulePath(), + ) + + def set_hostname(self, + *args: typing.Any, **kwargs: typing.Any + ) -> None: + self._set_domain(True, *args, **kwargs) def set_zone(self, - zone_str: str, - updated: int, - is_first_party: bool = None, - source: Path = None) -> None: - self.enter_step('set_zone_pack') - if is_first_party: - raise NotImplementedError - zone = self.pack_domain(zone_str) - self.enter_step('set_zone_brws') - dic = self.domtree - for part in zone.path: - if dic.match_zone: - # Refuse to add zone whose parent zone is already matching - return - if part not in dic.children: - dic.children[part] = DomainTreeNode() - dic = dic.children[part] - dic.match_zone = (updated, source or RulePath(), 0) + *args: typing.Any, **kwargs: typing.Any + ) -> None: + self._set_domain(False, *args, **kwargs) def set_asn(self, asn_str: str, updated: int, is_first_party: bool = None, source: Path = None) -> None: - self.enter_step('set_asn_pack') - if is_first_party or source: - # TODO updated + self.enter_step('set_asn') + if is_first_party: raise NotImplementedError - asn = self.pack_asn(asn_str) - self.enter_step('set_asn_brws') - self.asns.add(asn.asn) + path = self.pack_asn(asn_str) + match = AsnNode() + match.set( + updated, + 0, + source or RulePath() + ) + self.asns[path.asn] = match def _set_ip4(self, ip4: Ip4Path, @@ -407,7 +429,7 @@ class Database(Profiler): dic = self.ip4tree for i in reversed(range(ip4.prefixlen)): part = (ip4.value >> i) & 0b1 - if dic.match: + if dic.match.active(): # Refuse to add ip4* whose network is already matching return next_dic = dic.children[part] @@ -415,7 +437,11 @@ class Database(Profiler): next_dic = IpTreeNode() dic.children[part] = next_dic dic = next_dic - dic.match = (updated, source or RulePath(), 0) + dic.match.set( + updated, + 0, # TODO Level + source or RulePath(), + ) def set_ip4address(self, ip4address_str: str, diff --git a/feed_asn.py b/feed_asn.py index aa311f8..f34773f 100755 --- a/feed_asn.py +++ b/feed_asn.py @@ -40,7 +40,7 @@ if __name__ == '__main__': if parsed_prefix.version == 4: DB.set_ip4network( prefix, - # source=path, + source=path, updated=int(time.time()) ) log.info('Added %s from %s (%s)', prefix, asn_str, path)