"""
AnyDesk Sniffer v4 — hooks ALL AnyDesk processes
"""
import frida
import sys
import os
import json
import time
import subprocess
from datetime import datetime

CAPTURE_DIR = "captures"

AGENT_CODE = r"""
var hookCount = 0;

function getExport(mod, name) {
    try {
        var m = Process.findModuleByName(mod);
        if (!m) return null;
        return m.findExportByName(name);
    } catch(e) { return null; }
}

function analyzeBuffer(arr, len) {
    var tags = [];
    var run = 0, maxRun = 0, bestStr = '', curStr = '';
    for (var i = 0; i < len; i++) {
        var b = arr[i];
        if (b >= 0x20 && b < 0x7f) { run++; curStr += String.fromCharCode(b); }
        else { if (run > maxRun) { maxRun = run; bestStr = curStr; } run = 0; curStr = ''; }
    }
    if (run > maxRun) { maxRun = run; bestStr = curStr; }
    if (maxRun > 10) tags.push('STR:' + bestStr.substring(0, 80));
    if (len >= 8 && arr[0]===0x89 && arr[1]===0x50 && arr[2]===0x4E && arr[3]===0x47) tags.push('PNG');
    if (len >= 2 && arr[0]===0x42 && arr[1]===0x4D) tags.push('BMP');
    return tags;
}

// ---- FILE I/O ----
var p1 = getExport('KERNEL32.DLL', 'CreateFileW');
if (p1) { Interceptor.attach(p1, { onEnter: function(args) { try {
    var path = args[0].readUtf16String(); if (!path) return;
    var pl = path.toLowerCase();
    if (pl.indexOf('\\device\\')!==-1||pl.indexOf('\\pipe\\')!==-1||pl.indexOf('condrv')!==-1||pl.indexOf('\\registry')!==-1) return;
    var access = args[1].toInt32()>>>0, disp = args[4].toInt32();
    var isW = (access&0x40000000)!==0||(access&0x2)!==0, isC = disp===1||disp===2||disp===4;
    if (isW||isC||pl.indexOf('anydesk')!==-1||pl.indexOf('..')!==-1||pl.indexOf('desktop')!==-1||pl.indexOf('download')!==-1||pl.indexOf('startup')!==-1) {
        var fl=[]; if(path.indexOf('..\\')!==-1)fl.push('TRAV_BS'); if(path.indexOf('../')!==-1)fl.push('TRAV_FS'); if(pl.indexOf('startup')!==-1)fl.push('STARTUP');
        send({t:'file',op:'Create',path:path,w:isW,c:isC,fl:fl,ts:Date.now()});
    }} catch(e){} }}); hookCount++; }

var p2 = getExport('KERNEL32.DLL', 'MoveFileExW');
if(p2){Interceptor.attach(p2,{onEnter:function(args){try{send({t:'move',src:args[0].readUtf16String(),dst:args[1].readUtf16String(),ts:Date.now()});}catch(e){}}});hookCount++;}

var p3 = getExport('KERNEL32.DLL', 'CopyFileW');
if(p3){Interceptor.attach(p3,{onEnter:function(args){try{send({t:'copy',src:args[0].readUtf16String(),dst:args[1].readUtf16String(),ts:Date.now()});}catch(e){}}});hookCount++;}

var p4 = getExport('KERNEL32.DLL', 'CreateDirectoryW');
if(p4){Interceptor.attach(p4,{onEnter:function(args){try{send({t:'mkdir',path:args[0].readUtf16String(),ts:Date.now()});}catch(e){}}});hookCount++;}

// ---- WriteFile ----
var p5 = getExport('KERNEL32.DLL', 'WriteFile');
if(p5){ var wc=0; Interceptor.attach(p5,{onEnter:function(args){
    this.sz=args[2].toInt32(); if(this.sz>64&&wc<300){wc++;try{
    var preview=args[1].readByteArray(Math.min(this.sz,64));var arr=new Uint8Array(preview);
    var tags=analyzeBuffer(arr,arr.length);if(tags.length>0)send({t:'write',sz:this.sz,tags:tags,ts:Date.now()});
    }catch(e){}}}}); hookCount++; }

// ---- GDI Bitmaps ----
var p6 = getExport('GDI32.dll', 'CreateDIBSection');
if(p6){Interceptor.attach(p6,{onEnter:function(args){try{var pbmi=args[1];if(pbmi.isNull())return;
    var w=pbmi.add(4).readS32(),h=pbmi.add(8).readS32(),bpp=pbmi.add(14).readU16();
    var alloc=Math.abs(w)*Math.abs(h)*(bpp/8);
    send({t:'bmp',op:'DIB',w:w,h:h,bpp:bpp,alloc:alloc,ts:Date.now()});
    if(Math.abs(w)>8192||Math.abs(h)>8192)send({t:'alert',m:'HUGE BMP '+w+'x'+h+'@'+bpp});
}catch(e){}}});hookCount++;}

var p7 = getExport('GDI32.dll', 'CreateCompatibleBitmap');
if(p7){Interceptor.attach(p7,{onEnter:function(args){var cx=args[1].toInt32(),cy=args[2].toInt32();
    if(cx>0&&cy>0)send({t:'bmp',op:'Compat',w:cx,h:cy,bpp:0,alloc:0,ts:Date.now()});}});hookCount++;}

// ---- Clipboard ----
var p8=getExport('USER32.dll','SetClipboardData'),p9=getExport('USER32.dll','GetClipboardData');
var cfn={1:'TEXT',2:'BITMAP',7:'OEM',8:'DIB',13:'UNICODE',15:'HDROP',17:'DIBV5'};
if(p8){Interceptor.attach(p8,{onEnter:function(args){var f=args[0].toInt32();send({t:'clip',op:'SET',fmt:f,name:cfn[f]||('C'+f),ts:Date.now()});}});hookCount++;}
if(p9){Interceptor.attach(p9,{onEnter:function(args){this.f=args[0].toInt32();},onLeave:function(ret){if(!ret.isNull())send({t:'clip',op:'GET',fmt:this.f,name:cfn[this.f]||('C'+this.f),ts:Date.now()});}});hookCount++;}

// ---- ALL socket functions ----
var sc=0,rc=0,wsc=0,wrc=0;
var pS=getExport('WS2_32.dll','send');
if(pS){Interceptor.attach(pS,{onLeave:function(ret){var n=ret.toInt32();if(n>0){sc++;if(sc<=30||sc%100===0)send({t:'net',d:'S',sz:n,n:sc,ts:Date.now()});}}});hookCount++;}
var pR=getExport('WS2_32.dll','recv');
if(pR){Interceptor.attach(pR,{onLeave:function(ret){var n=ret.toInt32();if(n>0){rc++;if(rc<=30||rc%100===0)send({t:'net',d:'R',sz:n,n:rc,ts:Date.now()});}}});hookCount++;}
var pWS=getExport('WS2_32.dll','WSASend');
if(pWS){Interceptor.attach(pWS,{onEnter:function(args){try{var nBufs=args[2].toInt32(),lpBufs=args[1],total=0;
    for(var i=0;i<nBufs;i++)total+=lpBufs.add(i*8).readU32();wsc++;
    if(wsc<=30||wsc%100===0)send({t:'net',d:'WS',sz:total,n:wsc,ts:Date.now()});}catch(e){}}});hookCount++;}
var pWR=getExport('WS2_32.dll','WSARecv');
if(pWR){Interceptor.attach(pWR,{onEnter:function(args){try{var nBufs=args[2].toInt32(),lpBufs=args[1],total=0;
    for(var i=0;i<nBufs;i++)total+=lpBufs.add(i*8).readU32();wrc++;
    if(wrc<=30||wrc%100===0)send({t:'net',d:'WR',sz:total,n:wrc,ts:Date.now()});}catch(e){}}});hookCount++;}

// ---- Large heap alloc ----
var pH=getExport('ntdll.dll','RtlAllocateHeap');
if(pH){Interceptor.attach(pH,{onEnter:function(args){this.sz=args[2].toInt32()>>>0;},onLeave:function(ret){
    if(this.sz>50*1024*1024)send({t:'alert',m:'BIG ALLOC '+(this.sz/1024/1024).toFixed(1)+'MB',ts:Date.now()});}});hookCount++;}

setInterval(function(){send({t:'stats',s:sc,r:rc,ws:wsc,wr:wrc,ts:Date.now()});},10000);
send({t:'log',m:hookCount+' hooks OK'});
send({t:'ready'});
"""

