"""
AnyDesk Process Recon
Enumerates loaded modules, TLS-related exports, and interesting functions.
Run this first to understand AnyDesk's internals.
"""
import frida
import sys
import json
import time

AGENT_CODE = r"""
'use strict';

// ============================================================
// Phase 1: Enumerate all loaded modules
// ============================================================
function enumModules() {
    const mods = Process.enumerateModules();
    const result = [];
    for (const m of mods) {
        result.push({
            name: m.name,
            base: m.base.toString(),
            size: m.size,
            path: m.path
        });
    }
    return result;
}

// ============================================================
// Phase 2: Find TLS/crypto related exports across all modules
// ============================================================
function findTlsExports() {
    const patterns = [
        'SSL_read', 'SSL_write', 'SSL_new', 'SSL_connect',
        'SSL_accept', 'SSL_free', 'SSL_CTX_new',
        'EncryptMessage', 'DecryptMessage',
        'SslEncryptPacket', 'SslDecryptPacket',
        'BCryptEncrypt', 'BCryptDecrypt',
        'CryptEncrypt', 'CryptDecrypt',
        'mbedtls_ssl_read', 'mbedtls_ssl_write',
        // Winsock
        'send', 'recv', 'WSASend', 'WSARecv',
        'sendto', 'recvfrom',
        // File operations (for file transfer)
        'CreateFileW', 'CreateFileA', 'WriteFile', 'ReadFile',
        'MoveFileW', 'MoveFileExW', 'CopyFileW',
        'CreateDirectoryW', 'RemoveDirectoryW',
        // Memory alloc (for heap overflow detection)
        'HeapAlloc', 'HeapFree', 'HeapReAlloc',
        'VirtualAlloc', 'VirtualFree',
        'malloc', 'free', 'realloc', 'calloc',
        // Image/bitmap
        'CreateDIBSection', 'SetDIBits', 'GetDIBits',
        'CreateBitmap', 'CreateCompatibleBitmap',
        // Clipboard
        'SetClipboardData', 'GetClipboardData',
        'OpenClipboard', 'CloseClipboard',
        'EmptyClipboard',
    ];

    const found = [];
    const mods = Process.enumerateModules();

    for (const m of mods) {
        try {
            const exports = m.enumerateExports();
            for (const exp of exports) {
                for (const pat of patterns) {
                    if (exp.name && exp.name.includes(pat)) {
                        found.push({
                            module: m.name,
                            name: exp.name,
                            address: exp.address.toString(),
                            type: exp.type
                        });
                    }
                }
            }
        } catch (e) {
            // Some modules can't be enumerated
        }
    }
    return found;
}

// ============================================================
// Phase 3: Look for AnyDesk-specific strings in memory
// ============================================================
function findInterestingStrings() {
    const mainModule = Process.enumerateModules()[0]; // AnyDesk.exe
    const found = [];
    const patterns = [
        'DeskRT',
        'file_transfer',
        'FileTransfer',
        'clipboard',
        'Clipboard',
        'identity',
        'Identity',
        'user_image',
        'UserImage',
        'avatar',
        'codec',
        'Codec',
        'decompress',
        'Decompress',
        'decode_frame',
        'DecodeFrame',
        'bitmap',
        'Bitmap',
        'PNG',
        'png_decode',
        'traversal',
        'path_sanitize',
        'sanitize',
        'validate_path',
        'ad.security',
        'ssl_ctx',
        'TLS',
        'handshake',
        'Handshake',
    ];

    for (const pat of patterns) {
        try {
            const matches = Memory.scanSync(mainModule.base, mainModule.size,
                stringToPattern(pat));
            if (matches.length > 0) {
                found.push({
                    pattern: pat,
                    count: matches.length,
                    first_addr: matches[0].address.toString(),
                    // Read surrounding context
                    context: safeReadUtf8(matches[0].address.sub(16), 64)
                });
            }
        } catch (e) {
            // Scan failed for this pattern
        }
    }
    return found;
}

function stringToPattern(str) {
    let pat = '';
    for (let i = 0; i < str.length; i++) {
        if (i > 0) pat += ' ';
        pat += str.charCodeAt(i).toString(16).padStart(2, '0');
    }
    return pat;
}

function safeReadUtf8(addr, len) {
    try {
        const buf = addr.readByteArray(len);
        if (!buf) return '<null>';
        const arr = new Uint8Array(buf);
        let s = '';
        for (let i = 0; i < arr.length; i++) {
            const c = arr[i];
            if (c >= 0x20 && c < 0x7f) s += String.fromCharCode(c);
            else s += '.';
        }
        return s;
    } catch (e) {
        return '<unreadable>';
    }
}

// ============================================================
// Phase 4: Scan for OpenSSL-style function tables
// ============================================================
function findOpenSSLIndicators() {
    const mods = Process.enumerateModules();
    const indicators = [];

    for (const m of mods) {
        const nameLower = m.name.toLowerCase();
        if (nameLower.includes('ssl') || nameLower.includes('crypto') ||
            nameLower.includes('tls') || nameLower.includes('openssl') ||
            nameLower.includes('mbedtls') || nameLower.includes('boringssl') ||
            nameLower.includes('libcrypto') || nameLower.includes('libssl') ||
            nameLower.includes('schannel') || nameLower.includes('ncrypt') ||
            nameLower.includes('bcrypt') || nameLower.includes('sspicli')) {
            indicators.push({
                name: m.name,
                base: m.base.toString(),
                size: m.size,
                path: m.path
            });
        }
    }
    return indicators;
}

// ============================================================
// Run all phases
// ============================================================
const report = {};

send({type: 'status', msg: 'Phase 1: Enumerating modules...'});
report.modules = enumModules();
send({type: 'status', msg: `Found ${report.modules.length} modules`});

send({type: 'status', msg: 'Phase 2: Finding TLS/crypto exports...'});
report.tls_exports = findTlsExports();
send({type: 'status', msg: `Found ${report.tls_exports.length} interesting exports`});

send({type: 'status', msg: 'Phase 3: Scanning for interesting strings...'});
report.strings = findInterestingStrings();
send({type: 'status', msg: `Found ${report.strings.length} string patterns`});

send({type: 'status', msg: 'Phase 4: Looking for TLS library indicators...'});
report.tls_indicators = findOpenSSLIndicators();
send({type: 'status', msg: `Found ${report.tls_indicators.length} TLS-related modules`});

send({type: 'result', data: report});
"""

