"""
AnyDesk Blind Protocol Fuzzer
Hooks the RECEIVING side's TLS decryption and mutates incoming data
to trigger crashes in codec/protocol parsing.

Modes:
  --mode random     : Random byte flips in all incoming data
  --mode dimensions : Target likely dimension fields (2-byte and 4-byte values 100-8192)
  --mode overflow   : Replace small size values with large ones (integer overflow)
  --mode all        : Cycle through all mutation strategies

Run on the VICTIM side while connected to attacker.
The attacker just needs to move the mouse / show screen content to generate frames.

Usage: python 04_fuzzer.py --mode all --intensity medium
"""
import frida
import sys
import os
import json
import time
import random
import argparse
from datetime import datetime

FUZZER_AGENT = r"""
'use strict';

const FUZZ_MODE = '%FUZZ_MODE%';
const INTENSITY = %INTENSITY%; // 0.0-1.0, probability of mutating a packet
const SKIP_FIRST_N = %SKIP_FIRST_N%; // Skip first N packets (handshake)
let packetCount = 0;
let mutationCount = 0;
let crashDetected = false;

// Mutation strategies
const strategies = {
    // Flip random bytes
    random: function(arr, len) {
        const numFlips = Math.max(1, Math.floor(len * 0.01)); // 1% of bytes
        for (let i = 0; i < numFlips; i++) {
            const idx = Math.floor(Math.random() * len);
            arr[idx] = Math.floor(Math.random() * 256);
        }
        return `random_flip(${numFlips})`;
    },

    // Target dimension-like fields: find 2/4-byte LE values between 100-8192
    // and replace with edge cases
    dimensions: function(arr, len) {
        const edgeCases16 = [0, 1, 0x7FFF, 0x8000, 0xFFFF, 0xFFFE, 65535, 32768, 32767];
        const edgeCases32 = [0, 1, 0x7FFFFFFF, 0x80000000, 0xFFFFFFFF, 0xFFFFFFFE,
                             0x10000, 0xFFFF, 65536, 2147483647];
        let mutations = [];

        // Scan for 4-byte values that look like dimensions
        for (let i = 0; i < len - 4; i += 2) {
            const val = arr[i] | (arr[i+1] << 8) | (arr[i+2] << 16) | (arr[i+3] << 24);
            if (val >= 100 && val <= 8192) {
                // This might be a width/height — replace with edge case
                if (Math.random() < 0.3) {
                    const edge = edgeCases32[Math.floor(Math.random() * edgeCases32.length)];
                    arr[i] = edge & 0xFF;
                    arr[i+1] = (edge >> 8) & 0xFF;
                    arr[i+2] = (edge >> 16) & 0xFF;
                    arr[i+3] = (edge >> 24) & 0xFF;
                    mutations.push(`dim32@${i}:${val}->${edge}`);
                }
            }

            // Also check 2-byte values
            const val16 = arr[i] | (arr[i+1] << 8);
            if (val16 >= 100 && val16 <= 4096) {
                if (Math.random() < 0.2) {
                    const edge = edgeCases16[Math.floor(Math.random() * edgeCases16.length)];
                    arr[i] = edge & 0xFF;
                    arr[i+1] = (edge >> 8) & 0xFF;
                    mutations.push(`dim16@${i}:${val16}->${edge}`);
                }
            }
        }
        return mutations.length > 0 ? mutations.join(', ') : 'no_dims_found';
    },

    // Integer overflow: find size-like fields and make them huge
    overflow: function(arr, len) {
        let mutations = [];

        // Look for 4-byte values that could be sizes (reasonable range)
        for (let i = 0; i < len - 4; i += 4) {
            const val = arr[i] | (arr[i+1] << 8) | (arr[i+2] << 16) | (arr[i+3] << 24);

            // Sizes typically > 0 and < 10MB
            if (val > 0 && val < 10 * 1024 * 1024) {
                if (Math.random() < 0.15) {
                    // Integer overflow payloads
                    const overflows = [
                        0xFFFFFFFF,            // Max uint32
                        0x80000000,            // Int32 sign flip
                        val * 0x10001,         // width*height overflow pattern
                        0x7FFFFFFF,            // Max int32
                        val | 0xFF000000,      // High bytes set
                        (val << 16) | val,     // Doubled
                        0x01000000,            // 16MB (alloc stress)
                        0x10000000,            // 256MB
                        0xFFFFFFF0,            // Near-max aligned
                    ];
                    const ov = overflows[Math.floor(Math.random() * overflows.length)];
                    arr[i] = ov & 0xFF;
                    arr[i+1] = (ov >> 8) & 0xFF;
                    arr[i+2] = (ov >> 16) & 0xFF;
                    arr[i+3] = (ov >> 24) & 0xFF;
                    mutations.push(`overflow@${i}:${val}->0x${(ov >>> 0).toString(16)}`);
                }
            }
        }
        return mutations.length > 0 ? mutations.join(', ') : 'no_sizes_found';
    }
};

function mutatePacket(buf) {
    const arr = new Uint8Array(buf);
    const len = arr.length;

    if (len < 8) return null; // Too small to fuzz meaningfully

    let strategy;
    if (FUZZ_MODE === 'all') {
        const modes = ['random', 'dimensions', 'overflow'];
        strategy = modes[mutationCount % modes.length];
    } else {
        strategy = FUZZ_MODE;
    }

    const desc = strategies[strategy](arr, len);
    return {strategy: strategy, desc: desc};
}

// ============================================================
// Hook DecryptMessage (SChannel) — mutate data AFTER decryption
// ============================================================
function hookDecrypt() {
    const DecryptMessage = Module.findExportByName('sspicli.dll', 'DecryptMessage') ||
                           Module.findExportByName('secur32.dll', 'DecryptMessage');

    if (DecryptMessage) {
        Interceptor.attach(DecryptMessage, {
            onEnter(args) {
                this.pMessage = args[1];
            },
            onLeave(retval) {
                if (retval.toInt32() !== 0) return;
                packetCount++;

                if (packetCount <= SKIP_FIRST_N) return; // Skip handshake
                if (Math.random() > INTENSITY) return;   // Probabilistic

                try {
                    const pBufDesc = this.pMessage;
                    const cBuffers = pBufDesc.add(4).readU32();
                    const pBuffers = pBufDesc.add(8).readPointer();

                    for (let i = 0; i < cBuffers; i++) {
                        const bufPtr = pBuffers.add(i * 16);
                        const cbBuffer = bufPtr.readU32();
                        const bufType = bufPtr.add(4).readU32();
                        const pvBuffer = bufPtr.add(8).readPointer();

                        if (bufType === 1 && cbBuffer > 16 && cbBuffer < 1024 * 1024) {
                            const data = pvBuffer.readByteArray(cbBuffer);
                            const result = mutatePacket(data);

                            if (result) {
                                pvBuffer.writeByteArray(data);
                                mutationCount++;

                                if (mutationCount % 10 === 0 || mutationCount <= 5) {
                                    send({
                                        type: 'mutation',
                                        seq: packetCount,
                                        size: cbBuffer,
                                        strategy: result.strategy,
                                        desc: result.desc,
                                        totalMutations: mutationCount,
                                        timestamp: Date.now()
                                    });
                                }
                            }
                        }
                    }
                } catch(e) {
                    send({type: 'error', msg: e.toString()});
                }
            }
        });
        return true;
    }
    return false;
}

// ============================================================
// Hook SSL_read (OpenSSL) — mutate after read
// ============================================================
function hookSSLRead() {
    let hooked = false;
    const mods = Process.enumerateModules();

    for (const m of mods) {
        if (m.name.toLowerCase().includes('ssl')) {
            try {
                const exports = m.enumerateExports();
                for (const exp of exports) {
                    if (exp.name === 'SSL_read') {
                        Interceptor.attach(exp.address, {
                            onEnter(args) {
                                this.buf = args[1];
                            },
                            onLeave(retval) {
                                const read = retval.toInt32();
                                if (read <= 0) return;
                                packetCount++;

                                if (packetCount <= SKIP_FIRST_N) return;
                                if (Math.random() > INTENSITY) return;

                                const data = this.buf.readByteArray(read);
                                const result = mutatePacket(data);
                                if (result) {
                                    this.buf.writeByteArray(data);
                                    mutationCount++;
                                    if (mutationCount % 10 === 0 || mutationCount <= 5) {
                                        send({
                                            type: 'mutation',
                                            seq: packetCount,
                                            size: read,
                                            strategy: result.strategy,
                                            desc: result.desc,
                                            totalMutations: mutationCount,
                                            timestamp: Date.now()
                                        });
                                    }
                                }
                            }
                        });
                        hooked = true;
                    }
                }
            } catch(e) {}
        }
    }
    return hooked;
}

// ============================================================
// Fallback: Hook ws2_32!recv and mutate raw (may break TLS)
// ============================================================
function hookRawRecv() {
    const wsRecv = Module.findExportByName('ws2_32.dll', 'recv');
    if (wsRecv) {
        Interceptor.attach(wsRecv, {
            onEnter(args) {
                this.buf = args[1];
            },
            onLeave(retval) {
                const received = retval.toInt32();
                if (received <= 0) return;
                packetCount++;

                if (packetCount <= SKIP_FIRST_N) return;
                if (Math.random() > INTENSITY) return;

                const data = this.buf.readByteArray(received);
                const result = mutatePacket(data);
                if (result) {
                    this.buf.writeByteArray(data);
                    mutationCount++;
                    if (mutationCount % 50 === 0 || mutationCount <= 3) {
                        send({
                            type: 'mutation',
                            seq: packetCount,
                            size: received,
                            strategy: result.strategy,
                            desc: result.desc,
                            totalMutations: mutationCount,
                            note: 'RAW_SOCKET (may break TLS)',
                            timestamp: Date.now()
                        });
                    }
                }
            }
        });
        return true;
    }
    return false;
}

// ============================================================
// Crash detection: hook common crash paths
// ============================================================
function setupCrashDetection() {
    // Hook UnhandledExceptionFilter
    const uef = Module.findExportByName('kernel32.dll', 'UnhandledExceptionFilter');
    if (uef) {
        Interceptor.attach(uef, {
            onEnter(args) {
                const exceptionRecord = args[0];
                try {
                    const exceptionCode = exceptionRecord.readU32();
                    const exceptionAddress = exceptionRecord.add(Process.pointerSize * 2).readPointer();

                    send({
                        type: 'crash',
                        code: '0x' + (exceptionCode >>> 0).toString(16),
                        address: exceptionAddress.toString(),
                        totalPackets: packetCount,
                        totalMutations: mutationCount,
                        timestamp: Date.now()
                    });
                    crashDetected = true;
                } catch(e) {}
            }
        });
    }

    // Hook RtlReportException / RaiseException for earlier detection
    const raiseEx = Module.findExportByName('kernel32.dll', 'RaiseException');
    if (raiseEx) {
        Interceptor.attach(raiseEx, {
            onEnter(args) {
                const code = args[0].toInt32() >>> 0;
                // Filter: only report access violations, heap corruption, stack overflow
                if (code === 0xC0000005 || code === 0xC0000374 || code === 0xC00000FD ||
                    code === 0xC0000409) {
                    send({
                        type: 'exception',
                        code: '0x' + code.toString(16),
                        codeName: {
                            0xC0000005: 'ACCESS_VIOLATION',
                            0xC0000374: 'HEAP_CORRUPTION',
                            0xC00000FD: 'STACK_OVERFLOW',
                            0xC0000409: 'STACK_BUFFER_OVERRUN'
                        }[code] || 'UNKNOWN',
                        totalPackets: packetCount,
                        totalMutations: mutationCount,
                        stack: Thread.backtrace(this.context, Backtracer.ACCURATE)
                            .map(DebugSymbol.fromAddress).join('\n'),
                        timestamp: Date.now()
                    });
                }
            }
        });
    }

    send({type: 'status', msg: 'Crash detection hooks installed'});
}

// ============================================================
// Main
// ============================================================
send({type: 'status', msg: `Fuzzer starting (mode: ${FUZZ_MODE}, intensity: ${INTENSITY})`});
send({type: 'status', msg: `Skipping first ${SKIP_FIRST_N} packets (handshake protection)`});

setupCrashDetection();

let hooked = hookDecrypt();
if (!hooked) hooked = hookSSLRead();
if (!hooked) {
    send({type: 'status', msg: 'WARNING: No TLS hooks available. Falling back to raw socket (unstable).'});
    hookRawRecv();
}

send({type: 'status', msg: 'Fuzzer active. Make sure attacker screen is visible / moving...'});
send({type: 'ready'});

// Periodic stats
setInterval(function() {
    send({
        type: 'stats',
        packets: packetCount,
        mutations: mutationCount,
        crashDetected: crashDetected,
        timestamp: Date.now()
    });
}, 5000);
""";


