#!/usr/bin/env python3
"""Generate test fixture PDFs for CorbelPurge.

Creates four PDFs:
  - benign.pdf           — a clean PDF with just text
  - malicious_js.pdf     — a PDF with a /JavaScript action
  - malicious_launch.pdf — a PDF with a /Launch action
  - cve_writeup.pdf      — an educational PDF that mentions JS / shellcode

Uses pypdf for xref-safe object injection.
"""

from pathlib import Path
from reportlab.pdfgen import canvas
from reportlab.lib.pagesizes import letter

import pypdf
from pypdf.generic import (
    ArrayObject,
    DictionaryObject,
    NameObject,
    NumberObject,
    TextStringObject,
    IndirectObject,
)

FIXTURES_DIR = Path(__file__).parent.parent / "tests" / "fixtures"
FIXTURES_DIR.mkdir(parents=True, exist_ok=True)


def make_benign_pdf():
    """A clean PDF with just text — no JavaScript, no embedded files."""
    path = FIXTURES_DIR / "benign.pdf"
    c = canvas.Canvas(str(path), pagesize=letter)
    c.drawString(100, 700, "Hello, this is a benign PDF.")
    c.drawString(100, 680, "It contains only text and no active content.")
    c.setTitle("Benign Test PDF")
    c.setAuthor("CorbelPurge Tests")
    c.save()
    return path


def make_malicious_js_pdf():
    """A PDF with a /JavaScript action attached to the catalog's /OpenAction."""
    # First, generate a base PDF with reportlab.
    base_path = FIXTURES_DIR / "_base_js.pdf"
    c = canvas.Canvas(str(base_path), pagesize=letter)
    c.drawString(100, 700, "This PDF has a JavaScript action.")
    c.setTitle("Malicious JS Test PDF")
    c.save()

    # Now use pypdf to inject a /JavaScript action safely.
    path = FIXTURES_DIR / "malicious_js.pdf"
    reader = pypdf.PdfReader(str(base_path))
    writer = pypdf.PdfWriter()

    # Copy all pages.
    for page in reader.pages:
        writer.add_page(page)

    # Add a new /JavaScript action object.
    js_action = DictionaryObject({
        NameObject("/Type"): NameObject("/Action"),
        NameObject("/S"): NameObject("/JavaScript"),
        NameObject("/JS"): TextStringObject("app.alert('XSS from PDF');"),
    })
    js_action_ref = writer._add_object(js_action)

    # Attach it to the catalog's /OpenAction.
    writer._root_object[NameObject("/OpenAction")] = js_action_ref

    with open(path, "wb") as f:
        writer.write(f)

    base_path.unlink()
    return path


def make_malicious_launch_pdf():
    """A PDF with a /Launch action."""
    base_path = FIXTURES_DIR / "_base_launch.pdf"
    c = canvas.Canvas(str(base_path), pagesize=letter)
    c.drawString(100, 700, "This PDF has a Launch action.")
    c.setTitle("Malicious Launch Test PDF")
    c.save()

    path = FIXTURES_DIR / "malicious_launch.pdf"
    reader = pypdf.PdfReader(str(base_path))
    writer = pypdf.PdfWriter()

    for page in reader.pages:
        writer.add_page(page)

    # Add a /Launch action.
    launch_action = DictionaryObject({
        NameObject("/Type"): NameObject("/Action"),
        NameObject("/S"): NameObject("/Launch"),
        NameObject("/F"): TextStringObject("/bin/sh"),
        NameObject("/Win"): DictionaryObject({
            NameObject("/F"): TextStringObject("cmd.exe"),
        }),
    })
    launch_ref = writer._add_object(launch_action)

    writer._root_object[NameObject("/OpenAction")] = launch_ref

    with open(path, "wb") as f:
        writer.write(f)

    base_path.unlink()
    return path


def make_cve_writeup_pdf():
    """A PDF that *mentions* JavaScript and shellcode in an educational context."""
    path = FIXTURES_DIR / "cve_writeup.pdf"
    c = canvas.Canvas(str(path), pagesize=letter)
    c.setTitle("CVE-2024-1234 Writeup")
    c.setAuthor("Security Researcher")
    c.drawString(100, 750, "CVE-2024-1234: PDF JavaScript Injection Analysis")
    c.drawString(100, 720, "Abstract")
    text = c.beginText(100, 700)
    text.setFont("Helvetica", 10)
    text.textLines(
        "In this paper we describe a vulnerability in which a malicious PDF\n"
        "uses a /JavaScript action to execute arbitrary code. The eval()\n"
        "function is called with attacker-controlled input. Remediation:\n"
        "patch the reader to ignore /JavaScript actions in /OpenAction."
    )
    c.drawText(text)
    c.save()
    return path


def main():
    paths = [
        make_benign_pdf(),
        make_malicious_js_pdf(),
        make_malicious_launch_pdf(),
        make_cve_writeup_pdf(),
    ]
    for p in paths:
        print(f"  wrote {p} ({p.stat().st_size} bytes)")


if __name__ == "__main__":
    main()