diff --git a/Makefile b/Makefile index 9278ad9..d435994 100644 --- a/Makefile +++ b/Makefile @@ -4,6 +4,10 @@ lint: python3 -m flake8 --select F --per-file-ignores="__init__.py:F401" src/ tests/ example/ python3 -m mypy src/ tests/ example/ +diff: + python3 -m isort --diff -rc src/ tests/ example/ + python3 -m black --diff src/ tests/ example/ + format: python3 -m isort -rc src/ tests/ example/ python3 -m black src/ tests/ example/ diff --git a/setup.py b/setup.py index 609ac61..47c8eb3 100755 --- a/setup.py +++ b/setup.py @@ -9,7 +9,7 @@ setuptools.setup( name="enochecker", - version="0.2.1", + version="0.2.2", author="domenukk", author_email="dmaier@sect.tu-berlin.de", description="Library to build checker scripts for EnoEngine A/D CTF Framework in Python", diff --git a/src/enochecker/enochecker.py b/src/enochecker/enochecker.py index db462b1..14479ab 100644 --- a/src/enochecker/enochecker.py +++ b/src/enochecker/enochecker.py @@ -30,7 +30,7 @@ from .checkerservice import CHECKER_METHODS, init_service from .logging import ELKFormatter from .nosqldict import NoSqlDict -from .results import CheckerResult, EnoException, Result +from .results import CheckerResult, EnoException, OfflineException, Result from .storeddict import DB_DEFAULT_DIR, DB_GLOBAL_CACHE_SETTING, StoredDict from .useragents import random_useragent from .utils import SimpleSocket, snake_caseify @@ -739,6 +739,7 @@ def connect( host: Optional[str] = None, port: Optional[int] = None, timeout: Optional[int] = None, + retries: int = 3, ) -> SimpleSocket: """ Open a socket/telnet connection to the remote host. @@ -750,22 +751,51 @@ def connect( :param timeout: timeout on connection (defaults to self.timeout) :return: A connected Telnet instance """ + if timeout: - timeout_fun: Callable[[], int] = lambda: cast(int, timeout) + + def timeout_fun() -> int: + return cast(int, timeout) + else: - timeout = self.time_remaining // 2 - timeout_fun = lambda: self.time_remaining // 2 + + def timeout_fun() -> int: + return self.time_remaining // 2 if port is None: port = self.port if host is None: host = self.address - self.debug( - "Opening socket to {}:{} (timeout {} secs).".format(host, port, timeout) - ) - return SimpleSocket( - host, port, timeout=timeout, logger=self.logger, timeout_fun=timeout_fun - ) + + if retries < 0: + raise ValueError("Number of retries must be greater than zero.") + + for i in range(0, retries + 1): # + 1 for the initial try + try: + + timeout = timeout_fun() + self.debug( + "Opening socket to {}:{} (timeout {} secs).".format( + host, port, timeout + ) + ) + return SimpleSocket( + host, + port, + timeout=timeout, + logger=self.logger, + timeout_fun=timeout_fun, + ) + + except Exception as e: + self.warning( + f"Failed to establish connection to {host}:{port}, Try #{i+1} of {retries+1} ", + exc_info=e, + ) + continue + + self.error(f"Failed to establish connection to {host}:{port}") + raise OfflineException("Failed establishing connection to service.") @property def http_useragent(self) -> str: