diff --git a/runner/scripts/em_smb.py b/runner/scripts/em_smb.py index f98061b2..f8f193df 100644 --- a/runner/scripts/em_smb.py +++ b/runner/scripts/em_smb.py @@ -190,10 +190,15 @@ def _walk( new_path = str(Path(directory).joinpath(dirname)) yield from self._walk(new_path) - def __load_file(self, file_name: str) -> IO[str]: + def __load_file(self, file_name: str, index: int, length: int) -> IO[str]: + RunnerLog( + self.task, self.run_id, 10, f"({index} of {length}) downloading {file_name}" + ) + director = urllib.request.build_opener(SMBHandler) password = em_decrypt(self.password, app.config["PASS_KEY"]) + open_file_for_read = director.open( f"smb://{self.username}:{password}@{self.server_name},{self.server_ip}/{self.share_name}/{file_name}" ) @@ -279,9 +284,12 @@ def read(self, file_name: str) -> List[IO[str]]: ) # if a file was found, try to open. - return [self.__load_file(file_name) for file_name in file_list] + return [ + self.__load_file(file_name, i, len(file_list)) + for i, file_name in enumerate(file_list, 1) + ] - return [self.__load_file(file_name)] + return [self.__load_file(file_name, 1, 1)] except BaseException as e: raise RunnerException( self.task, diff --git a/runner/scripts/smb_fix.py b/runner/scripts/smb_fix.py index 817071d0..7cdaaede 100644 --- a/runner/scripts/smb_fix.py +++ b/runner/scripts/smb_fix.py @@ -8,6 +8,7 @@ import socket import sys import tempfile +import time import urllib.error import urllib.parse import urllib.request @@ -29,6 +30,7 @@ from urllib.response import addinfourl from nmb.NetBIOS import NetBIOS +from smb.base import NotConnectedError, SMBTimeout from smb.SMBConnection import SMBConnection USE_NTLM = True @@ -90,7 +92,28 @@ def smb_open(self, req): conn = SMBConnection( user, passwd, myname, server_name, domain=domain, use_ntlm_v2=USE_NTLM ) - conn.connect(host, port) + + # retry + retry = 0 + while True: + try: + connected = conn.connect(host, port, timeout=120) + if not connected: + raise AssertionError() + break + + except ( + AssertionError, + ConnectionResetError, + SMBTimeout, + NotConnectedError, + ) as e: + if retry <= 10: + retry += 1 + time.sleep(5) # wait 5 sec before retrying + continue + + raise ValueError(f"Connection failed.\n{e}") headers = email.message.Message() if req.data: