#!/usr/bin/env python3
"""
DecentraBrowser Key Management System

This script enhances DecentraBrowser with a secure key management system:
1. Creates a keychain folder structure for storing encryption keys
2. Implements public/private key infrastructure for secure key exchange
3. Adds server-to-server key transfer capabilities
4. Integrates with the existing file encryption system

Usage:
    python key_management_system.py
"""

import os
import sys
import shutil
import re
import json
import base64
import uuid
from pathlib import Path

def backup_file(file_path):
    """Create a backup of a file before modifying it."""
    backup_path = file_path + '.keychain.bak'
    print(f"Creating backup: {backup_path}")
    shutil.copy2(file_path, backup_path)
    return backup_path

def create_keychain_structure(base_dir):
    """Create the keychain directory structure."""
    # Data directory
    data_dir = os.path.join(base_dir, "data")
    if not os.path.exists(data_dir):
        data_dir = os.path.dirname(base_dir)  # Try parent directory
    
    # Create keychain directories
    keychain_dir = os.path.join(data_dir, "keychain")
    os.makedirs(keychain_dir, exist_ok=True)
    
    # Create subdirectories
    shared_keys_dir = os.path.join(keychain_dir, "shared")
    os.makedirs(shared_keys_dir, exist_ok=True)
    
    received_keys_dir = os.path.join(keychain_dir, "received")
    os.makedirs(received_keys_dir, exist_ok=True)
    
    trusted_servers_dir = os.path.join(keychain_dir, "trusted_servers")
    os.makedirs(trusted_servers_dir, exist_ok=True)
    
    print(f"Created keychain directory structure in {keychain_dir}")
    
    # Return paths for use in other functions
    return {
        "keychain": keychain_dir,
        "shared_keys": shared_keys_dir,
        "received_keys": received_keys_dir,
        "trusted_servers": trusted_servers_dir
    }

