Reworked match and node system
For level, and first_party later Next: add get_match to retrieve level of source and have correct levels ... am I going somewhere with all this?
This commit is contained in:
parent
aec8d3f8de
commit
a0e68f0848
134
database.py
134
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,
|
||||
def _set_domain(self,
|
||||
hostname: bool,
|
||||
domain_str: str,
|
||||
updated: int,
|
||||
is_first_party: bool = None,
|
||||
source: Path = None) -> None:
|
||||
self.enter_step('set_hostname_pack')
|
||||
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,
|
||||
|
|
|
@ -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)
|
||||
|
|
Loading…
Reference in a new issue