From f865ec912a21d407c2e713d1ae619dcad1c75995 Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= Date: Fri, 9 Apr 2021 14:33:40 +0200 Subject: [PATCH] fix: Move set/unset running to abstract Avoid issues when a script fails unexpectedly. --- bin/async_capture.py | 5 ++--- bin/background_indexer.py | 4 +--- lookyloo/abstractmanager.py | 5 ++++- 3 files changed, 7 insertions(+), 7 deletions(-) diff --git a/bin/async_capture.py b/bin/async_capture.py index 43e24cdd..5a002a4c 100755 --- a/bin/async_capture.py +++ b/bin/async_capture.py @@ -4,7 +4,7 @@ import logging from lookyloo.abstractmanager import AbstractManager -from lookyloo.helpers import set_running, unset_running, shutdown_requested +from lookyloo.helpers import shutdown_requested from lookyloo.lookyloo import Lookyloo logging.basicConfig(format='%(asctime)s %(name)s %(levelname)s:%(message)s', @@ -16,14 +16,13 @@ class AsyncCapture(AbstractManager): def __init__(self, loglevel: int=logging.INFO): super().__init__(loglevel) self.lookyloo = Lookyloo() + self.script_name = 'async_capture' def _to_run_forever(self): - set_running('async_capture') while True: url = self.lookyloo.process_capture_queue() if url is None or shutdown_requested(): break - unset_running('async_capture') def main(): diff --git a/bin/background_indexer.py b/bin/background_indexer.py index 4f4873b3..92340b2c 100755 --- a/bin/background_indexer.py +++ b/bin/background_indexer.py @@ -4,7 +4,6 @@ import logging from lookyloo.abstractmanager import AbstractManager -from lookyloo.helpers import set_running, unset_running from lookyloo.lookyloo import Lookyloo from lookyloo.exceptions import NoValidHarFile @@ -17,15 +16,14 @@ class BackgroundIndexer(AbstractManager): def __init__(self, loglevel: int=logging.INFO): super().__init__(loglevel) self.lookyloo = Lookyloo() + self.script_name = 'background_indexer' # make sure discarded captures dir exists self.discarded_captures_dir = self.lookyloo.capture_dir.parent / 'discarded_captures' self.discarded_captures_dir.mkdir(parents=True, exist_ok=True) def _to_run_forever(self): - set_running('background_indexer') self._build_missing_pickles() self._check_indexes() - unset_running('background_indexer') def _build_missing_pickles(self): for uuid_path in self.lookyloo.capture_dir.glob('*/uuid'): diff --git a/lookyloo/abstractmanager.py b/lookyloo/abstractmanager.py index 8b8ef243..ca60199a 100644 --- a/lookyloo/abstractmanager.py +++ b/lookyloo/abstractmanager.py @@ -4,7 +4,7 @@ from abc import ABC import logging -from .helpers import long_sleep, shutdown_requested +from .helpers import long_sleep, shutdown_requested, set_running, unset_running class AbstractManager(ABC): @@ -27,9 +27,12 @@ class AbstractManager(ABC): if shutdown_requested(): break try: + set_running(self.script_name) self._to_run_forever() except Exception: self.logger.exception(f'Something went terribly wrong in {self.__class__.__name__}.') + finally: + unset_running(self.script_name) if not long_sleep(sleep_in_sec): break self.logger.info(f'Shutting down {self.__class__.__name__}')