Update app.py

app.py

@@ -12,6 +12,7 @@ import huggingface_hub
 import torch
 import base64
 from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
+from cryptography.hazmat.primitives import padding as sym_padding # Use symmetric padding
 from cryptography.hazmat.backends import default_backend
 import yaml
 from gradio_logsview.logsview import Log, LogsView, LogsViewRunner
@@ -32,7 +33,7 @@ MARKDOWN_DESCRIPTION = """
 
 The fastest way to perform a model merge 🔥
 
-Specify a YAML configuration file (see examples below) and a HF token and this app will perform the merge and upload the merged model to your user profile.
+Specify a YAML configuration file (see examples below) and a HF token and this app will perform the merge and upload the merged model to your user profile. Includes encryption for the `mergekit_config.yml` and a tool to decrypt it.
 """
 
 MARKDOWN_ARTICLE = """
@@ -73,52 +74,119 @@ This Space is heavily inspired by LazyMergeKit by Maxime Labonne (see [Colab](ht
 
 examples = [[str(f)] for f in pathlib.Path("examples").glob("*.yaml")]
 
+def _prepare_key(key: str) -> bytes:
+    """Pads or truncates the key to 32 bytes (256 bits) for AES."""
+    key_bytes = key.encode('utf-8')
+    if len(key_bytes) < 32:
+        return key_bytes + b'\0' * (32 - len(key_bytes))
+    else:
+        return key_bytes[:32]
 
-def encrypt_file(file_path, key):
+def encrypt_file(file_path, key: str) -> bool:
     """
-    Encrypt the contents of a file using AES encryption with the provided key.
-    
+    Encrypt the contents of a file using AES-256-CBC encryption with the provided key.
+    The output is Base64 encoded.
+
     Args:
         file_path: Path to the file to encrypt (pathlib.Path or string)