def find_anydesk_pid():
    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


class FuzzerState:
    def __init__(self):
        self.ready = False
        self.mutations = 0
        self.packets = 0
        self.crashes = []
        self.exceptions = []
        self.start_time = time.time()

    def on_message(self, message, data):
        if message['type'] == 'send':
            payload = message['payload']
            msg_type = payload.get('type', '')

            if msg_type == 'status':
                print(f"  [*] {payload['msg']}")
            elif msg_type == 'ready':
                self.ready = True
            elif msg_type == 'mutation':
                self.mutations = payload.get('totalMutations', self.mutations + 1)
                print(f"  [FUZZ #{self.mutations}] pkt#{payload['seq']} "
                      f"{payload['size']}b {payload['strategy']}: {payload['desc']}")
            elif msg_type == 'stats':
                elapsed = time.time() - self.start_time
                self.packets = payload['packets']
                self.mutations = payload['mutations']
                rate = self.mutations / elapsed if elapsed > 0 else 0
                print(f"  [STATS] {self.packets} pkts, {self.mutations} mutations "
                      f"({rate:.1f}/s), {len(self.crashes)} crashes, {len(self.exceptions)} exceptions")
            elif msg_type == 'crash':
                self.crashes.append(payload)
                print(f"\n  {'!'*60}")
                print(f"  [CRASH] Exception 0x{payload['code']} at {payload['address']}")
                print(f"  [CRASH] After {payload['totalPackets']} packets, {payload['totalMutations']} mutations")
                print(f"  {'!'*60}\n")
            elif msg_type == 'exception':
                self.exceptions.append(payload)
                print(f"\n  [EXCEPTION] {payload['codeName']} (0x{payload['code']})")
                print(f"  After {payload['totalPackets']} packets, {payload['totalMutations']} mutations")
                if payload.get('stack'):
                    for frame in payload['stack'].split('\n')[:5]:
                        print(f"    {frame}")
                print()
            elif msg_type == 'error':
                print(f"  [!] {payload['msg']}")
        elif message['type'] == 'error':
            print(f"  [!] FRIDA ERROR: {message['description']}")

    def save_results(self, path):
        results = {
            'duration': time.time() - self.start_time,
            'total_packets': self.packets,
            'total_mutations': self.mutations,
            'crashes': self.crashes,
            'exceptions': self.exceptions
        }
        with open(path, 'w') as f:
            json.dump(results, f, indent=2)
        print(f"[+] Results saved to {path}")


