"""
Dump AnyDesk's main module from memory — uses Frida's module info directly.
"""
import frida
import sys
import subprocess
import time

AGENT = r"""
var mainMod = Process.enumerateModules()[0];
send({t:'info', msg: 'Module: ' + mainMod.name + ' @ ' + mainMod.base + ' size: ' + mainMod.size});

// Check if the base has a PE header
var mz = mainMod.base.readU16();
send({t:'info', msg: 'MZ check: 0x' + mz.toString(16) + (mz === 0x5A4D ? ' (valid)' : ' (invalid)')});

if (mz === 0x5A4D) {
    var peOff = mainMod.base.add(0x3C).readU32();
    var peSig = mainMod.base.add(peOff).readU32();
    send({t:'info', msg: 'PE sig: 0x' + peSig.toString(16)});

    var numSections = mainMod.base.add(peOff + 6).readU16();
    var sizeOfImage = mainMod.base.add(peOff + 0x50).readU32();
    var sizeOfHeaders = mainMod.base.add(peOff + 0x54).readU32();
    send({t:'info', msg: 'Sections: ' + numSections + ', SizeOfImage: ' + sizeOfImage + ', HeaderSize: ' + sizeOfHeaders});

    // Read section table to find real extent
    var optHeaderSize = mainMod.base.add(peOff + 0x14).readU16();
    var sectionTableOff = peOff + 0x18 + optHeaderSize;
    var maxEnd = 0;
    for (var i = 0; i < numSections; i++) {
        var secBase = mainMod.base.add(sectionTableOff + i * 40);
        var nameBytes = secBase.readByteArray(8);
        var nameArr = new Uint8Array(nameBytes);
        var name = '';
        for (var j = 0; j < 8; j++) { if (nameArr[j] === 0) break; name += String.fromCharCode(nameArr[j]); }
        var virtualSize = secBase.add(8).readU32();
        var virtualAddr = secBase.add(12).readU32();
        var rawSize = secBase.add(16).readU32();
        var chars = secBase.add(36).readU32();
        var end = virtualAddr + virtualSize;
        if (end > maxEnd) maxEnd = end;
        send({t:'section', name: name.replace(/\0/g,''), va: '0x'+virtualAddr.toString(16),
              vs: virtualSize, rs: rawSize, chars: '0x'+chars.toString(16)});
    }
    send({t:'info', msg: 'Max section end: 0x' + maxEnd.toString(16) + ' (' + maxEnd + ' bytes)'});
}

// Dump the full module
var dumpSize = mainMod.size;
send({t:'info', msg: 'Dumping ' + dumpSize + ' bytes...'});

var chunkSize = 512 * 1024; // 512KB chunks
var offset = 0;
while (offset < dumpSize) {
    var readSize = Math.min(chunkSize, dumpSize - offset);
    try {
        var data = mainMod.base.add(offset).readByteArray(readSize);
        send({t:'chunk', offset: offset}, data);
    } catch(e) {
        send({t:'info', msg: 'Failed at offset 0x' + offset.toString(16) + ': ' + e});
        var zeros = new ArrayBuffer(readSize);
        send({t:'chunk', offset: offset}, zeros);
    }
    offset += readSize;
}
send({t:'done', size: dumpSize});
"""

def find_pid():
    out = subprocess.check_output(
        ['tasklist', '/FI', 'IMAGENAME eq AnyDesk.exe', '/FO', 'CSV', '/NH'],
        text=True, stderr=subprocess.DEVNULL)
    for line in out.strip().split('\n'):
        if 'AnyDesk' in line:
            parts = line.strip().strip('"').split('","')
            if len(parts) >= 2:
                return int(parts[1].strip('"'))
    return None

def main():
    print("=" * 60)
    print("  AnyDesk Memory Dumper v2")
    print("=" * 60)

    pid = find_pid()
    if not pid:
        print("[!] AnyDesk not running")
        sys.exit(1)

    print(f"[+] Attaching to PID {pid}")

    dump_data = bytearray()
    dump_size = [0]
    done = [False]

    def on_msg(msg, data):
        if msg['type'] == 'send':
            p = msg['payload']
            if p['t'] == 'info':
                print(f"  [*] {p['msg']}")
            elif p['t'] == 'section':
                print(f"  [SEC] {p['name']:8s} VA={p['va']} VSize={p['vs']:>10,} RawSize={p['rs']:>10,} Chars={p['chars']}")
            elif p['t'] == 'chunk' and data:
                offset = p['offset']
                if len(dump_data) < offset + len(data):
                    dump_data.extend(b'\x00' * (offset + len(data) - len(dump_data)))
                dump_data[offset:offset+len(data)] = data
                mb = (offset + len(data)) / 1024 / 1024
                print(f"  [DUMP] {mb:.1f} MB...", end='\r')
            elif p['t'] == 'done':
                dump_size[0] = p['size']
                done[0] = True
                print(f"\n  [+] Dump complete: {p['size']:,} bytes")
        elif msg['type'] == 'error':
            print(f"  [!] {msg['description']}")

    session = frida.attach(pid)
    script = session.create_script(AGENT, runtime='v8')
    script.on('message', on_msg)
    script.load()

    for _ in range(1200):  # 2 min timeout
        if done[0]: break
        time.sleep(0.1)

    script.unload()
    session.detach()

    if dump_data:
        outpath = f"anydesk_dump_{pid}.bin"
        with open(outpath, 'wb') as f:
            f.write(dump_data)
        print(f"\n[+] Saved: {outpath} ({len(dump_data):,} bytes)")
        print(f"[+] Load in Ghidra: mcp__ghidra__analyze_binary with this file")
    else:
        print("[!] No data dumped")

if __name__ == '__main__':
    main()