# keyCocoon/keyBunker - Credential Management
# Copyright (c) 2025 Gradient Research LLC
# All rights reserved.

import keyring
import os
import yaml
import socket
from PyQt5.QtWidgets import (QApplication, QMainWindow, QWidget, QVBoxLayout, 
                             QHBoxLayout, QFormLayout, QComboBox, QLineEdit, 
                             QPushButton, QLabel, QMessageBox, QDialog, 
                             QDialogButtonBox, QCheckBox, QScrollArea)
from PyQt5.QtCore import Qt
from PyQt5.QtGui import QIcon
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes
from cryptography.fernet import Fernet
import base64
import sys
import requests
import threading
import pyotp
from datetime import datetime

class CustomFieldDialog(QDialog):
    def __init__(self, parent=None):
        super().__init__(parent)
        self.setWindowTitle("Add Custom Field")
        self.setModal(True)
        
        layout = QVBoxLayout()
        
        form = QFormLayout()
        self.field_name = QLineEdit()
        self.store_in_keyring = QCheckBox()
        
        form.addRow("Field Name:", self.field_name)
        form.addRow("Store in Keyring:", self.store_in_keyring)
        
        layout.addLayout(form)
        
        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)
        
        self.setLayout(layout)
    
    def get_field_data(self):
        return {
            'name': self.field_name.text().strip(),
            'store_in_keyring': self.store_in_keyring.isChecked()
        }
    
class PasswordDialog(QDialog):
    def __init__(self, prompt_text, parent=None):
        super().__init__(parent)
        self.setWindowTitle("KeyBunker Authentication")
        self.setModal(True)
        self.setWindowFlags(self.windowFlags() | Qt.WindowStaysOnTopHint)
        
        layout = QVBoxLayout()
        
        label = QLabel(prompt_text)
        layout.addWidget(label)
        
        self.password_input = QLineEdit()
        self.password_input.setEchoMode(QLineEdit.Password)
        layout.addWidget(self.password_input)
        
        buttons = QDialogButtonBox(QDialogButtonBox.Ok | QDialogButtonBox.Cancel)
        buttons.accepted.connect(self.accept)
        buttons.rejected.connect(self.reject)
        layout.addWidget(buttons)
        
        self.setLayout(layout)
    
    def showEvent(self, event):
        super().showEvent(event)
        from PyQt5.QtCore import QTimer
        QTimer.singleShot(0, self.set_focus)
        QTimer.singleShot(100, self.set_focus)
        QTimer.singleShot(200, self.set_focus)
    
    def set_focus(self):
        QApplication.processEvents()
        self.activateWindow()
        self.raise_()
        self.password_input.setFocus(Qt.OtherFocusReason)
    
    def get_password(self):
        return self.password_input.text()

class PasswordField(QWidget):
    def __init__(self, read_only=True):
        super().__init__()
        layout = QHBoxLayout()
        layout.setContentsMargins(0, 0, 0, 0)
        
        self.line_edit = QLineEdit()
        self.line_edit.setReadOnly(read_only)
        if read_only:
            self.line_edit.setEchoMode(QLineEdit.Password)
        
        self.toggle_btn = QPushButton("👁")
        self.toggle_btn.setFixedWidth(30)
        self.toggle_btn.clicked.connect(self.toggle_visibility)
        self.visible = False
        
        layout.addWidget(self.line_edit)
        if read_only:
            layout.addWidget(self.toggle_btn)
        
        self.setLayout(layout)
    
    def toggle_visibility(self):
        if self.visible:
            self.line_edit.setEchoMode(QLineEdit.Password)
        else:
            self.line_edit.setEchoMode(QLineEdit.Normal)
        self.visible = not self.visible
    
    def text(self):
        return self.line_edit.text()
    
    def setText(self, text):
        self.line_edit.setText(text)
    
    def clear(self):
        self.line_edit.clear()
    
    def setReadOnly(self, read_only):
        self.line_edit.setReadOnly(read_only)