class Capture:
    def __init__(self, d):
        self.d = d; os.makedirs(d, exist_ok=True)
        self.ev = []; self.ready_count = 0; self.t0 = time.time()

    def on_msg(self, pid, msg, data):
        if msg['type'] == 'send':
            p = msg['payload']
            t = p.get('t','')
            ts = (p.get('ts', self.t0*1000)/1000) - self.t0
            tag = f"P{pid}"

            if t == 'log': print(f"  [{tag}] {p['m']}")
            elif t == 'ready': self.ready_count += 1
            elif t == 'file':
                fl = ' '.join(f'[{f}]' for f in p.get('fl',[]))
                m = 'W' if p.get('w') else 'R'; c = '+C' if p.get('c') else ''
                print(f"  [{ts:7.2f}s] [{tag}] [FILE {m}{c}] {p['path']} {fl}")
                self.ev.append({**p, 'pid': pid})
            elif t == 'move':
                print(f"  [{ts:7.2f}s] [{tag}] [MOVE] {p['src']} -> {p['dst']}")
                self.ev.append({**p, 'pid': pid})
            elif t == 'copy':
                print(f"  [{ts:7.2f}s] [{tag}] [COPY] {p['src']} -> {p['dst']}")
                self.ev.append({**p, 'pid': pid})
            elif t == 'mkdir':
                print(f"  [{ts:7.2f}s] [{tag}] [MKDIR] {p['path']}")
                self.ev.append({**p, 'pid': pid})
            elif t == 'write':
                tg = ', '.join(p.get('tags',[]))
                if tg: print(f"  [{ts:7.2f}s] [{tag}] [WRITE] {p['sz']}b {tg}")
                self.ev.append({**p, 'pid': pid})
            elif t == 'bmp':
                print(f"  [{ts:7.2f}s] [{tag}] [BITMAP] {p['op']} {p['w']}x{p['h']} @{p['bpp']}bpp ({p['alloc']:,.0f}b)")
                self.ev.append({**p, 'pid': pid})
            elif t == 'clip':
                print(f"  [{ts:7.2f}s] [{tag}] [CLIP] {p['op']} {p['name']}")
                self.ev.append({**p, 'pid': pid})
            elif t == 'net':
                print(f"  [{ts:7.2f}s] [{tag}] [NET] {p['d']} {p['sz']}b #{p['n']}")
                self.ev.append({**p, 'pid': pid})
            elif t == 'alert':
                print(f"\n  {'!'*50}\n  [{tag}] ALERT: {p['m']}\n  {'!'*50}\n")
                self.ev.append({**p, 'pid': pid})
            elif t == 'stats':
                total = p['s'] + p['r'] + p.get('ws',0) + p.get('wr',0)
                if total > 0:  # Only print stats if there's activity
                    print(f"  [{ts:7.2f}s] [{tag}] send={p['s']} recv={p['r']} WSASend={p.get('ws',0)} WSARecv={p.get('wr',0)}")
        elif msg['type'] == 'error':
            print(f"  [P{pid}] ERROR: {msg['description']}")

    def save(self):
        path = os.path.join(self.d, "capture.json")
        with open(path, 'w') as f:
            json.dump({'dur': time.time()-self.t0, 'events': self.ev}, f, indent=2)
        fi = [e for e in self.ev if e.get('t')=='file']
        bm = [e for e in self.ev if e.get('t')=='bmp']
        cl = [e for e in self.ev if e.get('t')=='clip']
        nt = [e for e in self.ev if e.get('t')=='net']
        al = [e for e in self.ev if e.get('t')=='alert']
        print(f"\n[+] {len(self.ev)} events -> {path}")
        print(f"    Files:{len(fi)} Bitmaps:{len(bm)} Clip:{len(cl)} Net:{len(nt)} Alerts:{len(al)}")


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


