diff --git a/.env.default b/.env.default index 7090df7..1755c0f 100644 --- a/.env.default +++ b/.env.default @@ -1,3 +1,4 @@ RAPID7_API_KEY= CACHE_SIZE=536870912 MASSDNS_HASHMAP_SIZE=1000 +PROFILE=0 diff --git a/database.py b/database.py index 5a8bda2..9cbd330 100644 --- a/database.py +++ b/database.py @@ -11,6 +11,7 @@ import coloredlogs import pickle import numpy import math +import os TLD_LIST: typing.Set[str] = set() @@ -127,13 +128,23 @@ MatchCallable = typing.Callable[[Path, class Profiler(): def __init__(self) -> None: - self.log = logging.getLogger('profiler') - self.time_last = time.perf_counter() - self.time_step = 'init' - self.time_dict: typing.Dict[str, float] = dict() - self.step_dict: typing.Dict[str, int] = dict() + do_profile = int(os.environ.get('PROFILE', '0')) + if do_profile: + self.log = logging.getLogger('profiler') + self.time_last = time.perf_counter() + self.time_step = 'init' + self.time_dict: typing.Dict[str, float] = dict() + self.step_dict: typing.Dict[str, int] = dict() + self.enter_step = self.enter_step_real + self.profile = self.profile_real + else: + self.enter_step = self.enter_step_dummy + self.profile = self.profile_dummy - def enter_step(self, name: str) -> None: + def enter_step_dummy(self, name: str) -> None: + return + + def enter_step_real(self, name: str) -> None: now = time.perf_counter() try: self.time_dict[self.time_step] += now - self.time_last @@ -144,7 +155,10 @@ class Profiler(): self.time_step = name self.time_last = time.perf_counter() - def profile(self) -> None: + def profile_dummy(self) -> None: + return + + def profile_real(self) -> None: self.enter_step('profile') total = sum(self.time_dict.values()) for key, secs in sorted(self.time_dict.items(), key=lambda t: t[1]):