diff --git a/htpclient/binarydownload.py b/htpclient/binarydownload.py index d6a487e..6c6f3d8 100644 --- a/htpclient/binarydownload.py +++ b/htpclient/binarydownload.py @@ -7,6 +7,7 @@ from htpclient.config import Config from htpclient.download import Download +from htpclient.helpers import retrieveBinary from htpclient.initialize import Initialize from htpclient.jsonRequest import JsonRequest from htpclient.dicts import * @@ -63,7 +64,7 @@ def check_client_version(self): def __check_utils(self): path = '7zr' + Initialize.get_os_extension() - if not os.path.isfile(path): + if not retrieveBinary(path): query = copy_and_set_token(dict_downloadBinary, self.config.get_value('token')) query['type'] = '7zr' req = JsonRequest(query) @@ -80,7 +81,7 @@ def __check_utils(self): Download.download(ans['executable'], path) os.chmod(path, os.stat(path).st_mode | stat.S_IEXEC) path = 'uftpd' + Initialize.get_os_extension() - if not os.path.isfile(path) and self.config.get_value('multicast'): + if not retrieveBinary(path) and self.config.get_value('multicast'): query = copy_and_set_token(dict_downloadBinary, self.config.get_value('token')) query['type'] = 'uftpd' req = JsonRequest(query) @@ -121,10 +122,11 @@ def check_prince(self): logging.error("Download of prince failed!") sleep(5) return False - if Initialize.get_os() == 1: - os.system("7zr" + Initialize.get_os_extension() + " x -otemp prince.7z") - else: - os.system("./7zr" + Initialize.get_os_extension() + " x -otemp prince.7z") + zr_bin = retrieveBinary("7zr" + Initialize.get_os_extension()) + if not zr_bin: + logging.error("7zr not found, cannot extract archive") + return False + os.system(zr_bin + " x -otemp prince.7z") for name in os.listdir("temp"): # this part needs to be done because it is compressed with the main subfolder of prince if os.path.isdir("temp/" + name): os.rename("temp/" + name, "prince") @@ -160,10 +162,11 @@ def check_preprocessor(self, task): logging.error("Download of preprocessor failed!") sleep(5) return False - if Initialize.get_os() == 1: - os.system(f"7zr{Initialize.get_os_extension()} x -otemp temp.7z") - else: - os.system(f"./7zr{Initialize.get_os_extension()} x -otemp temp.7z") + zr_bin = retrieveBinary("7zr" + Initialize.get_os_extension()) + if not zr_bin: + logging.error("7zr not found, cannot extract archive") + return False + os.system(f"{zr_bin} x -otemp temp.7z") for name in os.listdir("temp"): # this part needs to be done because it is compressed with the main subfolder of prince if os.path.isdir(Path('temp', name)): os.rename(Path('temp', name), path) @@ -200,13 +203,12 @@ def check_version(self, cracker_id): # we need to extract the 7zip temp_folder = Path(self.config.get_value('crackers-path'), 'temp') zip_file = Path(self.config.get_value('crackers-path'), f'{cracker_id}.7z') - - if Initialize.get_os() == 1: - # Windows - cmd = f'7zr{Initialize.get_os_extension()} x -o"{temp_folder}" "{zip_file}"' - else: - # Linux - cmd = f"./7zr{Initialize.get_os_extension()} x -o'{temp_folder}' '{zip_file}'" + zrbinary = retrieveBinary("7zr" + Initialize.get_os_extension()) + if not zrbinary: + logging.error("7zr not found, cannot extract archive") + sleep(5) + return False + cmd = f'{zrbinary} x -o"{temp_folder}" "{zip_file}"' os.system(cmd) # Clean up 7zip diff --git a/htpclient/files.py b/htpclient/files.py index 4ac12f9..430e226 100644 --- a/htpclient/files.py +++ b/htpclient/files.py @@ -7,6 +7,7 @@ from htpclient.config import Config from htpclient.download import Download +from htpclient.helpers import retrieveBinary from htpclient.initialize import Initialize from htpclient.jsonRequest import JsonRequest from htpclient.dicts import * @@ -106,11 +107,10 @@ def check_files(self, files, task_id): if os.path.splitext(file_localpath)[1] == '.7z' and not os.path.isfile(txt_file): # extract if needed files_path = Path(self.config.get_value('files-path')) - if Initialize.get_os() == 1: - # Windows - cmd = f'7zr{Initialize.get_os_extension()} x -aoa -o"{files_path}" -y "{file_localpath}"' - else: - # Linux - cmd = f"./7zr{Initialize.get_os_extension()} x -aoa -o'{files_path}' -y '{file_localpath}'" + zr_bin = retrieveBinary("7zr" + Initialize.get_os_extension()) + if not zr_bin: + logging.error("7zr not found, cannot extract archive") + return False + cmd = f'{zr_bin} x -aoa -o"{files_path}" -y "{file_localpath}"' os.system(cmd) return True diff --git a/htpclient/helpers.py b/htpclient/helpers.py index eb4d1df..7f6dff9 100644 --- a/htpclient/helpers.py +++ b/htpclient/helpers.py @@ -8,6 +8,7 @@ from pathlib import Path import os +import shutil import subprocess from htpclient.dicts import copy_and_set_token, dict_clientError @@ -66,7 +67,10 @@ def start_uftpd(os_extension, config): subprocess.check_output("killall -s 9 uftpd", shell=True) # stop running service to make sure we can start it again except subprocess.CalledProcessError: pass # ignore in case uftpd was not running - path = './uftpd' + os_extension + path = retrieveBinary('uftpd' + os_extension) + if not path: + logging.error("uftpd binary not found, cant do multicast") + return cmd = path + ' ' if config.get_value('multicast-device'): cmd += "-I " + config.get_value('multicast-device') + ' ' @@ -132,3 +136,20 @@ def update_files(command, prince=False): def escape_ansi(line): ansi_escape = re.compile(r'(\x9B|\x1B\[)[0-?]*[ -/]*[@-~]') return ansi_escape.sub('', line) + +# function to retrieve a system or local binary. +def retrieveBinary(binary): + cwd = Path.cwd() + # use full path so that it works on Windows and Linux + local_binary = cwd / binary + + # First check if there is a local binary and use that if it is there + if local_binary.exists() and local_binary.is_file(): + return str(local_binary) + + # Fall back on system binary + system_binary = shutil.which(binary) + + if system_binary + return system_binary + return None \ No newline at end of file