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:
Geoffrey Frogeye 2019-12-15 23:13:25 +01:00
parent aec8d3f8de
commit a0e68f0848
Signed by: geoffrey
GPG key ID: D8A7ECA00A8CD3DD
2 changed files with 84 additions and 58 deletions

View file

@ -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,

View file

@ -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)