geoffrey
/
eulaurarien
1
0
Fork 0
Code Issues 6 Pull Requests Releases 1 Wiki Activity

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
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

140
database.py
View File

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

2
feed_asn.py
View File

@ -40,7 +40,7 @@ if __name__ == '__main__':
if parsed_prefix.version == 4: if parsed_prefix.version == 4:
DB.set_ip4network( DB.set_ip4network(
prefix, prefix,
# source=path, source=path,
updated=int(time.time()) updated=int(time.time())
) )
log.info('Added %s from %s (%s)', prefix, asn_str, path) log.info('Added %s from %s (%s)', prefix, asn_str, path)