Spaces:
Sleeping
Sleeping
| # keylock/app.py | |
| import gradio as gr | |
| from PIL import Image, ImageFont | |
| import tempfile | |
| import os | |
| import json | |
| import logging | |
| import traceback | |
| import base64 | |
| import io | |
| # Import specific exceptions from cryptography | |
| from cryptography.exceptions import InvalidTag, InvalidKey | |
| import core # Use relative import for core module | |
| __version__ ="v1"# Import version for footer | |
| app_logger = logging.getLogger("keylock_app") | |
| # Configure logger if not already configured | |
| if not app_logger.hasHandlers(): | |
| handler = logging.StreamHandler() | |
| # Using the format from the original code for consistency | |
| formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s') | |
| handler.setFormatter(formatter) | |
| app_logger.addHandler(handler) | |
| app_logger.setLevel(logging.INFO) # Keep INFO level for app logs | |
| # Theming colors (using fallback names for broader compatibility if needed) | |
| try: | |
| # Attempt to use Gradio's theme colors if available | |
| gr.themes.colors.blue # Check if colors exist | |
| blue_color = gr.themes.colors.blue | |
| sky_color = gr.themes.colors.sky | |
| slate_color = gr.themes.colors.slate | |
| cyan_color = gr.themes.colors.cyan | |
| neutral_color = gr.themes.colors.neutral | |
| except AttributeError: | |
| # Fallback for older Gradio versions or if colors structure changes | |
| app_logger.warning("gr.themes.colors not found. Using placeholder colors for themes.") | |
| class FallbackColors: # Basic fallback colors | |
| blue = "blue"; sky = "skyblue"; slate = "slategray"; cyan = "cyan"; neutral = "gray" | |
| blue_color = FallbackColors.blue | |
| sky_color = FallbackColors.sky | |
| slate_color = FallbackColors.slate | |
| cyan_color = FallbackColors.cyan | |
| neutral_color = FallbackColors.neutral | |
| ICON_EMBED = "🔒" # Changed icon to reflect encryption | |
| ICON_EXTRACT = "🔓" # Changed icon to reflect decryption | |
| ICON_KEY = "🔑" # New icon for key management | |
| def pil_to_base64_html(pil_image, max_width_px=None): | |
| buffered = io.BytesIO() | |
| # Save as PNG to preserve potential LSB data (though preview isn't used for stego) | |
| # Ensure it's RGB for saving if the original was RGBA (overlay conversion) | |
| if pil_image.mode == 'RGBA': | |
| pil_image = pil_image.convert('RGB') | |
| pil_image.save(buffered, format="PNG") | |
| img_str = base64.b64encode(buffered.getvalue()).decode() | |
| style = f"max-width:{max_width_px}px; height:auto; border:1px solid #ccc; display:block; margin-left:auto; margin-right:auto;" if max_width_px else "border:1px solid #ccc; display:block; margin-left:auto; margin-right:auto;" | |
| return f"<div style='text-align:center;'><img src='data:image/png;base64,{img_str}' alt='Stego Image Preview' style='{style}'/></div>" | |
| def gradio_embed_data(kv_string: str, recipient_public_key_pem: str, | |
| input_image_pil: Image.Image, generate_carrier_flag: bool, | |
| show_keys_on_image_flag: bool, output_filename_base: str): | |
| """Handles the embedding process in Gradio.""" | |
| output_html_img_str, status_msg, dl_file_path = None, "An error occurred.", None | |
| # Validate inputs | |
| if not recipient_public_key_pem or not recipient_public_key_pem.strip(): | |
| return None, "Error: Recipient's Public Key is required.", None | |
| if not kv_string or not kv_string.strip(): | |
| return None, "Error: Key-Value data cannot be empty.", None | |
| try: | |
| # Parse KV string | |
| data_dict = core.parse_kv_string_to_dict(kv_string) | |
| if not data_dict: | |
| return None, "Error: Parsed Key-Value data is empty or only contains comments/empty lines.", None | |
| serial_data = json.dumps(data_dict, indent=None).encode('utf-8') # Compact JSON for embedding | |
| # Prepare carrier image | |
| original_format_note = "" | |
| if generate_carrier_flag or input_image_pil is None: | |
| app_logger.info("Generating new carrier image.") | |
| carrier_img = core.generate_keylock_carrier_image() | |
| else: | |
| app_logger.info(f"Using uploaded carrier image (format: {getattr(input_image_pil, 'format', 'Unknown')}).") | |
| carrier_img = input_image_pil.copy() | |
| if hasattr(input_image_pil, 'format') and input_image_pil.format and input_image_pil.format.upper() != 'PNG': | |
| # Note about format conversion if not PNG | |
| original_format_note = ( | |
| f"Input carrier image was format '{input_image_pil.format}'. " | |
| f"It will be converted to RGB for LSB embedding and saved as PNG. " | |
| f"If original was lossy (e.g., JPEG), quality is preserved from upload; " | |
| f"if it had transparency (e.g., GIF, WebP with alpha), transparency will be lost." | |
| ) | |
| app_logger.warning(original_format_note) | |
| # Ensure carrier is RGB before overlay and embedding | |
| carrier_img = carrier_img.convert("RGB") | |
| # Encrypt the data using hybrid RSA/AES | |
| app_logger.info(f"Encrypting {len(serial_data)} bytes of data with hybrid encryption.") | |
| # Need bytes for public key input | |
| encrypted_payload = core.encrypt_data_hybrid(serial_data, recipient_public_key_pem.encode('utf-8')) | |
| app_logger.info(f"Encrypted payload size (from crypto layer): {len(encrypted_payload)} bytes.") | |
| # Add visual overlay | |
| keys_for_overlay = list(data_dict.keys()) if show_keys_on_image_flag else None | |
| overlay_title = "KeyLock: RSA Encrypted Data Embedded" | |
| final_carrier_with_overlay = core.draw_key_list_dropdown_overlay( | |
| carrier_img, | |
| keys=keys_for_overlay, | |
| title=overlay_title | |
| ) | |
| app_logger.info("Visual overlay added.") | |
| # Embed the encrypted payload using the LSB layer | |
| # The LSB function will add its own 4-byte header for this payload's length. | |
| app_logger.info(f"Embedding {len(encrypted_payload)} bytes (crypto payload) into image pixels.") | |
| stego_final_img = core.embed_data_in_image(final_carrier_with_overlay, encrypted_payload) | |
| # The total size embedded in the image LSBs is len(encrypted_payload) + core.LENGTH_HEADER_SIZE. | |
| app_logger.info(f"Data embedding complete. Total bytes embedded in LSB (including header): {len(encrypted_payload) + core.LENGTH_HEADER_SIZE} bytes.") | |
| # Save the stego image to a temporary file | |
| fname_base = "".join(c if c.isalnum() or c in ('_','-') else '_' for c in output_filename_base.strip()) or "keylock_img_rsa" | |
| temp_fp = None | |
| try: | |
| with tempfile.NamedTemporaryFile(prefix=fname_base+"_", suffix=".png", delete=False) as tmp: | |
| stego_final_img.save(tmp, format="PNG") | |
| temp_fp = tmp.name | |
| app_logger.info(f"Stego image saved to temp file: {temp_fp}") | |
| except Exception as e: | |
| app_logger.error(f"Error saving temp file: {e}", exc_info=True) | |
| status_msg = f"Error saving output image: {e}" | |
| return None, status_msg, None | |
| # Prepare output for Gradio | |
| output_html_img_str = pil_to_base64_html(stego_final_img, max_width_px=480) | |
| status_msg = (f"Data embedded into '{os.path.basename(temp_fp)}'.\n" | |
| f"{original_format_note}" | |
| f"Image contains visual \"{overlay_title}\" overlay " | |
| f"{'(with key list)' if show_keys_on_image_flag and keys_for_overlay else ''} " | |
| f"and your LSB-encoded encrypted data.\n" | |
| f"Secret data size: {len(serial_data)}B (raw JSON).\n" | |
| f"Crypto payload size: {len(encrypted_payload)}B (Hybrid RSA/AES).\n" | |
| f"Total bytes embedded in LSB (including header): {len(encrypted_payload) + core.LENGTH_HEADER_SIZE}B.\n" | |
| f"Recipient needs their Private Key to extract.") | |
| return output_html_img_str, status_msg, temp_fp | |
| except (ValueError, RuntimeError, TypeError, InvalidKey) as e: # Catch InvalidKey here too | |
| # Catch errors raised by core functions (e.g., invalid key, data too large, crypto errors) | |
| app_logger.error(f"Embed process error: {e}", exc_info=True) | |
| return None, f"Error during embedding: {str(e)}", None | |
| except Exception as e: | |
| # Catch any other unexpected errors | |
| app_logger.error(f"Unexpected Embed Error: {e}", exc_info=True) | |
| return None, f"An unexpected error occurred: {str(e)}", None | |
| def gradio_extract_data(stego_image_pil: Image.Image, recipient_private_key_pem: str): | |
| """Handles the extraction process in Gradio.""" | |
| if stego_image_pil is None: | |
| return "Error: No image provided.", "Error: No image." | |
| if not recipient_private_key_pem or not recipient_private_key_pem.strip(): | |
| return "Error: Your Private Key is required for decryption.", "Error: Private Key required." | |
| try: | |
| # Ensure image is RGB for extraction | |
| stego_image_rgb = stego_image_pil.convert("RGB") | |
| if hasattr(stego_image_pil, 'format') and stego_image_pil.format and stego_image_pil.format.upper() != "PNG": | |
| app_logger.warning(f"Uploaded image for extraction is format '{stego_image_pil.format}', not PNG. LSB data may be compromised if not the original KeyLock file.") | |
| # Extract the embedded payload using the LSB layer | |
| # This function (restored) reads the LSB header and returns *only* the data payload after it. | |
| app_logger.info("Extracting data from image pixels using LSB.") | |
| extracted_crypto_payload = core.extract_data_from_image(stego_image_rgb) | |
| app_logger.info(f"Extracted crypto payload size (from LSB layer): {len(extracted_crypto_payload)} bytes.") | |
| if not extracted_crypto_payload: | |
| # The original extract_data_from_image would return b"" if header length was 0. | |
| # This is handled in core's decrypt_data_hybrid, which checks min size. | |
| pass # Allow empty payload to proceed to core's decryption logic for specific errors | |
| # Decrypt the crypto payload using hybrid RSA/AES | |
| # This function now expects the raw crypto payload bytes, without the LSB header. | |
| app_logger.info("Decrypting extracted crypto payload with hybrid encryption.") | |
| # Need bytes for private key input | |
| decrypted_bytes = core.decrypt_data_hybrid(extracted_crypto_payload, recipient_private_key_pem.encode('utf-8')) | |
| app_logger.info(f"Decrypted data size (raw): {len(decrypted_bytes)} bytes.") | |
| # Attempt to parse decrypted data (usually JSON) | |
| try: | |
| data = json.loads(decrypted_bytes.decode('utf-8')) | |
| # Pretty print JSON | |
| txt, stat = json.dumps(data, indent=2), "Data extracted successfully (JSON)." | |
| app_logger.info("Successfully decrypted and parsed as JSON.") | |
| except (json.JSONDecodeError, UnicodeDecodeError): | |
| # Fallback if not valid UTF-8 JSON | |
| try: | |
| txt = "Decrypted (UTF-8, not JSON):\n" + decrypted_bytes.decode('utf-8', errors='replace') # Use replace for invalid chars | |
| stat = "Warning: Decrypted as UTF-8, but not valid JSON." | |
| app_logger.warning("Decrypted data is not valid JSON, but appears to be UTF-8.") | |
| except UnicodeDecodeError: | |
| # Fallback if not valid UTF-8 | |
| txt = "Decrypted (raw hex, not JSON/UTF-8):\n" + decrypted_bytes.hex() | |
| stat = "Warning: Decrypted bytes are not valid UTF-8 or JSON. Displaying as hex." | |
| app_logger.warning("Decrypted data is not valid UTF-8. Displaying as hex.") | |
| return txt, stat | |
| # Corrected the exception catching here to import InvalidTag and InvalidKey | |
| except (ValueError, RuntimeError, InvalidTag, InvalidKey, TypeError) as e: | |
| # Catch errors raised by core functions (e.g., extraction errors, decryption errors) | |
| app_logger.error(f"Extract process error: {e}", exc_info=True) | |
| # Pass through specific error messages from core | |
| # Use type(e).__name__ to show the actual exception type in the UI status | |
| return f"Error during extraction: {str(e)}", f"Extraction Failed: {type(e).__name__} Error." | |
| except Exception as e: | |
| # Catch any other unexpected errors | |
| app_logger.error(f"Unexpected Extract Error: {e}", exc_info=True) | |
| return f"An unexpected error occurred: {str(e)}", f"Extraction Failed: Unexpected Error." | |
| # Corrected the constant name here | |
| def gradio_generate_keys(key_size_bits: int = core.RSA_KEY_SIZE_DEFAULT): | |
| """Generates RSA keys and returns them in PEM format.""" | |
| status = "" | |
| try: | |
| private_pem, public_pem = core.generate_rsa_key_pair(key_size_bits) | |
| status = f"RSA {key_size_bits}-bit key pair generated successfully." | |
| app_logger.info(status) | |
| return private_pem.decode('utf-8'), public_pem.decode('utf-8'), status | |
| except Exception as e: | |
| app_logger.error(f"Key generation error: {e}", exc_info=True) | |
| status = f"Error generating keys: {str(e)}" | |
| return "", "", status | |
| def build_interface(): | |
| """Builds the Gradio interface with RSA key management.""" | |
| custom_theme = gr.themes.Base( | |
| primary_hue="indigo", | |
| secondary_hue="cyan", | |
| neutral_hue="zinc", | |
| text_size="md", | |
| spacing_size="md", | |
| radius_size="sm", | |
| font=["System UI", "sans-serif"] # Use system fonts first | |
| ).set( | |
| # Example adjustments for better dark theme appearance | |
| button_secondary_background_fill_hover=blue_color, # Use theme colors | |
| button_secondary_border_color_hover=blue_color, | |
| button_secondary_text_color_hover='white', | |
| # Add more theme adjustments as desired | |
| ) | |
| # Using inline CSS for custom backgrounds and colors as in original | |
| custom_css = """ | |
| body { | |
| background: linear-gradient(to right bottom, rgb(44, 62, 80), rgb(80 168 255)); /* Dark blue gradient */ | |
| color: #ecf0f1; /* Light text color */ | |
| } | |
| span { | |
| color: inherit; /* Inherit color from parent */ | |
| } | |
| .gradio-container { | |
| background: transparent !important; /* Make main container transparent */ | |
| } | |
| .gr-box, .gr-panel, .gr-pill { | |
| background-color: rgba(44, 62, 80, 0.8) !important; /* Semi-transparent dark boxes */ | |
| border-color: rgba(189, 195, 199, 0.2) !important; /* Light border */ | |
| } | |
| .gr-textbox, .gr-dropdown, .gr-button, .gr-code, .gr-chat-message, .gr-image { | |
| border-color: rgba(189, 195, 199, 0.3) !important; /* Lighter border for inputs */ | |
| background-color: rgba(52, 73, 94, 0.9) !important; /* Slightly lighter dark background for inputs */ | |
| color: #ecf0f1 !important; /* Light text color */ | |
| } | |
| /* Specific button colors using theme/fallback */ | |
| .gr-button { | |
| color: #c6c6fc !important; /* Default button text color */ | |
| } | |
| .gr-button.gr-button-primary { | |
| background-color: #1abc9c !important; /* Teal */ | |
| color: white !important; | |
| border-color: #16a085 !important; | |
| } | |
| .gr-button.gr-button-secondary { | |
| background-color: #9b59b6 !important; /* Purple */ | |
| color: white !important; | |
| border-color: #8e44ad !important; | |
| } | |
| .gr-button.gr-button-secondary:hover { | |
| background-color: #8e44ad !important; /* Darker purple on hover */ | |
| } | |
| .gr-button.gr-button-stop { | |
| background-color: #e74c3c !important; /* Red */ | |
| color: white !important; | |
| border-color: #c0392b !important; | |
| } | |
| .gr-markdown { | |
| background-color: rgba(44, 62, 80, 0.7) !important; /* Semi-transparent dark background for markdown */ | |
| padding: 10px; | |
| border-radius: 5px; | |
| } | |
| .gr-markdown h1, .gr-markdown h2, .gr-markdown h3, .gr-markdown h4, .gr-markdown h5, .gr-markdown h6 { | |
| color: #ecf0f1 !important; /* Light heading color */ | |
| border-bottom-color: rgba(189, 195, 199, 0.3) !important; /* Light border below headings */ | |
| } | |
| .gr-markdown p { /* Style for paragraph text in markdown */ | |
| color: #bdc3c7; /* Lighter gray for body text */ | |
| } | |
| .gr-markdown pre code { | |
| background-color: rgba(52, 73, 94, 0.95) !important; /* Dark code block background */ | |
| border-color: rgba(189, 195, 199, 0.3) !important; /* Light border for code */ | |
| } | |
| .gr-image div img { /* Style for image preview */ | |
| border: 1px solid #ccc; | |
| background-color: rgba(52, 73, 94, 0.9) !important; /* Dark background behind image */ | |
| } | |
| .gr-file div button { /* Style for file download button */ | |
| background-color: #1abc9c !important; /* Teal */ | |
| color: white !important; | |
| border: 1px solid #16a085 !important; | |
| } | |
| /* Style for the Key PEM textboxes to handle whitespace */ | |
| #key_pem_textbox textarea { | |
| white-space: pre !important; /* Preserve line breaks and spacing */ | |
| overflow-x: auto !important; /* Add scrollbar if lines are long */ | |
| } | |
| """ | |
| with gr.Blocks(theme=custom_theme, css=custom_css, title=f"KeyLock Steganography v{__version__}") as keylock_app_interface: | |
| gr.Markdown(f"<div align='center' style='margin-bottom:15px;'><span style='font-size:2.5em;font-weight:bold;'>🔑 KeyLock v{__version__}</span><h2 style='font-size:1.2em;color:#bdc3c7;margin-top:5px;'>Portable API Key Wallet in a PNG (RSA Encrypted)</h2></div>") | |
| gr.HTML("<div align='center' style='margin-bottom:10px;font-size:0.9em;color:#bdc3c7;'>Securely embed and extract API key-value pairs (or any text) within PNG images using LSB steganography and <span style='font-weight:bold;'>RSA Public-Key Encryption</span>, allowing recipient decryption without a shared password.</div>") | |
| gr.HTML("<div align='center' style='margin-bottom:15px;font-size:0.9em;'><span style='font-weight:bold;'>GitHub: <a href='https://github.com/broadfield-dev/KeyLock-API-Wallet' target='_blank' style='color:#1abc9c;'>KeyLock-API-Wallet</a> | Decoder Module: <a href='https://github.com/broadfield-dev/keylock-decode' target='_blank' style='color:#1abc9c;'>keylock-decode</a></span></div>") | |
| gr.HTML("<hr style='border-color: rgba(189, 195, 199, 0.2); margin-bottom:25px;'>") | |
| with gr.Tabs(): | |
| with gr.TabItem(f"{ICON_EMBED} Embed Data (Encrypt for Recipient)"): | |
| with gr.Row(): | |
| with gr.Column(scale=2): | |
| embed_kv_input = gr.Textbox( | |
| label="Secret Data (Key:Value Pairs, one per line)", | |
| placeholder="API_KEY_1: your_secret_value_1\nSERVICE_USER = 'user@example.com'\n# Lines starting with # are ignored", | |
| lines=7, | |
| info="Enter secrets as Key:Value or Key=Value. Each pair on a new line." | |
| ) | |
| embed_public_key_input = gr.Textbox( | |
| label="Recipient's Public Key (PEM Format)", | |
| placeholder="Paste the recipient's RSA public key here, or load from file...", | |
| lines=8, | |
| info="Data will be encrypted for the holder of the corresponding Private Key.", | |
| elem_id="key_pem_textbox" # Apply custom CSS | |
| ) | |
| # Optional: File upload for Public Key | |
| # embed_public_key_file = gr.File( | |
| # label="Load Public Key from .pem file", | |
| # type="filepath", # or "bytes" if reading in python | |
| # file_count="single", | |
| # # Add logic to read file content into the textbox | |
| # ) | |
| embed_output_filename_base = gr.Textbox( | |
| label="Base Name for Downloaded Stego Image", | |
| value="keylock_wallet_rsa", | |
| info="'.png' will be appended. e.g., 'project_secrets_for_john'" | |
| ) | |
| with gr.Accordion("Carrier Image Options", open=False): | |
| embed_generate_carrier_checkbox = gr.Checkbox( | |
| label="Generate new KeyLock Wallet image", | |
| value=True, | |
| info="Uncheck to upload your own PNG carrier image." | |
| ) | |
| embed_input_image_upload = gr.Image( | |
| label="Upload Your Own PNG Carrier (Optional)", | |
| type="pil", | |
| image_mode="RGB", # Will convert to RGB for LSB regardless | |
| sources=["upload","clipboard"], | |
| visible=False, # Initially hidden | |
| show_download_button=False, | |
| interactive=True | |
| ) | |
| embed_show_keys_checkbox = gr.Checkbox( | |
| label="Show list of key names on image overlay", | |
| value=True, | |
| info="Displays embedded key names (not values) on the image." | |
| ) | |
| embed_button = gr.Button("Embed Secrets "+ICON_EMBED, variant="primary") | |
| with gr.Column(scale=3): | |
| gr.Markdown("### Output Image & Status") | |
| embed_output_status = gr.Textbox( | |
| label="Embedding Status", | |
| lines=4, | |
| interactive=False, | |
| placeholder="Status messages will appear here..." | |
| ) | |
| embed_output_image_html = gr.HTML( | |
| label="Preview of Stego Image (Max 480px width)", | |
| value="<div style='text-align:center; color:#bdc3c7; padding:20px;'>Image preview will appear here after embedding.</div>" | |
| ) | |
| embed_download_file = gr.File( | |
| label="Download Your KeyLock Image (PNG)", | |
| interactive=False, | |
| file_count="single" | |
| ) | |
| # --- Interactivity --- | |
| def toggle_carrier_upload(generate_flag): | |
| return gr.update(visible=not generate_flag) | |
| embed_generate_carrier_checkbox.change( | |
| fn=toggle_carrier_upload, | |
| inputs=[embed_generate_carrier_checkbox], | |
| outputs=[embed_input_image_upload] | |
| ) | |
| embed_button.click( | |
| fn=gradio_embed_data, | |
| inputs=[ | |
| embed_kv_input, | |
| embed_public_key_input, | |
| embed_input_image_upload, | |
| embed_generate_carrier_checkbox, | |
| embed_show_keys_checkbox, | |
| embed_output_filename_base | |
| ], | |
| outputs=[ | |
| embed_output_image_html, | |
| embed_output_status, | |
| embed_download_file | |
| ] | |
| ) | |
| with gr.TabItem(f"{ICON_EXTRACT} Extract Data (Decrypt with Private Key)"): | |
| with gr.Row(): | |
| with gr.Column(scale=1): | |
| extract_stego_image_upload = gr.Image( | |
| label="Upload KeyLock PNG Image", | |
| type="pil", | |
| image_mode="RGB", # Will convert to RGB for LSB regardless | |
| sources=["upload","clipboard"], | |
| show_download_button=False, | |
| interactive=True, | |
| ) | |
| extract_private_key_input = gr.Textbox( | |
| label="Your Private Key (PEM Format)", | |
| placeholder="Paste your RSA private key here, or load from file...", | |
| lines=8, | |
| info="Required to decrypt the data.", | |
| elem_id="key_pem_textbox" # Apply custom CSS | |
| ) | |
| # Optional: File upload for Private Key | |
| # extract_private_key_file = gr.File( | |
| # label="Load Private Key from .pem file", | |
| # type="filepath", # or "bytes" if reading in python | |
| # file_count="single", | |
| # # Add logic to read file content into the textbox | |
| # ) | |
| extract_button = gr.Button("Extract Secrets "+ICON_EXTRACT, variant="primary") | |
| with gr.Column(scale=2): | |
| gr.Markdown("### Extracted Data & Status") | |
| extract_output_status = gr.Textbox( | |
| label="Extraction Status", | |
| lines=2, | |
| interactive=False, | |
| placeholder="Status messages will appear here..." | |
| ) | |
| extract_output_data = gr.Textbox( | |
| label="Extracted Secret Data", | |
| lines=10, | |
| interactive=False, | |
| placeholder="Extracted data (usually JSON) will appear here...", | |
| show_copy_button=True | |
| ) | |
| # --- Interactivity --- | |
| extract_button.click( | |
| fn=gradio_extract_data, | |
| inputs=[ | |
| extract_stego_image_upload, | |
| extract_private_key_input | |
| ], | |
| outputs=[ | |
| extract_output_data, | |
| extract_output_status | |
| ] | |
| ) | |
| with gr.TabItem(f"{ICON_KEY} Key Management (Generate Keys)"): | |
| with gr.Row(): | |
| with gr.Column(): | |
| gr.Markdown("### Generate RSA Key Pair") | |
| gr.Markdown( | |
| "Use this tab to generate a new pair of RSA Public and Private Keys.<br>" | |
| "🔑 **The Public Key** is used by others to encrypt data *for you*. You can share this freely.<br>" | |
| "🔐 **The Private Key** is used *only by you* to decrypt data sent to you. Keep this **absolutely secret and secure**! If you lose it, you cannot decrypt data encrypted with the corresponding Public Key.<br>" | |
| "Keys are generated in standard PEM format." | |
| ) | |
| key_size_dropdown = gr.Dropdown( | |
| label="Key Size (bits)", | |
| # Corrected the constant name here as well | |
| choices=[2048, 3072, 4096], | |
| value=core.RSA_KEY_SIZE_DEFAULT, | |
| info="Larger keys are more secure but result in larger encrypted data payloads (affecting image capacity)." | |
| ) | |
| generate_key_button = gr.Button("Generate New Key Pair "+ICON_KEY, variant="secondary") | |
| key_gen_status = gr.Textbox( | |
| label="Status", | |
| interactive=False, | |
| placeholder="Key generation status..." | |
| ) | |
| with gr.Column(): | |
| gr.Markdown("### Your Generated Keys") | |
| generated_private_key = gr.Textbox( | |
| label="Generated Private Key (Keep Secret!)", | |
| lines=10, | |
| interactive=False, | |
| show_copy_button=True, | |
| elem_id="key_pem_textbox" # Apply custom CSS | |
| ) | |
| generated_public_key = gr.Textbox( | |
| label="Generated Public Key (Share This)", | |
| lines=8, | |
| interactive=False, | |
| show_copy_button=True, | |
| elem_id="key_pem_textbox" # Apply custom CSS | |
| ) | |
| # Optional: Download buttons for keys | |
| # download_private_key_file = gr.File(label="Download Private Key", interactive=False) | |
| # download_public_key_file = gr.File(label="Download Public Key", interactive=False) | |
| # --- Interactivity --- | |
| generate_key_button.click( | |
| fn=gradio_generate_keys, | |
| inputs=[key_size_dropdown], | |
| outputs=[generated_private_key, generated_public_key, key_gen_status] | |
| # Add outputs for download files if implemented | |
| ) | |
| # Footer | |
| gr.Markdown("<hr style='border-color: rgba(189, 195, 199, 0.1); margin-top: 30px; margin-bottom:10px;'>") | |
| gr.Markdown(f"<div style='text-align:center; font-size:0.8em; color:#7f8c8d;'>KeyLock-API-Wallet v{__version__}. Use responsibly. <a href='https://github.com/broadfield-dev/KeyLock-API-Wallet/blob/main/README.md#important-notes--best-practices' target='_blank' style='color:#7f8c8d; text-decoration:underline;'>Read Important Notes & Best Practices</a>.</div>") | |
| return keylock_app_interface | |
| def main(): | |
| app_logger.info(f"Starting KeyLock Gradio Application v{__version__} (RSA Encryption Mode)...") | |
| try: | |
| # Attempt to load a common font to check PIL/Pillow font handling | |
| # These checks are more for diagnostic logging | |
| ImageFont.truetype("DejaVuSans.ttf", 10) # Common on Linux | |
| app_logger.info("DejaVuSans font found, PIL font rendering should be good.") | |
| except IOError: | |
| try: | |
| ImageFont.truetype("arial.ttf", 10) # Common on Windows | |
| app_logger.info("Arial font found, PIL font rendering should be good.") | |
| except IOError: | |
| app_logger.warning("Common system fonts (DejaVuSans/Arial) not found. PIL might use basic bitmap font if other preferred fonts in core.py are also unavailable. Overlay text quality might be affected.") | |
| keylock_app_interface = build_interface() | |
| # Prepare launch arguments | |
| launch_args = {} # Start with empty dict | |
| # Add allowed_paths for temp directory for file downloads | |
| temp_dir = tempfile.gettempdir() | |
| launch_args["allowed_paths"] = [temp_dir] | |
| app_logger.info(f"Allowing file downloads from temp directory: {temp_dir}") | |
| # Check environment variables for server name/port | |
| server_name = os.environ.get('GRADIO_SERVER_NAME') | |
| server_port = os.environ.get('GRADIO_SERVER_PORT') | |
| if server_name: | |
| launch_args["server_name"] = server_name | |
| app_logger.info(f"Using server_name from environment: {server_name}") | |
| if server_port: | |
| try: | |
| launch_args["server_port"] = int(server_port) | |
| app_logger.info(f"Using server_port from environment: {server_port}") | |
| except ValueError: | |
| app_logger.warning(f"Invalid GRADIO_SERVER_PORT: {server_port}. Using default.") | |
| # Check for --share argument (optional, for public sharing) | |
| # if "--share" in sys.argv: # Need to import sys if using this | |
| # launch_args["share"] = True | |
| # app_logger.info("Launching with share=True...") | |
| keylock_app_interface.launch(**launch_args) | |
| # Entry point for the Gradio app | |
| if __name__ == "__main__": | |
| # Add basic file logging for production if needed, alongside stream handler | |
| # file_handler = logging.FileHandler('keylock_app.log') | |
| # file_handler.setFormatter(formatter) # Use formatter defined earlier | |
| # app_logger.addHandler(file_handler) | |
| # logger.addHandler(file_handler) # Add to core logger too | |
| main() |