def on_message(message, data):
    if message['type'] == 'send':
        payload = message['payload']
        if payload.get('type') == 'status':
            print(f"  [*] {payload['msg']}")
        elif payload.get('type') == 'result':
            global result_data
            result_data = payload['data']
    elif message['type'] == 'error':
        print(f"  [!] ERROR: {message['description']}")

def find_anydesk_pid():
    """Find AnyDesk process ID"""
    import subprocess
    try:
        output = subprocess.check_output(
            ['tasklist', '/FI', 'IMAGENAME eq AnyDesk.exe', '/FO', 'CSV', '/NH'],
            text=True, stderr=subprocess.DEVNULL
        )
        for line in output.strip().split('\n'):
            if 'AnyDesk' in line:
                parts = line.strip().strip('"').split('","')
                if len(parts) >= 2:
                    return int(parts[1].strip('"'))
    except:
        pass
    return None

def main():
    global result_data
    result_data = None

    print("=" * 60)
    print("  AnyDesk Process Recon")
    print("=" * 60)

    pid = find_anydesk_pid()
    if not pid:
        print("[!] AnyDesk.exe not found. Make sure it's running.")
        sys.exit(1)

    print(f"[+] Found AnyDesk.exe (PID: {pid})")
    print(f"[+] Attaching Frida...")

    try:
        session = frida.attach(pid)
    except frida.ProcessNotFoundError:
        print("[!] Could not attach to AnyDesk. Run as Administrator.")
        sys.exit(1)
    except Exception as e:
        print(f"[!] Frida attach failed: {e}")
        print("[!] Make sure to run as Administrator.")
        sys.exit(1)

    print("[+] Injecting recon agent...")
    script = session.create_script(AGENT_CODE, runtime='v8')
    script.on('message', on_message)
    script.load()

    # Wait for results
    timeout = 30
    for i in range(timeout * 10):
        if result_data is not None:
            break
        time.sleep(0.1)

    if result_data is None:
        print("[!] Timed out waiting for results")
        session.detach()
        sys.exit(1)

    # Save and display results
    output_path = "recon_results.json"
    with open(output_path, 'w') as f:
        json.dump(result_data, f, indent=2)

    print(f"\n[+] Full results saved to {output_path}")

    # Summary
    print(f"\n{'=' * 60}")
    print("  SUMMARY")
    print(f"{'=' * 60}")

    print(f"\n[Modules] {len(result_data['modules'])} loaded")
    for m in result_data['modules'][:10]:
        print(f"  {m['name']:30s} @ {m['base']}  ({m['size']:>10,} bytes)")
    if len(result_data['modules']) > 10:
        print(f"  ... and {len(result_data['modules'])-10} more (see JSON)")

    print(f"\n[TLS/Crypto Related Modules]")
    if result_data['tls_indicators']:
        for m in result_data['tls_indicators']:
            print(f"  {m['name']:30s} @ {m['base']}  {m['path']}")
    else:
        print("  None found (AnyDesk may use statically linked TLS)")

    print(f"\n[Interesting Exports] {len(result_data['tls_exports'])} found")
    # Group by category
    tls_funcs = [e for e in result_data['tls_exports']
                 if any(k in e['name'] for k in ['SSL_', 'Encrypt', 'Decrypt', 'mbedtls', 'Crypt'])]
    net_funcs = [e for e in result_data['tls_exports']
                 if any(k in e['name'] for k in ['send', 'recv', 'WSA'])]
    file_funcs = [e for e in result_data['tls_exports']
                  if any(k in e['name'] for k in ['CreateFile', 'WriteFile', 'ReadFile', 'MoveFile', 'CopyFile', 'Directory'])]
    clip_funcs = [e for e in result_data['tls_exports']
                  if 'Clipboard' in e['name'] or 'clipboard' in e['name']]
    img_funcs = [e for e in result_data['tls_exports']
                 if any(k in e['name'] for k in ['DIB', 'Bitmap', 'bitmap'])]

    if tls_funcs:
        print(f"\n  TLS/Crypto ({len(tls_funcs)}):")
        for e in tls_funcs[:15]:
            print(f"    {e['module']:20s} :: {e['name']:40s} @ {e['address']}")

    if net_funcs:
        print(f"\n  Network ({len(net_funcs)}):")
        for e in net_funcs[:10]:
            print(f"    {e['module']:20s} :: {e['name']:40s} @ {e['address']}")

    if file_funcs:
        print(f"\n  File I/O ({len(file_funcs)}):")
        for e in file_funcs[:10]:
            print(f"    {e['module']:20s} :: {e['name']:40s} @ {e['address']}")

    if clip_funcs:
        print(f"\n  Clipboard ({len(clip_funcs)}):")
        for e in clip_funcs:
            print(f"    {e['module']:20s} :: {e['name']:40s} @ {e['address']}")

    if img_funcs:
        print(f"\n  Image/Bitmap ({len(img_funcs)}):")
        for e in img_funcs:
            print(f"    {e['module']:20s} :: {e['name']:40s} @ {e['address']}")

    print(f"\n[String Patterns in AnyDesk.exe]")
    if result_data['strings']:
        for s in result_data['strings']:
            print(f"  '{s['pattern']}' — {s['count']} hits, first @ {s['first_addr']}")
            if s['context']:
                print(f"    context: {s['context']}")
    else:
        print("  No patterns found in main module")

    session.detach()
    print(f"\n[+] Done. Review {output_path} for full details.")
    print("[+] Next step: run 02_sniffer.py to capture protocol traffic")

if __name__ == '__main__':
    main()