#!/usr/bin/env python3
"""
Analyze the structure of an encrypted file to understand the padding issue.
"""

import os
import sys
from pathlib import Path

def analyze_encrypted_file():
    """Analyze the encrypted file structure."""
    
    print("=== Analyzing Encrypted File Structure ===")
    
    # Load environment
    env_local_path = Path(".env.local")
    if env_local_path.exists():
        with open(env_local_path, 'r') as f:
            for line in f:
                line = line.strip()
                if line and not line.startswith('#') and '=' in line:
                    key, value = line.split('=', 1)
                    os.environ[key.strip()] = value.strip()
    
    password = os.getenv("ENCRYPTION_PASSWORD")
    print(f"Password: {password[:16]}..." if password else "No password")
    
    # Find encrypted file
    enc_files = list(Path("agent_v0/sessions").glob("*.enc"))
    if not enc_files:
        print("No .enc files found")
        return False
    
    enc_file = enc_files[0]
    print(f"Analyzing: {enc_file}")
    print(f"File size: {enc_file.stat().st_size} bytes")
    
    # Read file structure
    with open(enc_file, 'rb') as f:
        salt = f.read(16)
        iv = f.read(16)
        ciphertext = f.read()
    
    print(f"Salt: {len(salt)} bytes")
    print(f"IV: {len(iv)} bytes")
    print(f"Ciphertext: {len(ciphertext)} bytes")
    print(f"Ciphertext % 16: {len(ciphertext) % 16}")
    
    if len(ciphertext) % 16 != 0:
        print("Ciphertext length is not a multiple of 16!")
        return False
    
    # Try manual decryption to see where it fails
    try:
        from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
        from cryptography.hazmat.backends import default_backend
        from cryptography.hazmat.primitives import hashes
        from cryptography.hazmat.primitives.kdf.pbkdf2 import PBKDF2HMAC
        
        # Derive key
        kdf = PBKDF2HMAC(
            algorithm=hashes.SHA256(),
            length=32,
            salt=salt,
            iterations=100000,
            backend=default_backend()
        )
        key = kdf.derive(password.encode('utf-8'))
        print("Key derivation successful")
        
        # Decrypt
        cipher = Cipher(
            algorithms.AES(key),
            modes.CBC(iv),
            backend=default_backend()
        )
        decryptor = cipher.decryptor()
        plaintext = decryptor.update(ciphertext) + decryptor.finalize()
        print(f"Decryption successful, plaintext length: {len(plaintext)}")
        
        # Check padding
        if len(plaintext) == 0:
            print("Plaintext is empty!")
            return False
        
        padding_length = plaintext[-1]
        print(f"Last byte (padding length): {padding_length}")
        
        if padding_length < 1 or padding_length > 16:
            print(f"Invalid padding length: {padding_length}")
            return False
        
        # Check padding bytes
        padding_bytes = plaintext[-padding_length:]
        print(f"Padding bytes: {[b for b in padding_bytes]}")
        
        all_correct = all(b == padding_length for b in padding_bytes)
        if not all_correct:
            print("Padding bytes are not all the same!")
            print(f"Expected all bytes to be {padding_length}")
            return False
        
        print("Padding validation successful")
        return True
        
    except Exception as e:
        print(f"Manual decryption failed: {e}")
        import traceback
        traceback.print_exc()
        return False

if __name__ == "__main__":
    success = analyze_encrypted_file()
    sys.exit(0 if success else 1)