def create_keychain_manager(security_py, keychain_paths):
    """Create or update the KeychainManager class in the security.py file."""
    with open(security_py, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Check if KeychainManager already exists
    if "class KeychainManager" in content:
        print("KeychainManager already exists in security.py, updating...")
        
        # Extract existing class to preserve any custom modifications
        class_pattern = r'class KeychainManager.*?(?=\n\nclass|\Z)'
        match = re.search(class_pattern, content, re.DOTALL)
        if match:
            existing_class = match.group(0)
            # We'll modify this existing class
        else:
            print("Could not find existing KeychainManager class")
            existing_class = None
    else:
        existing_class = None
    
    # Create the new KeychainManager class
    keychain_manager_class = f"""
class KeychainManager:
    
    
    def __init__(self, security_manager, keychain_dir):
        
        self.security_manager = security_manager
        self.keychain_dir = keychain_dir
        self.shared_keys_dir = os.path.join(keychain_dir, "shared")
        self.received_keys_dir = os.path.join(keychain_dir, "received")
        self.trusted_servers_dir = os.path.join(keychain_dir, "trusted_servers")
        
        # Create directories if they don't exist
        os.makedirs(self.shared_keys_dir, exist_ok=True)
        os.makedirs(self.received_keys_dir, exist_ok=True)
        os.makedirs(self.trusted_servers_dir, exist_ok=True)
        
        # Load trusted servers
        self.trusted_servers = self._load_trusted_servers()
    
    def _load_trusted_servers(self):
        
        trusted_servers = {{}}
        
        for file_name in os.listdir(self.trusted_servers_dir):
            if file_name.endswith('.pub'):
                server_name = file_name[:-4]  # Remove .pub extension
                key_path = os.path.join(self.trusted_servers_dir, file_name)
                
                try:
                    with open(key_path, 'rb') as f:
                        public_key_data = f.read()
                    
                    # Load public key
                    public_key = serialization.load_pem_public_key(public_key_data)
                    trusted_servers[server_name] = {{
                        'public_key': public_key,
                        'public_key_pem': public_key_data
                    }}
                except Exception as e:
                    print(f"Error loading trusted server key {{file_name}}: {{str(e)}}")
        
        return trusted_servers
    
    def add_trusted_server(self, server_name, public_key_pem):
        
        try:
            # Validate the public key
            public_key = serialization.load_pem_public_key(public_key_pem)
            
            # Save the public key
            key_path = os.path.join(self.trusted_servers_dir, f"{{server_name}}.pub")
            with open(key_path, 'wb') as f:
                f.write(public_key_pem)
            
            # Add to trusted servers
            self.trusted_servers[server_name] = {{
                'public_key': public_key,
                'public_key_pem': public_key_pem
            }}
            
            return True, "Server added to trusted list"
        except Exception as e:
            return False, f"Error adding trusted server: {{str(e)}}"
    
    def remove_trusted_server(self, server_name):
        
        if server_name in self.trusted_servers:
            try:
                # Remove the public key file
                key_path = os.path.join(self.trusted_servers_dir, f"{{server_name}}.pub")
                if os.path.exists(key_path):
                    os.remove(key_path)
                
                # Remove from trusted servers
                del self.trusted_servers[server_name]
                
                return True, "Server removed from trusted list"
            except Exception as e:
                return False, f"Error removing trusted server: {{str(e)}}"
        else:
            return False, "Server not found in trusted list"
    
    def list_trusted_servers(self):
        
        return list(self.trusted_servers.keys())
    
    def store_key_for_file(self, file_path, aes_key, recipients=None):
        
        try:
            # Generate a unique key ID
            key_id = str(uuid.uuid4())
            
            # Create key metadata
            metadata = {{
                'key_id': key_id,
                'file_path': os.path.basename(file_path),
                'created': datetime.datetime.now().isoformat(),
                'recipients': []
            }}
            
            # Store the key in the shared keys directory
            key_path = os.path.join(self.shared_keys_dir, f"{{key_id}}.key")
            with open(key_path, 'wb') as f:
                f.write(aes_key)
            
            # Store metadata
            metadata_path = os.path.join(self.shared_keys_dir, f"{{key_id}}.meta")
            with open(metadata_path, 'w') as f:
                json.dump(metadata, f, indent=2)
            
            # If recipients specified, encrypt the key for each recipient
            if recipients:
                for recipient in recipients:
                    if recipient in self.trusted_servers:
                        server = self.trusted_servers[recipient]
                        public_key = server['public_key']
                        
                        # Encrypt the AES key with the recipient's public key
                        encrypted_key = self.security_manager.encrypt_with_rsa(public_key, aes_key)
                        
                        # Store the encrypted key
                        recipient_key_path = os.path.join(
                            self.shared_keys_dir, 
                            f"{{key_id}}.{{recipient}}.enc"
                        )
                        with open(recipient_key_path, 'w') as f:
                            f.write(encrypted_key)
                        
                        # Update metadata
                        metadata['recipients'].append(recipient)
                
                # Update metadata file with recipients
                with open(metadata_path, 'w') as f:
                    json.dump(metadata, f, indent=2)
            
            return True, key_id
        except Exception as e:
            return False, f"Error storing key: {{str(e)}}"
    
    def get_key_for_file(self, file_path):
        
        try:
            file_name = os.path.basename(file_path)
            
            # Search in shared keys
            for meta_file in os.listdir(self.shared_keys_dir):
                if meta_file.endswith('.meta'):
                    meta_path = os.path.join(self.shared_keys_dir, meta_file)
                    with open(meta_path, 'r') as f:
                        metadata = json.load(f)
                    
                    if metadata['file_path'] == file_name:
                        key_id = metadata['key_id']
                        key_path = os.path.join(self.shared_keys_dir, f"{{key_id}}.key")
                        
                        if os.path.exists(key_path):
                            with open(key_path, 'rb') as f:
                                return True, f.read()
            
            # Search in received keys
            for meta_file in os.listdir(self.received_keys_dir):
                if meta_file.endswith('.meta'):
                    meta_path = os.path.join(self.received_keys_dir, meta_file)
                    with open(meta_path, 'r') as f:
                        metadata = json.load(f)
                    
                    if metadata['file_path'] == file_name:
                        key_id = metadata['key_id']
                        key_path = os.path.join(self.received_keys_dir, f"{{key_id}}.key")
                        
                        if os.path.exists(key_path):
                            with open(key_path, 'rb') as f:
                                return True, f.read()
            
            return False, "Key not found for file"
        except Exception as e:
            return False, f"Error retrieving key: {{str(e)}}"
    
    def receive_encrypted_key(self, encrypted_key_data, sender, file_path):
        
        try:
            # Generate a key ID
            key_id = str(uuid.uuid4())
            
            # Load our private key
            private_key = self.security_manager.load_private_key('private_key.pem')
            
            # Decode if base64 encoded
            if isinstance(encrypted_key_data, str):
                encrypted_key_data = base64.b64decode(encrypted_key_data)
            
            # Decrypt the key
            decrypted_key = self.security_manager.decrypt_with_rsa(private_key, encrypted_key_data)
            
            # Store the decrypted key
            key_path = os.path.join(self.received_keys_dir, f"{{key_id}}.key")
            with open(key_path, 'wb') as f:
                f.write(decrypted_key)
            
            # Create metadata
            metadata = {{
                'key_id': key_id,
                'file_path': os.path.basename(file_path),
                'received_from': sender,
                'received': datetime.datetime.now().isoformat()
            }}
            
            # Store metadata
            metadata_path = os.path.join(self.received_keys_dir, f"{{key_id}}.meta")
            with open(metadata_path, 'w') as f:
                json.dump(metadata, f, indent=2)
            
            return True, key_id
        except Exception as e:
            return False, f"Error receiving encrypted key: {{str(e)}}"
    
    def export_encrypted_key(self, key_id, recipient):
        
        try:
            # Check if recipient is trusted
            if recipient not in self.trusted_servers:
                return False, "Recipient not in trusted servers list"
            
            # Find the key
            key_path = os.path.join(self.shared_keys_dir, f"{{key_id}}.key")
            if not os.path.exists(key_path):
                return False, "Key not found"
            
            # Check if already encrypted for this recipient
            recipient_key_path = os.path.join(self.shared_keys_dir, f"{{key_id}}.{{recipient}}.enc")
            if os.path.exists(recipient_key_path):
                # Return existing encrypted key
                with open(recipient_key_path, 'r') as f:
                    return True, f.read()
            
            # Load the key
            with open(key_path, 'rb') as f:
                aes_key = f.read()
            
            # Get recipient's public key
            public_key = self.trusted_servers[recipient]['public_key']
            
            # Encrypt the key
            encrypted_key = self.security_manager.encrypt_with_rsa(public_key, aes_key)
            
            # Store the encrypted key
            with open(recipient_key_path, 'w') as f:
                f.write(encrypted_key)
            
            # Update metadata
            meta_path = os.path.join(self.shared_keys_dir, f"{{key_id}}.meta")
            if os.path.exists(meta_path):
                with open(meta_path, 'r') as f:
                    metadata = json.load(f)
                
                if recipient not in metadata['recipients']:
                    metadata['recipients'].append(recipient)
                
                with open(meta_path, 'w') as f:
                    json.dump(metadata, f, indent=2)
            
            return True, encrypted_key
        except Exception as e:
            return False, f"Error exporting encrypted key: {{str(e)}}"
    
    def delete_key(self, key_id, received=False):
        
        try:
            # Determine the directory
            key_dir = self.received_keys_dir if received else self.shared_keys_dir
            
            # Delete key files
            key_path = os.path.join(key_dir, f"{{key_id}}.key")
            if os.path.exists(key_path):
                os.remove(key_path)
            
            meta_path = os.path.join(key_dir, f"{{key_id}}.meta")
            if os.path.exists(meta_path):
                os.remove(meta_path)
            
            # Delete encrypted keys for recipients
            if not received:
                for file_name in os.listdir(key_dir):
                    if file_name.startswith(f"{{key_id}}.") and file_name.endswith(".enc"):
                        os.remove(os.path.join(key_dir, file_name))
            
            return True, "Key deleted successfully"
        except Exception as e:
            return False, f"Error deleting key: {{str(e)}}"
    
    def list_keys(self, received=False):
        
        keys = []
        key_dir = self.received_keys_dir if received else self.shared_keys_dir
        
        for file_name in os.listdir(key_dir):
            if file_name.endswith('.meta'):
                meta_path = os.path.join(key_dir, file_name)
                try:
                    with open(meta_path, 'r') as f:
                        metadata = json.load(f)
                    keys.append(metadata)
                except Exception as e:
                    print(f"Error loading metadata {{file_name}}: {{str(e)}}")
        
        return keys
"""
    
    if existing_class:
        # Replace existing class with new one
        updated_content = content.replace(existing_class, keychain_manager_class)
    else:
        # Add class at the end of the file
        updated_content = content + "\n" + keychain_manager_class
    
    # Add any missing imports
    if "import uuid" not in updated_content:
        imports_end = updated_content.find("\n\n", updated_content.find("import"))
        if imports_end > 0:
            updated_imports = updated_content[:imports_end] + "\nimport uuid\nimport json\nimport datetime\n" + updated_content[imports_end:]
            updated_content = updated_imports
    
    # Write updated content
    with open(security_py, 'w', encoding='utf-8') as f:
        f.write(updated_content)
    
    print(f"Created/updated KeychainManager in {security_py}")
    return True

def update_security_manager(security_py, keychain_paths):
    """Update the SecurityManager to include the KeychainManager."""
    with open(security_py, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Check if keychain_manager is already in SecurityManager
    if "self.keychain_manager = KeychainManager" in content:
        print("SecurityManager already includes KeychainManager, skipping...")
        return True
    
    # Find the SecurityManager class
    class_pattern = r'class SecurityManager:.*?def __init__\(self, certs_dir\):.*?(?=\n\s+def)'
    match = re.search(class_pattern, content, re.DOTALL)
    
    if not match:
        print("Could not find SecurityManager.__init__ method")
        return False
    
    init_method = match.group(0)
    
    # Get the indentation level
    init_lines = init_method.split('\n')
    if len(init_lines) > 1:
        indentation = re.match(r'(\s*)', init_lines[1]).group(1)
    else:
        indentation = '        '  # Default indentation
    
    # Find where to insert keychain_manager initialization
    # Typically after self.certs_dir = certs_dir
    keychain_init = f"\n{indentation}# Initialize keychain manager\n"
    keychain_init += f"{indentation}keychain_dir = os.path.join(os.path.dirname(certs_dir), 'keychain')\n"
    keychain_init += f"{indentation}os.makedirs(keychain_dir, exist_ok=True)\n"
    keychain_init += f"{indentation}self.keychain_manager = KeychainManager(self, keychain_dir)\n"
    
    # Insert after certs_dir assignment
    certs_dir_pattern = r'(self\.certs_dir = certs_dir\s*)'
    if re.search(certs_dir_pattern, init_method):
        updated_init = re.sub(certs_dir_pattern, f"\\1{keychain_init}", init_method)
        updated_content = content.replace(init_method, updated_init)
    else:
        # If pattern not found, add at the end of init
        last_line_pattern = r'(__init__.*?\n)(\s+)([^\n]+\n)'
        match = re.search(last_line_pattern, init_method)
        if match:
            indent = match.group(2)
            updated_init = init_method + keychain_init
            updated_content = content.replace(init_method, updated_init)
        else:
            print("Could not find a place to insert keychain_manager initialization")
            return False
    
    # Write updated content
    with open(security_py, 'w', encoding='utf-8') as f:
        f.write(updated_content)
    
    print(f"Updated SecurityManager in {security_py}")
    return True

def update_file_creation(main_window_py, keychain_paths):
    """Update the file creation process to use the keychain system."""
    with open(main_window_py, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Find the create_new_file method
    method_pattern = r'def create_new_file\(self\):.*?(?=\n    def|\Z)'
    match = re.search(method_pattern, content, re.DOTALL)
    
    if not match:
        print("Could not find create_new_file method")
        return False
    
    create_file_method = match.group(0)
    
    # Check if the method already uses keychain
    if "self.security_manager.keychain_manager.store_key_for_file" in create_file_method:
        print("create_new_file already uses keychain system, skipping...")
        return True
    
    # Find the save_file function within create_new_file
    save_file_pattern = r'def save_file\(\):.*?(?=\n        dialog\.exec_\(\))'
    save_file_match = re.search(save_file_pattern, create_file_method, re.DOTALL)
    
    if not save_file_match:
        print("Could not find save_file function in create_new_file method")
        return False
    
    save_file_function = save_file_match.group(0)
    
    # Update the save_file function to use keychain
    if "# Save the key alongside the file" in save_file_function:
        # Replace the key storage part
        old_key_storage = r'# Save the key alongside the file.*?key_path = file_path \+ \'\.key\'.*?with open\(key_path, \'wb\'\) as f:.*?f\.write\(aes_key\)'
        new_key_storage = """
                # Store key in keychain
                success, key_id = self.security_manager.keychain_manager.store_key_for_file(file_path, aes_key)
                if not success:
                    logger.error(f"Error storing key in keychain: {key_id}")
"""
        updated_save_file = re.sub(old_key_storage, new_key_storage, save_file_function, flags=re.DOTALL)
        
        # Update the create_file_method with the new save_file function
        updated_create_file = create_file_method.replace(save_file_function, updated_save_file)
        
        # Update the content with the updated method
        updated_content = content.replace(create_file_method, updated_create_file)
        
        # Write updated content
        with open(main_window_py, 'w', encoding='utf-8') as f:
            f.write(updated_content)
        
        print(f"Updated file creation process in {main_window_py}")
        return True
    else:
        print("Could not find key storage code in save_file function")
        return False

def update_file_download(client_py, keychain_paths):
    """Update the file download process to use the keychain system."""
    with open(client_py, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Find the _download_file_thread method
    method_pattern = r'def _download_file_thread\(self, download_id, url, session, save_dir, file_path\):.*?(?=\n    def|\Z)'
    match = re.search(method_pattern, content, re.DOTALL)
    
    if not match:
        print("Could not find _download_file_thread method")
        return False
    
    download_method = match.group(0)
    
    # Check if the method already uses keychain
    if "self.security_manager.keychain_manager.receive_encrypted_key" in download_method:
        print("_download_file_thread already uses keychain system, skipping...")
        return True
    
    # Update the method to use keychain
    if "# Check if encryption key is provided" in download_method:
        # Replace the key handling part
        old_key_handling = r'# Check if encryption key is provided.*?encryption_key = None.*?if \'X-Encryption-Key\' in response\.headers:.*?encryption_key = base64\.b64decode\(response\.headers\[\'X-Encryption-Key\'\]\)'
        new_key_handling = """
            # Check if encryption key is provided
            encryption_key = None
            if 'X-Encryption-Key' in response.headers:
                try:
                    encrypted_key = base64.b64decode(response.headers['X-Encryption-Key'])
                    
                    # Store the key in keychain
                    success, key_id_or_error = self.security_manager.keychain_manager.receive_encrypted_key(
                        encrypted_key, 
                        domain_name, 
                        file_path
                    )
                    
                    if success:
                        # Retrieve the decrypted key
                        success, encryption_key = self.security_manager.keychain_manager.get_key_for_file(file_path)
                        if not success:
                            logger.error(f"Error retrieving key from keychain: {encryption_key}")
                    else:
                        logger.error(f"Error storing encrypted key in keychain: {key_id_or_error}")
                except Exception as e:
                    logger.error(f"Error processing encryption key: {str(e)}")"""
        
        updated_download_method = re.sub(old_key_handling, new_key_handling, download_method, flags=re.DOTALL)
        
        # Update the content with the updated method
        updated_content = content.replace(download_method, updated_download_method)
        
        # Write updated content
        with open(client_py, 'w', encoding='utf-8') as f:
            f.write(updated_content)
        
        print(f"Updated file download process in {client_py}")
        return True
    else:
        print("Could not find key handling code in _download_file_thread method")
        return False

def add_keychain_ui(main_window_py, keychain_paths):
    """Add a keychain management UI to the main window."""
    with open(main_window_py, 'r', encoding='utf-8') as f:
        content = f.read()
    
    # Check if the keychain tab already exists
    if "self.tab_keychain" in content:
        print("Keychain UI already exists in main_window.py, skipping...")
        return True
    
    # Find the tab initialization code
    tab_init_pattern = r'(self\.tabs\.addTab\(self\.tab_decryption, "Decryption Arena"\)\s*)'
    match = re.search(tab_init_pattern, content)
    
    if not match:
        print("Could not find tab initialization code")
        return False
    
    tab_init_code = match.group(1)
    new_tab_code = tab_init_code + '        self.tab_keychain = QWidget()\n        self.tabs.addTab(self.tab_keychain, "Keychain")\n'
    updated_content = content.replace(tab_init_code, new_tab_code)
    
    # Find the setup_decryption_tab method to add our method after it
    decryption_tab_pattern = r'(def setup_decryption_tab\(self\):.*?)(?=\n    def|\Z)'
    match = re.search(decryption_tab_pattern, updated_content, re.DOTALL)
    
    if not match:
        print("Could not find setup_decryption_tab method")
        return False
    
    decryption_tab_method = match.group(1)
    
    # Add keychain tab setup method
    keychain_tab_method = """
    def setup_keychain_tab(self):
       
        layout = QVBoxLayout()
        
        # Description
        description = QLabel("Keychain Manager - Manage encryption keys and trusted servers")
        description.setStyleSheet("font-size: 14px; font-weight: bold; margin: 10px;")
        layout.addWidget(description)
        
        # Tabs for different sections
        keychain_tabs = QTabWidget()
        
        # Create tabs
        tab_shared_keys = QWidget()
        tab_received_keys = QWidget()
        tab_trusted_servers = QWidget()
        
        keychain_tabs.addTab(tab_shared_keys, "Shared Keys")
        keychain_tabs.addTab(tab_received_keys, "Received Keys")
        keychain_tabs.addTab(tab_trusted_servers, "Trusted Servers")
        
        # Setup shared keys tab
        shared_layout = QVBoxLayout()
        
        self.tbl_shared_keys = QTableWidget()
        self.tbl_shared_keys.setColumnCount(4)
        self.tbl_shared_keys.setHorizontalHeaderLabels([
            "Key ID", "File", "Created", "Actions"
        ])
        self.tbl_shared_keys.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        
        shared_buttons = QHBoxLayout()
        btn_refresh_shared = QPushButton("Refresh")
        btn_refresh_shared.clicked.connect(self.refresh_shared_keys)
        
        btn_export_key = QPushButton("Export Key")
        btn_export_key.clicked.connect(self.export_shared_key)
        
        btn_delete_shared = QPushButton("Delete Key")
        btn_delete_shared.clicked.connect(self.delete_shared_key)
        
        shared_buttons.addWidget(btn_refresh_shared)
        shared_buttons.addWidget(btn_export_key)
        shared_buttons.addWidget(btn_delete_shared)
        
        shared_layout.addWidget(self.tbl_shared_keys)
        shared_layout.addLayout(shared_buttons)
        tab_shared_keys.setLayout(shared_layout)
        
        # Setup received keys tab
        received_layout = QVBoxLayout()
        
        self.tbl_received_keys = QTableWidget()
        self.tbl_received_keys.setColumnCount(4)
        self.tbl_received_keys.setHorizontalHeaderLabels([
            "Key ID", "File", "From", "Actions"
        ])
        self.tbl_received_keys.horizontalHeader().setSectionResizeMode(QHeaderView.Stretch)
        
        received_buttons = QHBoxLayout()
        btn_refresh_received = QPushButton("Refresh")
        btn_refresh_received.clicked.connect(self.refresh_received_keys)
        
        btn_delete_received = QPushButton("Delete Key")
        btn_delete_received.clicked.connect(self.delete_received_key)
        
        received_buttons.addWidget(btn_refresh_received)
        received_buttons.addWidget(btn_delete_received)
        
        received_layout.addWidget(self.tbl_received_keys)
        received_layout.addLayout(received_buttons)
        tab_received_keys.setLayout(received_layout)
        
        # Setup trusted servers tab
        trusted_layout = QVBoxLayout()
        
        self.lst_trusted_servers = QListWidget()
        
        trusted_buttons = QHBoxLayout()
        btn_refresh_servers = QPushButton("Refresh")
        btn_refresh_servers.clicked.connect(self.refresh_trusted_servers)
        
        btn_add_server = QPushButton("Add Server")
        btn_add_server.clicked.connect(self.add_trusted_server)
        
        btn_remove_server = QPushButton("Remove Server")
        btn_remove_server.clicked.connect(self.remove_trusted_server)
        
        trusted_buttons.addWidget(btn_refresh_servers)
        trusted_buttons.addWidget(btn_add_server)
        trusted_buttons.addWidget(btn_remove_server)
        
        trusted_layout.addWidget(self.lst_trusted_servers)
        trusted_layout.addLayout(trusted_buttons)
        tab_trusted_servers.setLayout(trusted_layout)
        
        # Add tabs to layout
        layout.addWidget(keychain_tabs)
        
        # Initial data load
        self.refresh_shared_keys()
        self.refresh_received_keys()
        self.refresh_trusted_servers()
        
        self.tab_keychain.setLayout(layout)
    
    def refresh_shared_keys(self):
        
        try:
            # Get shared keys
            shared_keys = self.security_manager.keychain_manager.list_keys(received=False)
            
            # Clear table
            self.tbl_shared_keys.setRowCount(0)
            
            # Fill table
            for i, key_data in enumerate(shared_keys):
                self.tbl_shared_keys.insertRow(i)
                
                key_id = key_data.get('key_id', '')
                file_path = key_data.get('file_path', '')
                created = key_data.get('created', '')
                
                # Format date if it's ISO format
                if created and 'T' in created:
                    try:
                        dt = datetime.datetime.fromisoformat(created)
                        created = dt.strftime('%Y-%m-%d %H:%M:%S')
                    except:
                        pass
                
                self.tbl_shared_keys.setItem(i, 0, QTableWidgetItem(key_id))
                self.tbl_shared_keys.setItem(i, 1, QTableWidgetItem(file_path))
                self.tbl_shared_keys.setItem(i, 2, QTableWidgetItem(created))
                
                # Actions cell
                actions_widget = QWidget()
                actions_layout = QHBoxLayout(actions_widget)
                actions_layout.setContentsMargins(2, 2, 2, 2)
                
                btn_view = QPushButton("View")
                btn_view.setProperty("key_id", key_id)
                btn_view.clicked.connect(self.view_shared_key_details)
                
                actions_layout.addWidget(btn_view)
                
                self.tbl_shared_keys.setCellWidget(i, 3, actions_widget)
        except Exception as e:
            self.statusBar.showMessage(f"Error refreshing shared keys: {str(e)}", 5000)
    
    def refresh_received_keys(self):
        
        try:
            # Get received keys
            received_keys = self.security_manager.keychain_manager.list_keys(received=True)
            
            # Clear table
            self.tbl_received_keys.setRowCount(0)
            
            # Fill table
            for i, key_data in enumerate(received_keys):
                self.tbl_received_keys.insertRow(i)
                
                key_id = key_data.get('key_id', '')
                file_path = key_data.get('file_path', '')
                sender = key_data.get('received_from', '')
                
                self.tbl_received_keys.setItem(i, 0, QTableWidgetItem(key_id))
                self.tbl_received_keys.setItem(i, 1, QTableWidgetItem(file_path))
                self.tbl_received_keys.setItem(i, 2, QTableWidgetItem(sender))
                
                # Actions cell
                actions_widget = QWidget()
                actions_layout = QHBoxLayout(actions_widget)
                actions_layout.setContentsMargins(2, 2, 2, 2)
                
                btn_view = QPushButton("View")
                btn_view.setProperty("key_id", key_id)
                btn_view.clicked.connect(self.view_received_key_details)
                
                actions_layout.addWidget(btn_view)
                
                self.tbl_received_keys.setCellWidget(i, 3, actions_widget)
        except Exception as e:
            self.statusBar.showMessage(f"Error refreshing received keys: {str(e)}", 5000)
    
    def refresh_trusted_servers(self):
        
        try:
            # Get trusted servers
            trusted_servers = self.security_manager.keychain_manager.list_trusted_servers()
            
            # Clear list
            self.lst_trusted_servers.clear()
            
            # Fill list
            for server in trusted_servers:
                self.lst_trusted_servers.addItem(server)
        except Exception as e:
            self.statusBar.showMessage(f"Error refreshing trusted servers: {str(e)}", 5000)
    
    def view_shared_key_details(self):
        
        sender = self.sender()
        key_id = sender.property("key_id")
        
        # Get shared keys
        shared_keys = self.security_manager.keychain_manager.list_keys(received=False)
        
        # Find the key
        key_data = None
        for key in shared_keys:
            if key.get('key_id') == key_id:
                key_data = key
                break
        
        if not key_data:
            self.statusBar.showMessage(f"Key {key_id} not found", 5000)
            return
        
        # Display details
        details = f"Key ID: {key_id}\\n"
        details += f"File: {key_data.get('file_path', '')}\\n"
        details += f"Created: {key_data.get('created', '')}\\n"
        details += f"Recipients: {', '.join(key_data.get('recipients', []))}\\n"
        
        QMessageBox.information(self, "Key Details", details)
    
    def view_received_key_details(self):
        
        sender = self.sender()
        key_id = sender.property("key_id")
        
        # Get received keys
        received_keys = self.security_manager.keychain_manager.list_keys(received=True)
        
        # Find the key
        key_data = None
        for key in received_keys:
            if key.get('key_id') == key_id:
                key_data = key
                break
        
        if not key_data:
            self.statusBar.showMessage(f"Key {key_id} not found", 5000)
            return
        
        # Display details
        details = f"Key ID: {key_id}\\n"
        details += f"File: {key_data.get('file_path', '')}\\n"
        details += f"Received from: {key_data.get('received_from', '')}\\n"
        details += f"Received: {key_data.get('received', '')}\\n"
        
        QMessageBox.information(self, "Key Details", details)
    
    def export_shared_key(self):
        
        selected_items = self.tbl_shared_keys.selectedItems()
        if not selected_items:
            self.statusBar.showMessage("No key selected", 5000)
            return
        
        row = selected_items[0].row()
        key_id = self.tbl_shared_keys.item(row, 0).text()
        
        # Get list of trusted servers
        trusted_servers = self.security_manager.keychain_manager.list_trusted_servers()
        if not trusted_servers:
            QMessageBox.warning(self, "Export Key", "No trusted servers available. Add a server first.")
            return
        
        # Ask user to select a server
        server, ok = QInputDialog.getItem(
            self, "Select Recipient", "Export key to server:", trusted_servers, 0, False
        )
        
        if not ok or not server:
            return
        
        # Export the key
        success, result = self.security_manager.keychain_manager.export_encrypted_key(key_id, server)
        
        if success:
            # Save to file
            file_path, _ = QFileDialog.getSaveFileName(
                self, "Save Encrypted Key", "", "Encrypted Keys (*.ekey);;All Files (*.*)"
            )
            
            if file_path:
                with open(file_path, 'w') as f:
                    f.write(result)
                
                self.statusBar.showMessage(f"Key exported to {file_path}", 5000)
        else:
            QMessageBox.warning(self, "Export Error", result)
    
    def delete_shared_key(self):
        
        selected_items = self.tbl_shared_keys.selectedItems()
        if not selected_items:
            self.statusBar.showMessage("No key selected", 5000)
            return
        
        row = selected_items[0].row()
        key_id = self.tbl_shared_keys.item(row, 0).text()
        
        # Confirm deletion
        reply = QMessageBox.question(
            self, "Confirm Deletion",
            f"Are you sure you want to delete key {key_id}?",
            QMessageBox.Yes | QMessageBox.No, QMessageBox.No
        )
        
        if reply != QMessageBox.Yes:
            return
        
        # Delete the key
        success, message = self.security_manager.keychain_manager.delete_key(key_id, received=False)
        
        if success:
            self.statusBar.showMessage("Key deleted successfully", 5000)
            self.refresh_shared_keys()
        else:
            QMessageBox.warning(self, "Delete Error", message)
    
    def delete_received_key(self):
        
        selected_items = self.tbl_received_keys.selectedItems()
        if not selected_items:
            self.statusBar.showMessage("No key selected", 5000)
            return
        
        row = selected_items[0].row()
        key_id = self.tbl_received_keys.item(row, 0).text()
        
        # Confirm deletion
        reply = QMessageBox.question(
            self, "Confirm Deletion",
            f"Are you sure you want to delete key {key_id}?",
            QMessageBox.Yes | QMessageBox.No, QMessageBox.No
        )
        
        if reply != QMessageBox.Yes:
            return
        
        # Delete the key
        success, message = self.security_manager.keychain_manager.delete_key(key_id, received=True)
        
        if success:
            self.statusBar.showMessage("Key deleted successfully", 5000)
            self.refresh_received_keys()
        else:
            QMessageBox.warning(self, "Delete Error", message)
    
    def add_trusted_server(self):
        
        # Get server name
        name, ok = QInputDialog.getText(self, "Add Trusted Server", "Server name:")
        if not ok or not name:
            return
        
        # Get public key file
        file_path, _ = QFileDialog.getOpenFileName(
            self, "Select Public Key", "", "Public Keys (*.pub);;All Files (*.*)"
        )
        
        if not file_path:
            return
        
        try:
            # Read public key
            with open(file_path, 'rb') as f:
                public_key_pem = f.read()
            
            # Add to trusted servers
            success, message = self.security_manager.keychain_manager.add_trusted_server(name, public_key_pem)
            
            if success:
                self.statusBar.showMessage(f"Added {name} to trusted servers", 5000)
                self.refresh_trusted_servers()
            else:
                QMessageBox.warning(self, "Add Server Error", message)
        except Exception as e:
            QMessageBox.critical(self, "Error", f"Failed to add server: {str(e)}")
    
    def remove_trusted_server(self):
        
        selected_items = self.lst_trusted_servers.selectedItems()
        if not selected_items:
            self.statusBar.showMessage("No server selected", 5000)
            return
        
        server_name = selected_items[0].text()
        
        # Confirm removal
        reply = QMessageBox.question(
            self, "Confirm Removal",
            f"Are you sure you want to remove {server_name} from trusted servers?",
            QMessageBox.Yes | QMessageBox.No, QMessageBox.No
        )
        
        if reply != QMessageBox.Yes:
            return
        
        # Remove the server
        success, message = self.security_manager.keychain_manager.remove_trusted_server(server_name)
        
        if success:
            self.statusBar.showMessage(f"Removed {server_name} from trusted servers", 5000)
            self.refresh_trusted_servers()
        else:
            QMessageBox.warning(self, "Remove Server Error", message)
"""
    
    # Add the keychain tab method after decryption tab method
    updated_content = updated_content.replace(decryption_tab_method, decryption_tab_method + keychain_tab_method)
    
    # Add the keychain tab setup call in init_ui
    setup_pattern = r'(# Setup decryption tab\s*self\.setup_decryption_tab\(\)\s*)'
    match = re.search(setup_pattern, updated_content)
    
    if match:
        setup_code = match.group(1)
        new_setup = setup_code + '\n        # Setup keychain tab\n        self.setup_keychain_tab()\n'
        updated_content = updated_content.replace(setup_code, new_setup)
    else:
        print("Could not find setup_decryption_tab call in init_ui")
        return False
    
    # Write updated content
    with open(main_window_py, 'w', encoding='utf-8') as f:
        f.write(updated_content)
    
    print(f"Added keychain management UI to {main_window_py}")
    return True

def main():
    """Main function to set up the key management system."""
    # Find DecentraBrowser installation directory
    base_dir = None
    possible_paths = [
        os.path.join(".", "decentrabrowser"),
        os.path.join("..", "decentrabrowser"),
        os.path.join(os.path.dirname(os.path.abspath(__file__)), "decentrabrowser")
    ]
    
    for path in possible_paths:
        if os.path.isdir(path):
            base_dir = path
            break
    
    if not base_dir:
        print("Error: Could not find DecentraBrowser directory.")
        print("Please run this script from the same directory as the 'decentrabrowser' folder.")
        return 1
    
    print(f"Found DecentraBrowser at: {base_dir}")
    
    # Target files
    security_py = os.path.join(base_dir, "shared", "security.py")
    main_window_py = os.path.join(base_dir, "client", "gui", "main_window.py")
    client_py = os.path.join(base_dir, "client", "network", "client.py")
    
    if not all(os.path.exists(f) for f in [security_py, main_window_py, client_py]):
        print("Error: Required files not found.")
        return 1
    
    # Create backups
    backup_file(security_py)
    backup_file(main_window_py)
    backup_file(client_py)
    
    # Create keychain structure
    keychain_paths = create_keychain_structure(base_dir)
    
    # Create/update KeychainManager in security.py
    create_keychain_manager(security_py, keychain_paths)
    
    # Update SecurityManager to include KeychainManager
    update_security_manager(security_py, keychain_paths)
    
    # Update file creation process
    update_file_creation(main_window_py, keychain_paths)
    
    # Update file download process
    update_file_download(client_py, keychain_paths)
    
    # Add keychain UI
    add_keychain_ui(main_window_py, keychain_paths)
    
    print("\nKey Management System installed successfully!")
    print("\nThe system includes:")
    print("1. Secure keychain storage for encryption keys")
    print("2. Public/private key infrastructure for key exchange")
    print("3. Trusted server management")
    print("4. UI for managing keys and trusted servers")
    print("\nPlease restart DecentraBrowser to use the new features.")
    
    return 0

if __name__ == "__main__":
    sys.exit(main())
