#!/usr/bin/env python3
"""
Secure Password Vault for Python Applications
Uses industry-standard encryption (AES-256-GCM) with PBKDF2 key derivation

pip install cryptography

# Create a new vault
python vault.py create

# Store a password
python vault.py store gmail john.doe@gmail.com

# Retrieve a password
python vault.py get gmail

# List all services
python vault.py list

# Delete a password
python vault.py delete gmail

# Change master password
python vault.py change-master

"""

import os
import json
import base64
import getpass
import hashlib
from pathlib import Path
from typing import Dict, Optional, Any
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
from cryptography.exceptions import InvalidSignature, InvalidTag
import secrets
import argparse
import sys
from datetime import datetime

class SecureVault:
    """
    Secure password vault using AES-256-GCM encryption with PBKDF2 key derivation.
    
    Security Features:
    - AES-256-GCM encryption (authenticated encryption)
    - PBKDF2-SHA256 key derivation (600,000 iterations)
    - Cryptographically secure random salts and nonces
    - Master password verification without storing password
    - Secure memory handling where possible
    """
    
    def __init__(self, vault_path: str = None):
        resolved = vault_path or os.getenv('VPATH')
        self.vault_path = Path(resolved) if resolved else Path.home() / ".secure_vault" / "vault.enc"
        self.vault_path.parent.mkdir(exist_ok=True)
        if os.name != 'nt':
            self.vault_path.parent.chmod(0o700)
        self._master_key = None
        
    def _derive_key(self, password: str, salt: bytes) -> bytes:
        """Derive encryption key from master password using PBKDF2"""
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,  # 256 bits for AES-256
            salt=salt,
            iterations=600000,  # OWASP recommended minimum for 2023+
        )
        return kdf.derive(password.encode('utf-8'))
    
    def _create_password_hash(self, password: str, salt: bytes) -> str:
        """Create a hash of the master password for verification"""
        return hashlib.pbkdf2_hmac('sha256', password.encode('utf-8'), salt, 600000).hex()
    
    def _encrypt_data(self, data: Dict[str, Any], master_password: str) -> bytes:
        """Encrypt vault data"""
        # Generate cryptographically secure salt and nonce
        salt = secrets.token_bytes(32)
        nonce = secrets.token_bytes(12)  # 96-bit nonce for GCM
        
        # Derive encryption key
        key = self._derive_key(master_password, salt)
        
        # Create password hash for verification
        password_hash = self._create_password_hash(master_password, salt)
        
        # Prepare data with metadata
        vault_data = {
            'version': '1.0',
            'created': datetime.now().isoformat(),
            'password_hash': password_hash,
            'entries': data
        }
        
        # Encrypt the data
        aesgcm = AESGCM(key)
        ciphertext = aesgcm.encrypt(nonce, json.dumps(vault_data).encode('utf-8'), None)
        
        # Package everything together
        encrypted_package = {
            'salt': base64.b64encode(salt).decode('utf-8'),
            'nonce': base64.b64encode(nonce).decode('utf-8'),
            'ciphertext': base64.b64encode(ciphertext).decode('utf-8')
        }
        
        return json.dumps(encrypted_package).encode('utf-8')
    
    def _decrypt_data(self, encrypted_data: bytes, master_password: str) -> Dict[str, Any]:
        """Decrypt vault data"""
        try:
            # Parse encrypted package
            package = json.loads(encrypted_data.decode('utf-8'))
            salt = base64.b64decode(package['salt'])
            nonce = base64.b64decode(package['nonce'])
            ciphertext = base64.b64decode(package['ciphertext'])
            
            # Derive decryption key
            key = self._derive_key(master_password, salt)
            
            # Decrypt the data
            aesgcm = AESGCM(key)
            decrypted_data = aesgcm.decrypt(nonce, ciphertext, None)
            vault_data = json.loads(decrypted_data.decode('utf-8'))
            
            # Verify master password
            expected_hash = self._create_password_hash(master_password, salt)
            if vault_data['password_hash'] != expected_hash:
                raise ValueError("Invalid master password")
            
            return vault_data['entries']
            
        except InvalidTag:
            raise ValueError("Incorrect vault master password.")
        except (json.JSONDecodeError, KeyError) as e:
            raise ValueError("Vault file is corrupt or unreadable.") from e
        except InvalidSignature:
            raise ValueError("Vault data signature invalid -- file may have been tampered with.")
    
    def create_vault(self, master_password: str = None) -> bool:
        """Create a new vault"""
        if self.vault_path.exists():
            response = input(f"Vault already exists at {self.vault_path}. Overwrite? (y/N): ")
            if response.lower() != 'y':
                print("Vault creation cancelled.")
                return False
        
        if not master_password:
            master_password = os.getenv('ID')

        if not master_password:
            while True:
                master_password = getpass.getpass("Enter NEW vault master password: ")
                confirm_password = getpass.getpass("Confirm vault master password: ")

                if master_password != confirm_password:
                    print("Passwords don't match. Please try again.")
                    continue

                if len(master_password) < 12:
                    print("Master password must be at least 12 characters long.")
                    continue

                break

        if len(master_password) < 12:
            raise ValueError("Master password must be at least 12 characters long.")

        # Create empty vault
        encrypted_data = self._encrypt_data({}, master_password)
        
        # Write to file with secure permissions
        self.vault_path.write_bytes(encrypted_data)
        if os.name != 'nt':
            self.vault_path.chmod(0o600)
        
        print(f"Vault created successfully at {self.vault_path}")
        return True
    
    def unlock_vault(self, master_password: str = None) -> bool:
        """Unlock the vault with master password"""
        if not self.vault_path.exists():
            print(f"Vault not found at {self.vault_path}")
            return False
        
        if not master_password:
            master_password = os.getenv('ID') or getpass.getpass("Enter vault master password: ")
        
        try:
            encrypted_data = self.vault_path.read_bytes()
            self._vault_data = self._decrypt_data(encrypted_data, master_password)
            self._master_password = master_password
            print("Vault unlocked successfully.")
            return True
        except ValueError as e:
            print(f"Failed to unlock vault: {e}")
            return False
    
    def _save_vault(self):
        """Save current vault state"""
        if not hasattr(self, '_vault_data') or not self._master_password:
            raise RuntimeError("Vault is not unlocked")
        
        encrypted_data = self._encrypt_data(self._vault_data, self._master_password)
        self.vault_path.write_bytes(encrypted_data)
    
    def store_password(self, service: str, username: str, password: str, notes: str = "") -> bool:
        """Store a password in the vault"""
        if not hasattr(self, '_vault_data'):
            print("Vault is not unlocked. Please unlock first.")
            return False
        
        entry = {
            'username': username,
            'password': password,
            'notes': notes,
            'created': datetime.now().isoformat(),
            'modified': datetime.now().isoformat()
        }
        
        if service in self._vault_data:
            entry['created'] = self._vault_data[service].get('created', entry['created'])
        
        self._vault_data[service] = entry
        self._save_vault()
        print(f"Password for '{service}' stored successfully.")
        return True
    
    def get_password(self, service: str) -> Optional[Dict[str, str]]:
        """Retrieve a password from the vault"""
        if not hasattr(self, '_vault_data'):
            print("Vault is not unlocked. Please unlock first.")
            return None
        
        return self._vault_data.get(service)
    
    def list_services(self) -> list:
        """List all services in the vault"""
        if not hasattr(self, '_vault_data'):
            print("Vault is not unlocked. Please unlock first.")
            return []
        
        return list(self._vault_data.keys())
    
    def delete_password(self, service: str) -> bool:
        """Delete a password from the vault"""
        if not hasattr(self, '_vault_data'):
            print("Vault is not unlocked. Please unlock first.")
            return False
        
        if service in self._vault_data:
            del self._vault_data[service]
            self._save_vault()
            print(f"Password for '{service}' deleted successfully.")
            return True
        else:
            print(f"Service '{service}' not found in vault.")
            return False
    
    def change_master_password(self, new_password: str = None) -> bool:
        """Change the master password"""
        if not hasattr(self, '_vault_data'):
            print("Vault is not unlocked. Please unlock first.")
            return False
        
        if not new_password:
            while True:
                new_password = getpass.getpass("Enter NEW vault master password: ")
                confirm_password = getpass.getpass("Confirm new vault master password: ")
                
                if new_password != confirm_password:
                    print("Passwords don't match. Please try again.")
                    continue
                
                if len(new_password) < 12:
                    print("Master password must be at least 12 characters long.")
                    continue
                
                break
        
        self._master_password = new_password
        self._save_vault()
        print("Master password changed successfully.")
        return True
    
    def lock_vault(self):
        """Lock the vault and clear sensitive data from memory"""
        if hasattr(self, '_vault_data'):
            # Clear sensitive data
            self._vault_data.clear()
            del self._vault_data
        
        if hasattr(self, '_master_password'):
            # Overwrite password in memory (best effort)
            self._master_password = 'x' * len(self._master_password)
            del self._master_password
        
        print("Vault locked.")