-        key: Encryption key
-        
+        key: Encryption key string.
+
     Returns:
         bool: True if encryption was successful, False otherwise
     """
     try:
         file_path = pathlib.Path(file_path)
         if not file_path.exists():
+            print(f"Encryption error: File not found at {file_path}")
             return False
-        
-        # Ensure key is 32 bytes (256 bits)
-        key_bytes = key.encode('utf-8')
-        key_bytes = key_bytes + b'\0' * (32 - len(key_bytes)) if len(key_bytes) < 32 else key_bytes[:32]
-        
-        # Generate a random IV
+
+        key_bytes = _prepare_key(key)
+
+        # Generate a random IV (Initialization Vector) - 16 bytes for AES
         iv = os.urandom(16)
-        
-        # Create an encryptor
+
+        # Create an AES cipher instance with CBC mode
         cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv), backend=default_backend())
         encryptor = cipher.encryptor()
-        
-        # Read file content
+
+        # Use PKCS7 padding
+        padder = sym_padding.PKCS7(algorithms.AES.block_size).padder()
+
         with open(file_path, 'rb') as f:
-            content = f.read()
-        
-        # Pad the content to be a multiple of 16 bytes
-        padding = 16 - (len(content) % 16)
-        content += bytes([padding]) * padding
-        
-        # Encrypt and write back
-        encrypted = iv + encryptor.update(content) + encryptor.finalize()
+            plaintext = f.read()
+
+        # Pad the data
+        padded_data = padder.update(plaintext) + padder.finalize()
+
+        # Encrypt the padded data
+        ciphertext = encryptor.update(padded_data) + encryptor.finalize()
+
+        # Prepend the IV to the ciphertext and base64 encode the result
+        encrypted_data_with_iv = base64.b64encode(iv + ciphertext)
+
+        # Write the base64 encoded encrypted data back to the file
         with open(file_path, 'wb') as f:
-            f.write(base64.b64encode(encrypted))
-            
+            f.write(encrypted_data_with_iv)
+
         return True
     except Exception as e:
         print(f"Encryption error: {e}")
         return False
 
+def decrypt_file_content(file_input, key: str) -> str:
+    """
+    Decrypts the content of an uploaded file using AES-256-CBC and returns the result.
+    Assumes the file content is Base64 encoded IV + ciphertext.
+
+    Args:
+        file_input: Gradio File component output (temporary file object).
+        key: Decryption key string.
+
+    Returns:
+        str: Decrypted content as a UTF-8 string, or an error message.
+    """
+    if file_input is None:
+        return "Error: No file provided for decryption."
+    if not key:
+        return "Error: Decryption key cannot be empty."
+
+    try:
+        file_path = file_input.name # Get the temporary file path from Gradio
+        key_bytes = _prepare_key(key)
+
+        with open(file_path, 'rb') as f:
+            base64_encoded_data = f.read()
+
+        # Decode from Base64
+        encrypted_data_with_iv = base64.b64decode(base64_encoded_data)
+
+        # Extract the IV (first 16 bytes)
+        iv = encrypted_data_with_iv[:16]
+        # Extract the ciphertext (the rest)
+        ciphertext = encrypted_data_with_iv[16:]
+
+        # Create an AES cipher instance with CBC mode for decryption
+        cipher = Cipher(algorithms.AES(key_bytes), modes.CBC(iv), backend=default_backend())
+        decryptor = cipher.decryptor()
+
+        # Decrypt the data
+        padded_plaintext = decryptor.update(ciphertext) + decryptor.finalize()
+
+        # Unpad the data using PKCS7
+        unpadder = sym_padding.PKCS7(algorithms.AES.block_size).unpadder()
+        plaintext = unpadder.update(padded_plaintext) + unpadder.finalize()
+
+        # Decode the plaintext from bytes to string (assuming UTF-8)
+        return plaintext.decode('utf-8')
+
+    except (ValueError, TypeError) as e:
+         # Catches Base64 decoding errors, incorrect key type errors
+        return f"Decryption Error: Invalid input data or key format. ({e})"
+    except Exception as e:
+        # Catches padding errors (often due to wrong key), or other crypto issues
+        print(f"Decryption error details: {e}")
+        return f"Decryption Failed: Likely incorrect key or corrupted file. Error: {type(e).__name__}"
+
 
 def merge(yaml_config: str, hf_token: str, repo_name: str, cipher_key: str) -> Iterable[List[Log]]:
     runner = LogsViewRunner()
@@ -136,7 +204,7 @@ def merge(yaml_config: str, hf_token: str, repo_name: str, cipher_key: str) -> I
     if not hf_token:
         yield runner.log("No HF token provided. A valid token is required for uploading.", level="ERROR")
         return
-    
+
     # Validate that the token works by trying to get user info
     try:
         api = huggingface_hub.HfApi(token=hf_token)
@@ -146,10 +214,14 @@ def merge(yaml_config: str, hf_token: str, repo_name: str, cipher_key: str) -> I
         yield runner.log(f"Invalid HF token: {e}", level="ERROR")
         return
 
-    # Set default cipher key if none provided
+    # Use default key if none provided, but log a warning
     if not cipher_key:
-        cipher_key = "default_key"  # Fallback key, though we should encourage users to set their own
-        yield runner.log("No cipher key provided. Using default key (not recommended).", level="WARNING")
+        cipher_key = "default_insecure_key" # Make default explicitely insecure sounding
+        yield runner.log("No cipher key provided. Using a default, insecure key. Please provide your own key for security.", level="WARNING")
+    elif cipher_key == "Default": # Check against the placeholder value
+         cipher_key = "default_insecure_key" # Treat placeholder as no key provided
+         yield runner.log("Default placeholder key detected. Using an insecure key. Please provide your own key.", level="WARNING")
+
 
     with tempfile.TemporaryDirectory(ignore_cleanup_errors=True) as tmpdirname:
         tmpdir = pathlib.Path(tmpdirname)
@@ -182,37 +254,49 @@ def merge(yaml_config: str, hf_token: str, repo_name: str, cipher_key: str) -> I
 
         if runner.exit_code != 0:
             yield runner.log("Merge failed. Deleting repo as no model is uploaded.", level="ERROR")
-            api.delete_repo(repo_url.repo_id)
+            try:
+                api.delete_repo(repo_url.repo_id)
+                yield runner.log(f"Repo {repo_url.repo_id} deleted.")
+            except Exception as delete_e:
+                 yield runner.log(f"Failed to delete repo {repo_url.repo_id}: {delete_e}", level="WARNING")
             return
 
-        yield runner.log("Model merged successfully. Uploading to HF.")
-        
-        # Delete Readme.md if it exists (case-insensitive check)
+        yield runner.log("Model merged successfully. Preparing for upload.")
+
+        # ---- Encryption Step ----
         merge_dir = merged_path / "merge"
+        config_yml_path = merge_dir / "mergekit_config.yml"
+
+        if config_yml_path.exists():
+            yield runner.log(f"Found {config_yml_path.name}. Encrypting...")
+            if encrypt_file(config_yml_path, cipher_key):
+                yield runner.log(f"Successfully encrypted {config_yml_path.name} with provided key.")
+            else:
+                yield runner.log(f"Failed to encrypt {config_yml_path.name}. Uploading unencrypted.", level="ERROR")
+        else:
+            yield runner.log(f"{config_yml_path.name} not found in merge output, nothing to encrypt.", level="INFO")
+        # ---- End Encryption Step ----
+
+        # Delete Readme.md if it exists (case-insensitive check) before upload
         readme_deleted = False
-        for file in merge_dir.glob("*"):
-            if file.name.lower() == "readme.md":
-                try:
+        try:
+            for file in merge_dir.glob("*"):
+                if file.name.lower() == "readme.md":
                     file.unlink()
                     readme_deleted = True
                     yield runner.log(f"Deleted {file.name} file before upload")
-                except Exception as e:
-                    yield runner.log(f"Error deleting {file.name}: {e}", level="WARNING")
-        
+                    break # Assume only one readme
+        except Exception as e:
+            yield runner.log(f"Error deleting Readme.md: {e}", level="WARNING")
+
         if not readme_deleted:
-            yield runner.log("No Readme.md file found to delete", level="INFO")
-        
-        # Encrypt mergekit_config.yml if it exists
-        config_yml_path = merged_path / "merge" / "mergekit_config.yml"
-        if not config_yml_path.exists():
-            yield runner.log("mergekit_config.yml not found, nothing to encrypt", level="INFO")
-        elif encrypt_file(config_yml_path, cipher_key):
-            yield runner.log("Encrypted mergekit_config.yml with provided key")
-        
+            yield runner.log("No Readme.md file found to delete.", level="INFO")
+
+        yield runner.log("Uploading merged model files to HF.")
         yield from runner.run_python(
             api.upload_folder,
             repo_id=repo_url.repo_id,
-            folder_path=merged_path / "merge",
+            folder_path=merge_dir, # Upload from the 'merge' subdirectory
         )
         yield runner.log(f"Model successfully uploaded to HF: {repo_url.repo_id}")
 
@@ -234,43 +318,65 @@ NEXT_RESTART = f"Next Restart: {next_run_time_utc.strftime('%Y-%m-%d %H:%M:%S')}
 with gr.Blocks() as demo:
     gr.Markdown(MARKDOWN_DESCRIPTION)
     gr.Markdown(NEXT_RESTART)
-    
-    with gr.Row():
-        filename = gr.Textbox(visible=False, label="filename")
-        config = gr.Code(language="yaml", lines=10, label="config.yaml")
-        with gr.Column():
-            token = gr.Textbox(
-                lines=1,
-                label="HF Write Token",
-                info="https://hf.co/settings/token",
-                type="password",
-                placeholder="Required for model upload.",
-            )
-            repo_name = gr.Textbox(
-                lines=1,
-                label="Repo name",
-                placeholder="Optional. Will create a random name if empty.",
+
+    with gr.Tabs():
+        with gr.TabItem("Merge Model"):
+            with gr.Row():
+                filename = gr.Textbox(visible=False, label="filename")
+                config = gr.Code(language="yaml", lines=10, label="config.yaml")
+                with gr.Column():
+                    token = gr.Textbox(
+                        lines=1,
+                        label="HF Write Token",
+                        info="https://hf.co/settings/token",
+                        type="password",
+                        placeholder="Required for model upload.",
+                    )
+                    repo_name = gr.Textbox(
+                        lines=1,
+                        label="Repo name",
+                        placeholder="Optional. Will create a random name if empty.",
+                    )
+                    cipher_key = gr.Textbox(
+                        lines=1,
+                        label="Encryption Key",
+                        type="password",
+                        info="Key used to encrypt the generated mergekit_config.yml file before upload. Leave blank or 'Default' for no encryption (or insecure default).",
+                        placeholder="Enter your secret key here",
+                        value="Default" # Set a default placeholder
+                    )
+            button = gr.Button("Merge and Upload", variant="primary")
+            logs = LogsView(label="Merge Progress / Terminal output")
+            gr.Examples(
+                examples,
+                fn=lambda s: (s,),
+                run_on_click=True,
+                label="Merge Examples",
+                inputs=[filename],
+                outputs=[config],
             )
-            cipher_key = gr.Textbox(
-                lines=1,
-                label="Encryption Key",
-                type="password",
-                placeholder="Key used to encrypt the config file.",
-                value="Default"
+            gr.Markdown(MARKDOWN_ARTICLE)
+
+            button.click(fn=merge, inputs=[config, token, repo_name, cipher_key], outputs=[logs])
+
+        with gr.TabItem("Decrypt Configuration"):
+            gr.Markdown("Upload an encrypted `mergekit_config.yml` file and provide the key to decrypt it.")
+            with gr.Row():
+                decrypt_file_input = gr.File(label="Upload Encrypted mergekit_config.yml")
+                decrypt_key_input = gr.Textbox(
+                    lines=1,
+                    label="Decryption Key",
+                    type="password",
+                    placeholder="Enter the key used for encryption",
+                )
+            decrypt_button = gr.Button("Decrypt File", variant="secondary")
+            decrypted_output = gr.Code(language="yaml", label="Decrypted Configuration", lines=15, interactive=False)
+
+            decrypt_button.click(
+                fn=decrypt_file_content,
+                inputs=[decrypt_file_input, decrypt_key_input],
+                outputs=[decrypted_output]
             )
-    button = gr.Button("Merge", variant="primary")
-    logs = LogsView(label="Terminal output")
-    gr.Examples(
-        examples,
-        fn=lambda s: (s,),
-        run_on_click=True,
-        label="Examples",
-        inputs=[filename],
-        outputs=[config],
-    )
-    gr.Markdown(MARKDOWN_ARTICLE)
-
-    button.click(fn=merge, inputs=[config, token, repo_name, cipher_key], outputs=[logs])
-
-
-demo.queue(default_concurrency_limit=1).launch()
+
+
+demo.queue(default_concurrency_limit=1).launch()