Made it possible to let the agent use system binaries if exists (#64)

* Made it possible to let the agent use system binaries if they already exists


---------

Co-authored-by: jessevz <jesse.van.zutphen@nfi.nl>

Author: jessevz

htpclient/binarydownload.py

@@ -7,6 +7,7 @@
 
 from htpclient.config import Config
 from htpclient.download import Download
+from htpclient.helpers import retrieve_binary
 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 retrieve_binary(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 retrieve_binary(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 = retrieve_binary("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 = retrieve_binary("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}'"
+                zr_bin = retrieve_binary("7zr" + Initialize.get_os_extension())
+                if not zr_bin:
+                    logging.error("7zr not found, cannot extract archive")
+                    sleep(5)
+                    return False
+                cmd = f'{zr_bin} x -o"{temp_folder}" "{zip_file}"'
                 os.system(cmd)
 
                 # Clean up 7zip

htpclient/files.py

@@ -7,6 +7,7 @@
 
 from htpclient.config import Config
 from htpclient.download import Download
+from htpclient.helpers import retrieve_binary
 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 = retrieve_binary("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

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 = retrieve_binary('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 retrieve_binary(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