def main():
    """Command-line interface for the password vault"""
    parser = argparse.ArgumentParser(description="Secure Password Vault")
    parser.add_argument("--vault-path", help="Path to vault file")
    
    subparsers = parser.add_subparsers(dest="command", help="Available commands")
    
    # Create vault
    create_parser = subparsers.add_parser("create", help="Create a new vault")
    
    # Store password
    store_parser = subparsers.add_parser("store", help="Store a password")
    store_parser.add_argument("service", help="Service name")
    store_parser.add_argument("username", help="Username")
    store_parser.add_argument("--notes", default="", help="Additional notes")
    
    # Get password
    get_parser = subparsers.add_parser("get", help="Retrieve a password")
    get_parser.add_argument("service", help="Service name")
    
    # List services
    list_parser = subparsers.add_parser("list", help="List all services")
    
    # Delete password
    delete_parser = subparsers.add_parser("delete", help="Delete a password")
    delete_parser.add_argument("service", help="Service name")
    
    # Change master password
    change_parser = subparsers.add_parser("change-master", help="Change master password")
    
    args = parser.parse_args()
    
    if not args.command:
        parser.print_help()
        return
    
    vault = SecureVault(args.vault_path)
    
    if args.command == "create":
        vault.create_vault()
    
    elif args.command == "store":
        if vault.unlock_vault():
            password = getpass.getpass(f"Enter password to store for '{args.service}': ")
            vault.store_password(args.service, args.username, password, args.notes)
            vault.lock_vault()
    
    elif args.command == "get":
        if vault.unlock_vault():
            entry = vault.get_password(args.service)
            if entry:
                print(f"Service: {args.service}")
                print(f"Username: {entry['username']}")
                print(f"Password: {entry['password']}")
                if entry['notes']:
                    print(f"Notes: {entry['notes']}")
                print(f"Created: {entry['created']}")
                print(f"Modified: {entry['modified']}")
            else:
                print(f"No password found for '{args.service}'")
            vault.lock_vault()
    
    elif args.command == "list":
        if vault.unlock_vault():
            services = vault.list_services()
            if services:
                print("Stored services:")
                for service in sorted(services):
                    entry = vault.get_password(service)
                    print(f"  - {service} ({entry['username']})")
            else:
                print("No passwords stored in vault.")
            vault.lock_vault()
    
    elif args.command == "delete":
        if vault.unlock_vault():
            confirm = input(f"Delete password for '{args.service}'? (y/N): ")
            if confirm.lower() == 'y':
                vault.delete_password(args.service)
            vault.lock_vault()
    
    elif args.command == "change-master":
        if vault.unlock_vault():
            vault.change_master_password()
            vault.lock_vault()


# Example usage as a library
def example_library_usage():
    """Example of how to use the vault in your Python applications"""
    
    # Initialize vault
    vault = SecureVault("my_app_vault.enc")
    
    # Create vault (first time only)
    if not vault.vault_path.exists():
        vault.create_vault("your_secure_master_password")
    
    # Unlock vault
    if vault.unlock_vault("your_secure_master_password"):
        
        # Store a password
        vault.store_password("database", "admin", "super_secret_password", "Production DB")
        
        # Retrieve a password
        db_creds = vault.get_password("database")
        if db_creds:
            print(f"DB Username: {db_creds['username']}")
            print(f"DB Password: {db_creds['password']}")
        
        # Lock vault when done
        vault.lock_vault()


if __name__ == "__main__":
    try:
        main()
    except KeyboardInterrupt:
        print("\nAborted.")
        sys.exit(1)
    except PermissionError as e:
        print(f"ERROR: Permission denied accessing vault file -- {e}")
        sys.exit(1)
    except FileNotFoundError as e:
        print(f"ERROR: File not found -- {e}")
        sys.exit(1)
    except RuntimeError as e:
        print(f"ERROR: {e}")
        sys.exit(1)
    except Exception as e:
        print(f"ERROR: Unexpected error -- {e}")
        sys.exit(1)