#!/usr/bin/env python3
import base64
import json
import os
import sys
from scapy.all import rdpcap, DNS, DNSQR, TCP, Raw

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))

PCAP_CANDIDATES = [
    os.path.join(SCRIPT_DIR, "..", "public", "pinecraft_dns_tunnel.pcap"),
]

pcap_path = None
for candidate in PCAP_CANDIDATES:
    if os.path.exists(candidate):
        pcap_path = candidate
        break

if pcap_path is None:
    print(f"[!] PCAP not found in {PCAP_CANDIDATES}", file=sys.stderr)
    sys.exit(1)

packets = rdpcap(pcap_path)

# Flag 1: C2 domain from suspicious DNS queries

c2_domain = None
dns_queries = set()

for pkt in packets:
    if DNS in pkt and pkt[DNS].qr == 0 and DNSQR in pkt:
        qname = pkt[DNSQR].qname.decode(errors="ignore").rstrip(".")
        dns_queries.add(qname)

for qname in sorted(dns_queries):
    if "evil" in qname.lower() and ".data." not in qname:
        c2_domain = qname
        break

if c2_domain:
    print(f"Flag 1: vsosh{{{c2_domain}}}")
else:
    print("[!] Flag 1 not found — no C2 domain with 'evil' in DNS queries", file=sys.stderr)

# Flag 2: Hex-encoded data from DNS exfiltration subdomains

data_suffix = f".data.{c2_domain}" if c2_domain else ".data.mc-textures.evil.net"
hex_chunks = []

for pkt in packets:
    if DNS in pkt and pkt[DNS].qr == 0 and DNSQR in pkt:
        qname = pkt[DNSQR].qname.decode(errors="ignore").rstrip(".")
        if qname.endswith(data_suffix):
            chunk = qname[: -(len(data_suffix))].rstrip(".")
            hex_chunks.append(chunk)

if hex_chunks:
    hex_str = "".join(hex_chunks)
    try:
        decoded = bytes.fromhex(hex_str).decode()
        print(f"Flag 2: vsosh{{{decoded}}}")
    except (ValueError, UnicodeDecodeError) as e:
        print(f"[!] Flag 2 hex decode error: {e}, raw hex: {hex_str}", file=sys.stderr)
else:
    print("[!] Flag 2 not found — no DNS exfil queries to *.data.* domain", file=sys.stderr)

# Flag 3: Base64-encoded JSON from HTTP POST /collect

flag3_found = False
for pkt in packets:
    if TCP in pkt and Raw in pkt:
        payload = bytes(pkt[Raw]).decode("utf-8", errors="ignore")
        if "POST /collect" in payload:
            parts = payload.split("\r\n\r\n", 1)
            if len(parts) > 1:
                body = parts[1].strip().split("\x00")[0].strip()
                try:
                    decoded_bytes = base64.b64decode(body)
                    decoded_json = json.loads(decoded_bytes)
                    token = decoded_json.get("admin_token", "")
                    if token:
                        print(f"Flag 3: vsosh{{{token}}}")
                        flag3_found = True
                        break
                except Exception:
                    # Try with padding adjustments
                    for pad in range(4):
                        try:
                            decoded_bytes = base64.b64decode(body + "=" * pad)
                            decoded_json = json.loads(decoded_bytes)
                            token = decoded_json.get("admin_token", "")
                            if token:
                                print(f"Flag 3: vsosh{{{token}}}")
                                flag3_found = True
                                break
                        except Exception:
                            continue
                    if flag3_found:
                        break

if not flag3_found:
    print("[!] Flag 3 not found — no POST /collect with base64 JSON body", file=sys.stderr)