def main():
    parser = argparse.ArgumentParser(description='AnyDesk Blind Protocol Fuzzer')
    parser.add_argument('--mode', choices=['random', 'dimensions', 'overflow', 'all'],
                       default='all', help='Mutation strategy')
    parser.add_argument('--intensity', choices=['low', 'medium', 'high', 'max'],
                       default='medium', help='Mutation probability per packet')
    parser.add_argument('--skip', type=int, default=50,
                       help='Skip first N packets (handshake protection)')
    args = parser.parse_args()

    intensity_map = {'low': 0.05, 'medium': 0.15, 'high': 0.4, 'max': 0.9}
    intensity = intensity_map[args.intensity]

    print("=" * 60)
    print("  AnyDesk Blind Protocol Fuzzer")
    print("=" * 60)
    print(f"  Mode: {args.mode}")
    print(f"  Intensity: {args.intensity} ({intensity*100:.0f}% of packets)")
    print(f"  Skip first: {args.skip} packets")
    print()

    pid = find_anydesk_pid()
    if not pid:
        print("[!] AnyDesk.exe not found.")
        sys.exit(1)

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

    state = FuzzerState()

    try:
        session = frida.attach(pid)
    except Exception as e:
        print(f"[!] Frida attach failed: {e}")
        print("[!] Run as Administrator.")
        sys.exit(1)

    agent_code = FUZZER_AGENT.replace('%FUZZ_MODE%', args.mode)
    agent_code = agent_code.replace('%INTENSITY%', str(intensity))
    agent_code = agent_code.replace('%SKIP_FIRST_N%', str(args.skip))

    script = session.create_script(agent_code, runtime='v8')
    script.on('message', state.on_message)
    script.load()

    for _ in range(100):
        if state.ready:
            break
        time.sleep(0.1)

    print(f"\n{'=' * 60}")
    print("  FUZZING — move mouse on attacker screen to generate frames")
    print("  Press Ctrl+C to stop and save results")
    print(f"{'=' * 60}\n")

    try:
        while True:
            time.sleep(0.1)
    except KeyboardInterrupt:
        print("\n\n[+] Stopping fuzzer...")

    timestamp = datetime.now().strftime('%Y%m%d_%H%M%S')
    state.save_results(f"fuzz_results_{timestamp}.json")

    if state.crashes:
        print(f"\n[!!!] {len(state.crashes)} CRASHES DETECTED — review results file")
    if state.exceptions:
        print(f"[!!!] {len(state.exceptions)} EXCEPTIONS DETECTED — review results file")

    session.detach()

if __name__ == '__main__':
    main()