class KeyBunkerGUI(QMainWindow):
    def __init__(self, kb_instance):
        super().__init__()
        self.kb = kb_instance
        
        # Ensure authentication and vault initialization
        if self.kb._session_key is None:
            vault_exists = self.kb._locate_vault()
            prompt = "Enter keyBunker password:" if vault_exists else "No vault exists. Enter new password:"
            
            dialog = PasswordDialog(prompt, self)
            if dialog.exec_() == QDialog.Accepted:
                password = dialog.get_password()
                if password:
                    self.kb.register(password)
            else:
                QMessageBox.critical(self, "Error", "Authentication required")
                sys.exit(1)
        
        self.kb.complete_initialization()
        
        self.editing = False
        self.current_account = None
        self.standard_fields = ['Account Name', 'Username', 'Password', 'URL', 
                               'Account Number', 'TOTP Secret']
        self.keyring_fields = {'Password', 'TOTP Secret'}
        self.custom_fields = {}
        
        self.initUI()
        self.load_accounts()
    
    def initUI(self):
        self.setWindowTitle("KeyBunker Credential Manager")
        self.setGeometry(100, 100, 700, 600)
        
        central_widget = QWidget()
        self.setCentralWidget(central_widget)
        main_layout = QVBoxLayout()
        
        dropdown_layout = QHBoxLayout()
        dropdown_layout.addWidget(QLabel("Account:"))
        self.account_dropdown = QComboBox()
        self.account_dropdown.currentTextChanged.connect(self.on_account_selected)
        dropdown_layout.addWidget(self.account_dropdown, 1)
        main_layout.addLayout(dropdown_layout)
        
        scroll = QScrollArea()
        scroll.setWidgetResizable(True)
        scroll.setHorizontalScrollBarPolicy(Qt.ScrollBarAlwaysOff)
        
        form_container = QWidget()
        self.form_layout = QFormLayout()
        form_container.setLayout(self.form_layout)
        scroll.setWidget(form_container)
        
        main_layout.addWidget(scroll)
        
        button_layout = QHBoxLayout()
        self.new_btn = QPushButton("New Account")
        self.new_btn.clicked.connect(self.new_account)
        self.edit_btn = QPushButton("Edit")
        self.edit_btn.clicked.connect(self.enter_edit_mode)
        self.delete_btn = QPushButton("Delete")
        self.delete_btn.clicked.connect(self.delete_account)
        self.save_btn = QPushButton("Save")
        self.save_btn.clicked.connect(self.save_account)
        self.save_btn.hide()
        self.cancel_btn = QPushButton("Cancel")
        self.cancel_btn.clicked.connect(self.cancel_edit)
        self.cancel_btn.hide()
        self.add_field_btn = QPushButton("Add Custom Field")
        self.add_field_btn.clicked.connect(self.add_custom_field)
        self.add_field_btn.hide()
        
        button_layout.addWidget(self.new_btn)
        button_layout.addWidget(self.edit_btn)
        button_layout.addWidget(self.delete_btn)
        button_layout.addWidget(self.add_field_btn)
        button_layout.addStretch()
        button_layout.addWidget(self.save_btn)
        button_layout.addWidget(self.cancel_btn)
        
        main_layout.addLayout(button_layout)
        central_widget.setLayout(main_layout)
        
        self.field_widgets = {}
    
    def load_accounts(self):
        self.account_dropdown.clear()
        if self.kb._credentials_map and 'accounts' in self.kb._credentials_map:
            accounts = sorted(self.kb._credentials_map['accounts'].keys())
            self.account_dropdown.addItems(accounts)
    
    def on_account_selected(self, account_name):
        if not account_name or self.editing:
            return
        self.current_account = account_name
        self.display_account(account_name, read_only=True)
    
    def display_account(self, account_name, read_only=True):
        while self.form_layout.rowCount() > 0:
            self.form_layout.removeRow(0)
        self.field_widgets.clear()
        self.custom_fields.clear()
        
        if not account_name:
            return
        
        account_data = self.kb._credentials_map['accounts'].get(account_name, {})
        
        for field in self.standard_fields:
            if field == 'Account Name':
                widget = QLineEdit(account_name)
                widget.setReadOnly(read_only)
                self.field_widgets[field] = widget
                self.form_layout.addRow(field + ":", widget)
            elif field == 'Password' or field == 'TOTP Secret':
                widget = PasswordField(read_only=read_only)
                service_name = account_data.get('service_name', f"keyBunker_{account_name}")
                username = account_data.get('username', '')
                
                if field == 'Password':
                    value = keyring.get_password(service_name, username) or ''
                else:
                    totp_key = f"{service_name}_totp"
                    value = keyring.get_password(totp_key, username) or ''
                
                widget.setText(value)
                self.field_widgets[field] = widget
                self.form_layout.addRow(field + ":", widget)
            else:
                key = field.lower().replace(' ', '_')
                value = account_data.get(key, '')
                widget = QLineEdit(value or '')
                widget.setReadOnly(read_only)
                self.field_widgets[field] = widget
                self.form_layout.addRow(field + ":", widget)
        
        for key, value in account_data.items():
            if key not in ['service_name', 'username', 'url', 'account_number']:
                service_name = account_data.get('service_name', f"keyBunker_{account_name}")
                username = account_data.get('username', '')
                custom_key = f"{service_name}_{key}"
                keyring_value = keyring.get_password(custom_key, username)
                
                is_keyring_field = keyring_value is not None
                self.custom_fields[key] = is_keyring_field
                
                if is_keyring_field:
                    widget = PasswordField(read_only=read_only)
                    widget.setText(keyring_value or '')
                else:
                    widget = QLineEdit(value or '')
                    widget.setReadOnly(read_only)
                
                row_widget = QWidget()
                row_layout = QHBoxLayout()
                row_layout.setContentsMargins(0, 0, 0, 0)
                row_layout.addWidget(widget)
                
                if not read_only:
                    delete_btn = QPushButton("Delete")
                    delete_btn.clicked.connect(lambda checked, f=key: self.delete_custom_field(f))
                    row_layout.addWidget(delete_btn)
                
                row_widget.setLayout(row_layout)
                self.field_widgets[key] = widget
                self.form_layout.addRow(key + ":", row_widget)
    
    def enter_edit_mode(self):
        if not self.current_account:
            return
        
        self.editing = True
        self.account_dropdown.setEnabled(False)
        self.new_btn.setEnabled(False)
        self.edit_btn.hide()
        self.delete_btn.hide()
        self.save_btn.show()
        self.cancel_btn.show()
        self.add_field_btn.show()
        
        self.display_account(self.current_account, read_only=False)
    
    def cancel_edit(self):
        reply = QMessageBox.question(self, 'Confirm Cancel',
                                     'Discard all changes?',
                                     QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.No:
            return
        
        self.editing = False
        self.account_dropdown.setEnabled(True)
        self.new_btn.setEnabled(True)
        self.edit_btn.show()
        self.delete_btn.show()
        self.save_btn.hide()
        self.cancel_btn.hide()
        self.add_field_btn.hide()
        
        if self.current_account:
            self.display_account(self.current_account, read_only=True)
        else:
            while self.form_layout.rowCount() > 0:
                self.form_layout.removeRow(0)
            self.field_widgets.clear()
    
    def save_account(self):
        account_name = self.field_widgets['Account Name'].text().strip()
        if not account_name:
            QMessageBox.warning(self, "Validation Error", "Account Name is required")
            return
        
        password = self.field_widgets['Password'].text()
        if not password:
            QMessageBox.warning(self, "Validation Error", "Password is required")
            return
        
        existing_accounts = [a.lower() for a in self.kb._credentials_map['accounts'].keys()]
        if self.current_account:
            existing_accounts.remove(self.current_account.lower())
        
        if account_name.lower() in existing_accounts:
            QMessageBox.warning(self, "Validation Error", 
                              "An account with this name already exists (case-insensitive)")
            return
        
        service_name = f"keyBunker_{account_name}"
        username = self.field_widgets['Username'].text().strip()
        
        account_data = {
            'service_name': service_name,
            'username': username,
            'url': self.field_widgets['URL'].text().strip(),
            'account_number': self.field_widgets['Account Number'].text().strip()
        }
        
        if self.current_account and self.current_account != account_name:
            old_service = f"keyBunker_{self.current_account}"
            old_username = self.kb._credentials_map['accounts'][self.current_account]['username']
            
            try:
                keyring.delete_password(old_service, old_username)
            except:
                pass
            
            try:
                keyring.delete_password(f"{old_service}_totp", old_username)
            except:
                pass
            
            old_data = self.kb._credentials_map['accounts'][self.current_account]
            for key in old_data.keys():
                if key not in ['service_name', 'username', 'url', 'account_number']:
                    try:
                        keyring.delete_password(f"{old_service}_{key}", old_username)
                    except:
                        pass
            
            del self.kb._credentials_map['accounts'][self.current_account]
        
        keyring.set_password(service_name, username, password)
        
        totp_secret = self.field_widgets['TOTP Secret'].text().strip()
        keyring.set_password(f"{service_name}_totp", username, totp_secret)
        
        for field_name, is_keyring in self.custom_fields.items():
            widget = self.field_widgets[field_name]
            value = widget.text() if hasattr(widget, 'text') else widget.line_edit.text()
            
            if is_keyring:
                keyring.set_password(f"{service_name}_{field_name}", username, value)
            else:
                account_data[field_name] = value
        
        self.kb._credentials_map['accounts'][account_name] = account_data
        self.kb._write_vault()
        
        self.current_account = account_name
        self.editing = False
        self.account_dropdown.setEnabled(True)
        self.new_btn.setEnabled(True)
        self.edit_btn.show()
        self.delete_btn.show()
        self.save_btn.hide()
        self.cancel_btn.hide()
        self.add_field_btn.hide()
        
        self.load_accounts()
        
        index = self.account_dropdown.findText(account_name)
        if index >= 0:
            self.account_dropdown.setCurrentIndex(index)
    
    def new_account(self):
        self.current_account = None
        self.editing = True
        self.account_dropdown.setEnabled(False)
        self.new_btn.setEnabled(False)
        self.edit_btn.hide()
        self.delete_btn.hide()
        self.save_btn.show()
        self.cancel_btn.show()
        self.add_field_btn.show()
        
        while self.form_layout.rowCount() > 0:
            self.form_layout.removeRow(0)
        self.field_widgets.clear()
        self.custom_fields.clear()
        
        for field in self.standard_fields:
            if field == 'Password' or field == 'TOTP Secret':
                widget = PasswordField(read_only=False)
            else:
                widget = QLineEdit()
            
            self.field_widgets[field] = widget
            self.form_layout.addRow(field + ":", widget)
    
    def delete_account(self):
        if not self.current_account:
            return
        
        reply = QMessageBox.question(self, 'Confirm Delete',
                                     f'Delete account "{self.current_account}"?',
                                     QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.No:
            return
        
        account_data = self.kb._credentials_map['accounts'][self.current_account]
        service_name = account_data['service_name']
        username = account_data.get('username', '')
        
        try:
            keyring.delete_password(service_name, username)
        except:
            pass
        
        try:
            keyring.delete_password(f"{service_name}_totp", username)
        except:
            pass
        
        for key in account_data.keys():
            if key not in ['service_name', 'username', 'url', 'account_number']:
                try:
                    keyring.delete_password(f"{service_name}_{key}", username)
                except:
                    pass
        
        del self.kb._credentials_map['accounts'][self.current_account]
        self.kb._write_vault()
        
        self.current_account = None
        self.load_accounts()
        
        while self.form_layout.rowCount() > 0:
            self.form_layout.removeRow(0)
        self.field_widgets.clear()
    
    def add_custom_field(self):
        dialog = CustomFieldDialog(self)
        if dialog.exec_() == QDialog.Accepted:
            field_data = dialog.get_field_data()
            field_name = field_data['name']
            
            if not field_name:
                QMessageBox.warning(self, "Error", "Field name cannot be empty")
                return
            
            if field_name in self.field_widgets:
                QMessageBox.warning(self, "Error", "Field already exists")
                return
            
            is_keyring = field_data['store_in_keyring']
            self.custom_fields[field_name] = is_keyring
            
            if is_keyring:
                widget = PasswordField(read_only=False)
            else:
                widget = QLineEdit()
            
            row_widget = QWidget()
            row_layout = QHBoxLayout()
            row_layout.setContentsMargins(0, 0, 0, 0)
            row_layout.addWidget(widget)
            
            delete_btn = QPushButton("Delete")
            delete_btn.clicked.connect(lambda: self.delete_custom_field(field_name))
            row_layout.addWidget(delete_btn)
            
            row_widget.setLayout(row_layout)
            self.field_widgets[field_name] = widget
            self.form_layout.addRow(field_name + ":", row_widget)
    
    def delete_custom_field(self, field_name):
        reply = QMessageBox.question(self, 'Confirm Delete',
                                     f'Delete custom field "{field_name}"?',
                                     QMessageBox.Yes | QMessageBox.No)
        if reply == QMessageBox.No:
            return
        
        for i in range(self.form_layout.rowCount()):
            label = self.form_layout.itemAt(i, QFormLayout.LabelRole)
            if label and label.widget().text() == field_name + ":":
                self.form_layout.removeRow(i)
                break
        
        del self.field_widgets[field_name]
        del self.custom_fields[field_name]

class keyBunker:
    _instance = None
    _session_key = None
    _credentials_map = None
    _initialized = False
    _vault_path = None

    def __new__(cls):
        if cls._instance is None:
            cls._instance = super(keyBunker, cls).__new__(cls)
            cls._instance._basic_initialize()
        return cls._instance

    def _basic_initialize(self):
        self._lock = threading.Lock()
        #print("Basic initialization of keyBunker", flush=True)
        self._credentials_map = {"accounts": {}, "salt": None}
        self._ipc_server_thread = None
        self._search_locations = [
            os.path.join(os.path.expanduser("~"), ".keybunker", "vault.yaml"),
            os.path.abspath(os.path.join(os.path.dirname(__file__), "vault.yaml")),
            "D:\\agent\\vault.yaml",
            os.path.join(os.getcwd(), "credentials_map.yaml"),
            os.path.join(os.getcwd(), "vault.yaml"),
            os.path.join(os.path.dirname(os.getcwd()), "credentials_map.yaml"),
            os.path.join(os.path.dirname(os.getcwd()), "vault.yaml")
        ]

    def _locate_vault(self):
        if self._vault_path and os.path.exists(self._vault_path):
            #print(f"Using existing vault path: {self._vault_path}", flush=True)
            return True
        #print("Searching for vault file:", flush=True)
        for loc in self._search_locations:
            if os.path.exists(loc):
                #print(f"Found vault at: {loc}", flush=True)
                self._vault_path = loc
                return True
        self._vault_path = self._search_locations[0]
        os.makedirs(os.path.dirname(self._vault_path), exist_ok=True)
        #print(f"Will create new vault at: {self._vault_path}", flush=True)
        return False

    def _encrypt_data(self, data: dict) -> bytes:
        if not self._session_key:
            raise ValueError("Session key required for encryption")
        fernet = Fernet(self._session_key)
        yaml_str = yaml.safe_dump(data)
        return fernet.encrypt(yaml_str.encode())
    
    def _decrypt_data(self, encrypted_data: bytes) -> dict:
        if not self._session_key:
            raise ValueError("Session key required for decryption")
        fernet = Fernet(self._session_key)
        decrypted_yaml = fernet.decrypt(encrypted_data).decode()
        return yaml.safe_load(decrypted_yaml)
    
    def _write_vault(self):
        #print(f"Writing encrypted vault to {self._vault_path}", flush=True)
        os.makedirs(os.path.dirname(self._vault_path), exist_ok=True)
        
        # Extract the salt for storing separately
        salt = self._credentials_map.get("salt")
        
        # Create a copy of credentials_map for encryption
        encrypt_data = self._credentials_map.copy() 
        
        # Encrypt the data
        encrypted_data = self._encrypt_data(encrypt_data)
        
        # Create the vault structure with separate salt and encrypted data
        vault_data = {
            "salt": salt,  # Store salt unencrypted
            "encrypted_data": base64.b64encode(encrypted_data).decode()  # Store encrypted data as base64
        }
        
        # Write the vault to file
        temp_path = f"{self._vault_path}.tmp"
        with open(temp_path, "w") as f:
            yaml.safe_dump(vault_data, f)
        
        os.replace(temp_path, self._vault_path)
        os.chmod(self._vault_path, 0o600)
        #print("Vault written successfully", flush=True)
    
    def _read_vault(self):
        #print(f"Reading encrypted vault from {self._vault_path}", flush=True)
        
        # Load the file content
        with open(self._vault_path, "r") as f:
            vault_data = yaml.safe_load(f)
        
        # Check if this is the new format with separate salt and encrypted data
        if isinstance(vault_data, dict) and "salt" in vault_data and "encrypted_data" in vault_data:
            # Extract the salt
            salt = vault_data["salt"]
            
            # If we have a session key, try to decrypt the data
            if self._session_key:
                try:
                    # Decode the base64 encrypted data
                    encrypted_data = base64.b64decode(vault_data["encrypted_data"])
                    
                    # Decrypt the data
                    decrypted_data = self._decrypt_data(encrypted_data)
                    
                    # Use the decrypted data
                    self._credentials_map = decrypted_data
                    
                    # Ensure the salt is in the credentials map
                    self._credentials_map["salt"] = salt
                except Exception as e:
                    #print(f"Error decrypting vault: {e}", flush=True)
                    self._credentials_map = {"accounts": {}, "salt": salt}
            else:
                # If no session key, just return the salt in a basic structure
                self._credentials_map = {"accounts": {}, "salt": salt}
        else:
            # Legacy format - just use the data directly
            self._credentials_map = vault_data if vault_data else {"accounts": {}, "salt": None}
        
        # Ensure the structure is valid
        if self._credentials_map is None:
            #print("Warning: Loaded credentials map is None", flush=True)
            self._credentials_map = {"accounts": {}, "salt": None}
        elif "accounts" not in self._credentials_map:
            #print("Warning: No 'accounts' key in loaded map", flush=True)
            self._credentials_map["accounts"] = {}
        
        #print(f"Loaded {len(self._credentials_map['accounts'])} accounts", flush=True)
        #print("Vault read completed", flush=True)

    def complete_initialization(self):
        if not self._initialized:
            #print("Completing keyBunker initialization", flush=True)
            vault_exists = self._locate_vault()
            if vault_exists and self._session_key:
                self._read_vault()
            self._initialized = True

    def _derive_key(self, password: str, salt: bytes = None) -> tuple[bytes, bytes]:
        #print("Deriving key", flush=True)
        if salt is None:
            salt = os.urandom(16)
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
        )
        key = base64.urlsafe_b64encode(kdf.derive(password.encode()))
        #print("Key derived successfully", flush=True)
        return key, salt

    def register(self, password: str):
        #print("Registering with password", flush=True)
        if self._session_key:
            #print("Already registered", flush=True)
            return
        
        vault_exists = self._locate_vault()
        
        if not vault_exists:
            # New vault - generate salt and key
            #print("Creating new vault", flush=True)
            key, salt = self._derive_key(password)
            self._session_key = key
            self._credentials_map = {"accounts": {}, "salt": base64.b64encode(salt).decode()}
            self._write_vault()
        else:
            # Existing vault - read to get salt but don't try to decrypt
            #print(f"Verifying vault at {self._vault_path}", flush=True)
            
            # Read the vault file to get the salt
            with open(self._vault_path, "r") as f:
                vault_data = yaml.safe_load(f)
            
            # Extract the salt
            if isinstance(vault_data, dict) and "salt" in vault_data:
                salt_b64 = vault_data["salt"]
            else:
                # Legacy format or missing salt
                #print("Warning: Could not find salt in vault", flush=True)
                key, salt = self._derive_key(password)
                self._session_key = key
                self._credentials_map = {"accounts": {}, "salt": base64.b64encode(salt).decode()}
                self._write_vault()
                return
            
            # Derive the key using the salt
            stored_salt = base64.b64decode(salt_b64)
            key, _ = self._derive_key(password, salt=stored_salt)
            
            # Set the session key
            self._session_key = key
            
            # Now read and decrypt the vault
            self._read_vault()
        
        #print("Registration complete", flush=True)

    def clearSession(self):
        #print("Clearing session", flush=True)
        self._session_key = None
        self._credentials_map = None
        if self._ipc_server_thread:
            self._ipc_server_thread = None

    def _get_session_key_from_ipc(self):
        """Attempt to fetch session key from IPC."""
        if sys.platform == "win32":
            import win32pipe
            import win32file
            
            pipe_name = r'\\.\pipe\keybunker'
            pipe_handle = win32file.CreateFile(
                pipe_name,
                win32file.GENERIC_READ,
                0, None,
                win32file.OPEN_EXISTING,
                0, None
            )
            
            # Read from the pipe
            result, data = win32file.ReadFile(pipe_handle, 1024)
            win32file.CloseHandle(pipe_handle)
            
            if data:
                #print("Session key retrieved via Windows named pipe", flush=True)
                return data
        else:
            socket_path = "/tmp/keybunker_socket"
            client = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
            client.connect(socket_path)
            session_key = client.recv(1024)
            client.close()
            if session_key:
                #print("Session key retrieved via Unix socket", flush=True)
                return session_key

    def getSessionKey(self) -> str | None:
        #print("Getting session key", flush=True)
        if self._session_key:
            return self._session_key
        is_main_module = __name__ == "__main__" or sys.modules["__main__"].__file__.endswith("keyBunker.py")
        if not is_main_module:
            try:
                #print("Attempting to get session key via IPC...", flush=True)
                ipc_key = self._get_session_key_from_ipc()
                if ipc_key:
                    #print("IPC session key retrieved, setting as session key", flush=True)
                    self._session_key = ipc_key
                    
                    # Explicitly trigger vault loading after getting key via IPC
                    #print("Explicitly loading vault after IPC connection", flush=True)
                    self._locate_vault()
                    self._read_vault()
                    
                    return self._session_key
            except Exception as e:
                print(f"IPC failed: {e}", flush=True)
        vault_exists = self._locate_vault()
        prompt = "Enter keyBunker password: " if vault_exists else "No vault exists. Enter new password: "
        password = input(prompt)
        if password:
            self.register(password)
            return self._session_key
        #print("Password prompt skipped", flush=True)
        return None

    def _load_vault(self):
        #print("Loading vault", flush=True)
        if self._session_key is None:
            self._session_key = self.getSessionKey()
            if self._session_key is None:
                raise ValueError("No session key to decrypt vault")
        vault_exists = self._locate_vault()
        if vault_exists:
            self._read_vault()
        else:
            #print("No vault found, initializing empty accounts", flush=True)
            self._credentials_map = {"accounts": {}, "salt": base64.b64encode(os.urandom(16)).decode()}
            self._write_vault()
        #print("Vault load completed", flush=True)

    def getPassword(self, account: str, username: str) -> str:
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("Not authenticated")
        if self._credentials_map is None:
            self._load_vault()
        account_data = self._credentials_map["accounts"].get(account)
        if not account_data or account_data["username"] != username:
            raise ValueError(f"Account {account} or username {username} not found")
        password = keyring.get_password(account_data["service_name"], username)
        if not password:
            raise ValueError(f"No password for {account}")
        return password
    
    def getURL(self, account: str, username: str) -> str | None:
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("keyBunker not authenticated; cannot retrieve URL")
        if self._credentials_map is None:
            self._load_vault()
        account_data = self._credentials_map["accounts"].get(account)
        if not account_data or account_data["username"] != username:
            raise ValueError(f"Account {account} or username {username} not found")
        return account_data.get("url")

    def getAccountNumber(self, account: str, username: str) -> str | None:
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("keyBunker not authenticated; cannot retrieve account number")
        if self._credentials_map is None:
            self._load_vault()
        account_data = self._credentials_map["accounts"].get(account)
        if not account_data or account_data["username"] != username:
            raise ValueError(f"Account {account} or username {username} not found")
        return account_data.get("account_number")
    
    def extract_root_domain(self, url_or_domain):
        try:
            # Handle both full URLs and domain-only strings
            if url_or_domain.startswith(('http://', 'https://')):
                from urllib.parse import urlparse
                hostname = urlparse(url_or_domain).hostname
            else:
                hostname = url_or_domain
            
            parts = hostname.split('.')
            return '.'.join(parts[-2:])  # Return domain.tld
        except:
            return url_or_domain
    
    def getCredentialsForWebsite(self, domain: str, username: str) -> dict:
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("KeyBunker not authenticated")
        if self._credentials_map is None:
            self._load_vault()
        domain_root = self.extract_root_domain(domain)
        
        # Find accounts where URL contains domain and username matches
        matches = []
        for account_name, account_data in self._credentials_map["accounts"].items():
            account_url = account_data.get("url", "")
            account_username = account_data.get("username", "")
            
            account_root = self.extract_root_domain(account_url)
            
            if domain_root == account_root and username == account_username:
                matches.append((account_name, account_data))
        
        if not matches:
            raise ValueError(f"No credentials found for domain '{domain}' with username '{username}'")
        
        if len(matches) > 1:
            raise ValueError(f"Multiple accounts found for domain '{domain}' with username '{username}'")
        
        account_name, account_data = matches[0]
        password = keyring.get_password(account_data["service_name"], username)
        if not password:
            raise ValueError(f"Password not found in keyring for account '{account_name}'")
        
        return {
            "username": username,
            "password": password,
            "url": account_data.get("url")
        }

    # Token handling methods
    def hasToken(self, account: str, username: str) -> bool:
        """
        Check if an auth token exists for the given account and username.
        
        Args:
            account (str): The account name
            username (str): The username
            
        Returns:
            bool: True if a token exists, False otherwise
        
        Raises:
            ValueError: If account not found or not authenticated
        """
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("keyBunker not authenticated; cannot check token")
        if self._credentials_map is None:
            self._load_vault()
            
        account_data = self._credentials_map["accounts"].get(account)
        if not account_data or account_data["username"] != username:
            raise ValueError(f"Account {account} or username {username} not found")
            
        # Check if auth_token exists in account metadata
        return account_data.get("has_token", False)

    def getToken(self, account: str, username: str) -> str:
        """
        Retrieve the authentication token for the given account and username.
        
        Args:
            account (str): The account name
            username (str): The username
            
        Returns:
            str: The authentication token
            
        Raises:
            ValueError: If account not found, token doesn't exist, or not authenticated
        """
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("keyBunker not authenticated; cannot retrieve token")
        if self._credentials_map is None:
            self._load_vault()
            
        account_data = self._credentials_map["accounts"].get(account)
        if not account_data or account_data["username"] != username:
            raise ValueError(f"Account {account} or username {username} not found")
            
        # Check if the account has a token flag set
        if not account_data.get("has_token", False):
            raise ValueError(f"No token exists for {account}")
            
        # Get token from keyring
        token_key = f"{account_data['service_name']}_token"
        token = keyring.get_password(token_key, username)
        if not token:
            # If token doesn't exist in keyring but flag is set, fix the inconsistency
            account_data["has_token"] = False
            self._write_vault()
            raise ValueError(f"Token not found for {account} (inconsistent state fixed)")
            
        return token

    def setToken(self, account: str, username: str, token: str) -> None:
        """
        Set an authentication token for the given account and username.
        
        Args:
            account (str): The account name
            username (str): The username
            token (str): The authentication token to store
            
        Raises:
            ValueError: If account not found or not authenticated
        """
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("keyBunker not authenticated; cannot set token")
        if self._credentials_map is None:
            self._load_vault()
            
        with self._lock:  # Use lock to prevent race conditions
            account_data = self._credentials_map["accounts"].get(account)
            if not account_data or account_data["username"] != username:
                raise ValueError(f"Account {account} or username {username} not found")
                
            # Set token in keyring
            token_key = f"{account_data['service_name']}_token"
            keyring.set_password(token_key, username, token)
            
            # Update metadata to indicate token exists
            account_data["has_token"] = True
            self._write_vault()
            #print(f"Token set for {account}", flush=True)

    def removeToken(self, account: str, username: str) -> None:
        """
        Remove the authentication token for the given account and username.
        
        Args:
            account (str): The account name
            username (str): The username
            
        Raises:
            ValueError: If account not found or not authenticated
        """
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("keyBunker not authenticated; cannot remove token")
        if self._credentials_map is None:
            self._load_vault()
            
        with self._lock:
            account_data = self._credentials_map["accounts"].get(account)
            if not account_data or account_data["username"] != username:
                raise ValueError(f"Account {account} or username {username} not found")
                
            # Only attempt to delete if a token exists
            if account_data.get("has_token", False):
                token_key = f"{account_data['service_name']}_token"
                try:
                    keyring.delete_password(token_key, username)
                except keyring.errors.PasswordDeleteError:
                    print(f"Warning: Could not delete token from keyring for {account}", flush=True)
                
                # Update metadata
                account_data["has_token"] = False
                self._write_vault()
                #print(f"Token removed for {account}", flush=True)

    def generateTOTP(self, account: str, username: str) -> str:
        """
        Generate a TOTP code for the given account and username.
        
        Args:
            account (str): The account name
            username (str): The username
            
        Returns:
            str: A 6-digit TOTP code
            
        Raises:
            ValueError: If account not found, TOTP secret doesn't exist, or not authenticated
        """
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("keyBunker not authenticated; cannot generate TOTP")
        if self._credentials_map is None:
            self._load_vault()
            
        account_data = self._credentials_map["accounts"].get(account)
        if not account_data or account_data["username"] != username:
            raise ValueError(f"Account {account} or username {username} not found")
            
        # Get the TOTP secret from keyring
        service_name = account_data["service_name"]
        totp_key = f"{service_name}_totp"
        totp_secret = keyring.get_password(totp_key, username)
        
        if not totp_secret:
            raise ValueError(f"No TOTP secret exists for {account}")
        
        # Generate TOTP code
        totp = pyotp.TOTP(totp_secret)
        return totp.now()

    def getTOTPForWebsite(self, domain: str, username: str) -> str:
        """
        Generate a TOTP code for the given domain and username.
        This parallels getCredentialsForWebsite for TOTP generation.
        
        Args:
            domain (str): The domain or URL
            username (str): The username
            
        Returns:
            str: A 6-digit TOTP code
            
        Raises:
            ValueError: If no matching credentials found or TOTP secret doesn't exist
        """
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("KeyBunker not authenticated")
        if self._credentials_map is None:
            self._load_vault()
        
        domain_root = self.extract_root_domain(domain)
        
        # Find accounts where URL contains domain and username matches
        matches = []
        for account_name, account_data in self._credentials_map["accounts"].items():
            account_url = account_data.get("url", "")
            account_username = account_data.get("username", "")
            
            account_root = self.extract_root_domain(account_url)
            
            if domain_root == account_root and username == account_username:
                matches.append(account_name)
        
        if not matches:
            raise ValueError(f"No credentials found for domain '{domain}' with username '{username}'")
        
        if len(matches) > 1:
            raise ValueError(f"Multiple accounts found for domain '{domain}' with username '{username}'")
        
        # Generate TOTP using the found account
        return self.generateTOTP(matches[0], username)

    def registerScript(self, script_name: str, script_hash: str, manifest: dict) -> None:
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("Not authenticated")
        if self._credentials_map is None:
            self._load_vault()
        if not isinstance(script_name, str) or not script_name:
            raise ValueError(f"registerScript: script_name must be non-empty string, got {type(script_name).__name__}")
        if not isinstance(script_hash, str) or len(script_hash) != 64:
            raise ValueError(f"registerScript: script_hash must be 64-char hex, got len={len(script_hash) if isinstance(script_hash, str) else type(script_hash).__name__}")
        if not isinstance(manifest, dict):
            raise ValueError(f"registerScript: manifest must be dict, got {type(manifest).__name__}")
        if "script_registry" not in self._credentials_map:
            self._credentials_map["script_registry"] = {}
        self._credentials_map["script_registry"][script_name] = {
            "hash": script_hash,
            "manifest": manifest,
            "registered_at": datetime.now().isoformat()
        }
        self._write_vault()

    def getScriptRegistration(self, script_name: str) -> dict | None:
        if self._session_key is None:
            self._session_key = self.getSessionKey()
        if self._session_key is None:
            raise ValueError("Not authenticated")
        if self._credentials_map is None:
            self._load_vault()
        registry = self._credentials_map.get("script_registry") or {}
        return registry.get(script_name)

    def test_credentials_in_vault(self, account: str, username: str) -> bool:
        """Test if stored credentials work against their URL."""
        if self._credentials_map is None:
            self._load_vault()
        account_data = self._credentials_map["accounts"].get(account)
        if not account_data or account_data["username"] != username:
            raise ValueError(f"Account {account} or username {username} not found")
        url = account_data.get("url")
        if not url or not url.startswith("https://"):
            raise ValueError("Valid HTTPS URL required for testing")
            
        # Modified to check for token first if available
        auth_headers = {}
        try:
            if account_data.get("has_token", False):
                # Try to use token authentication first
                token = self.getToken(account, username)
                auth_headers = {"Authorization": f"Bearer {token}"}
                response = requests.get(url, headers=auth_headers, timeout=10)
                return response.status_code in (200, 302)
            else:
                # Fall back to password authentication
                password = self.getPassword(account, username)
                response = requests.post(url, data={"username": username, "password": password}, timeout=10)
                return response.status_code in (200, 302)
        except requests.RequestException:
            return False

if __name__ == "__main__":
    app = QApplication(sys.argv)
    KB = keyBunker()
    gui = KeyBunkerGUI(KB)
    gui.show()
    sys.exit(app.exec_())