def main():
    print("=" * 60)
    print("  AnyDesk Sniffer v4 — ALL processes")
    print("=" * 60)

    pids = find_all_pids()
    if not pids:
        print("[!] No AnyDesk processes found")
        sys.exit(1)

    print(f"[+] Found {len(pids)} AnyDesk processes: {pids}")

    ts = datetime.now().strftime('%Y%m%d_%H%M%S')
    cap = Capture(os.path.join(CAPTURE_DIR, f"cap_{ts}"))

    sessions = []
    scripts = []

    for pid in pids:
        try:
            print(f"[+] Attaching to PID {pid}...")
            session = frida.attach(pid)
            script = session.create_script(AGENT_CODE, runtime='v8')
            script.on('message', lambda msg, data, p=pid: cap.on_msg(p, msg, data))
            script.load()
            sessions.append(session)
            scripts.append(script)
            print(f"[+] PID {pid} hooked")
        except Exception as e:
            print(f"[!] PID {pid} failed: {e}")

    # Wait for all to be ready
    for _ in range(100):
        if cap.ready_count >= len(scripts):
            break
        time.sleep(0.1)

    print(f"\n{'='*60}")
    print(f"  {cap.ready_count}/{len(scripts)} agents ready")
    print("  NOW: connect from VM, transfer files, move mouse")
    print("  Ctrl+C to stop")
    print(f"{'='*60}\n")

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

    cap.save()
    for s in sessions:
        try:
            s.detach()
        except:
            pass

if __name__ == '__main__':
    main()