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 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,
|
||||||
|
domain_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_hostname_pack')
|
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,
|
||||||
|
|
|
@ -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)
|
||||||
|
|
Loading…
Reference in a new issue