encryption.py

from cryptography.fernet import Fernet, InvalidToken
from cryptography.hazmat.primitives import hashes, hmac
from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
import base64
import json
import os
import random
import string
import secrets
from config import SecurityConfig

# Constants
KDF_ITERATIONS = 480000
VAULT_KDF_ITERATIONS = 100000
MIN_FAKE_PASSWORDS = 5
MAX_FAKE_PASSWORDS = 15
MIN_PASSWORD_LENGTH = 8
MAX_PASSWORD_LENGTH = 16
PEPPER = b'S3cur3P3pp3r2024'
STORAGE_KEY = b'Tr1v1alK3y2024'

def simple_encrypt(data: bytes, key: bytes) -> bytes:
    """Simple XOR encryption for .enc file obfuscation."""
    return bytes(b ^ key[i % len(key)] for i, b in enumerate(data))

def simple_decrypt(data: bytes, key: bytes) -> bytes:
    """Simple XOR decryption for .enc file obfuscation."""
    return simple_encrypt(data, key)  # XOR is its own inverse

class EncryptionManager:
    def __init__(self):
        self.fernet = None
        self.vault_id = None
        self.vault_id_hash = None
        self.encryption_salt = None
        self.vault_salt = None

    def setup_encryption(self, master_password: str) -> None:
        """Set up encryption with the master password."""
        if not master_password:
            raise ValueError("Master password cannot be empty")
        
        # Generate vault ID using only pepper for consistency
        vault_kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=16,
            salt=PEPPER,
            iterations=VAULT_KDF_ITERATIONS,
        )
        self.vault_id = base64.urlsafe_b64encode(
            vault_kdf.derive(master_password.encode())
        ).decode('utf-8')
        
        digest = hashes.Hash(hashes.SHA256())
        digest.update(self.vault_id.encode())
        self.vault_id_hash = digest.finalize().hex()
        
        # Try to load existing vault to get its salts
        if os.path.exists('passwords.enc'):
            with open('passwords.enc', 'rb') as f:
                encrypted_data = f.read()
                decrypted_data = simple_decrypt(encrypted_data, STORAGE_KEY)
                vaults = json.loads(decrypted_data)
                
                # Look for existing vault
                for vault in vaults:
                    if vault.get('id_hash') == self.vault_id_hash:
                        self.encryption_salt = base64.b64decode(vault['encryption_salt'])
                        self.vault_salt = base64.b64decode(vault['vault_salt'])
                        break
                else:
                    # No existing vault found, generate new salts
                    self.encryption_salt = secrets.token_bytes(16)
                    self.vault_salt = secrets.token_bytes(16)
        else:
            # No vaults file exists, generate new salts
            self.encryption_salt = secrets.token_bytes(16)
            self.vault_salt = secrets.token_bytes(16)
        
        # Derive encryption key with salt and pepper for real passwords
        params = SecurityConfig.get_key_derivation_params()
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=self.encryption_salt + PEPPER,
            iterations=KDF_ITERATIONS,
        )
        key = base64.urlsafe_b64encode(kdf.derive(master_password.encode()))
        self.fernet = Fernet(key)

    def generate_fake_passwords(self):
        h = hmac.HMAC(self.vault_id.encode(), hashes.SHA256())
        h.update(b'fake_password_seed')
        seed = int.from_bytes(h.finalize(), 'big')
        random.seed(seed)
        
        num_passwords = random.randint(MIN_FAKE_PASSWORDS, MAX_FAKE_PASSWORDS)
        fake_passwords = []
        for _ in range(num_passwords):
            length = random.randint(MIN_PASSWORD_LENGTH, MAX_PASSWORD_LENGTH)
            chars = string.ascii_letters + string.digits + string.punctuation
            password = ''.join(random.choice(chars) for _ in range(length))
            fake_passwords.append(password)
        return fake_passwords

    def load_passwords(self):
        try:
            if os.path.exists('passwords.enc'):
                with open('passwords.enc', 'rb') as f:
                    encrypted_data = f.read()
                    decrypted_data = simple_decrypt(encrypted_data, STORAGE_KEY)
                    vaults = json.loads(decrypted_data)
            else:
                return self.generate_fake_passwords()
            
            for vault in vaults:
                if vault.get('id_hash') == self.vault_id_hash:
                    try:
                        self.encryption_salt = base64.b64decode(vault['encryption_salt'])
                        self.vault_salt = base64.b64decode(vault['vault_salt'])
                        
                        decrypted_data = self.fernet.decrypt(vault['data'].encode())
                        return json.loads(decrypted_data)
                    except InvalidToken:
                        return self.generate_fake_passwords()
            
            return self.generate_fake_passwords()
        
        except Exception as e:
            print(f"Error loading passwords: {e}")
            return []

    def save_passwords(self, passwords):
        try:
            if os.path.exists('passwords.enc'):
                with open('passwords.enc', 'rb') as f:
                    encrypted_data = f.read()
                    decrypted_data = simple_decrypt(encrypted_data, STORAGE_KEY)
                    vaults = json.loads(decrypted_data)
            else:
                vaults = []
            
            encrypted_data = self.fernet.encrypt(json.dumps(passwords).encode()).decode('utf-8')
            
            new_vault = {
                'id_hash': self.vault_id_hash,
                'encryption_salt': base64.b64encode(self.encryption_salt).decode('utf-8'),
                'vault_salt': base64.b64encode(self.vault_salt).decode('utf-8'),
                'data': encrypted_data
            }
            
            # Update or add vault
            for i, vault in enumerate(vaults):
                if vault.get('id_hash') == self.vault_id_hash:
                    vaults[i] = new_vault
                    break
            else:
                vaults.append(new_vault)
            
            final_data = json.dumps(vaults).encode()
            encrypted_storage = simple_encrypt(final_data, STORAGE_KEY)
            with open('passwords.enc', 'wb') as f:
                f.write(encrypted_storage)
                
        except Exception as e:
            print(f"Error saving passwords: {e}")