#!/usr/bin/env python3
"""Generate DOCX test fixtures for CorbelPurge.

Creates:
  - benign.docx        — clean DOCX with just text
  - malicious_macro.docx — DOCX with a VBA macro stub
  - malicious_ole.docx — DOCX with an embedded OLE object (PE)
  - malicious_link.docx — DOCX with an external phishing hyperlink
"""

from pathlib import Path
import zipfile

from docx import Document
from docx.opc.constants import RELATIONSHIP_TYPE as RT

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


def make_benign_docx():
    """A clean DOCX with just text — no macros, no embedded objects."""
    path = FIXTURES_DIR / "benign.docx"
    doc = Document()
    doc.core_properties.title = "Benign Test DOCX"
    doc.core_properties.author = "CorbelPurge Tests"
    doc.core_properties.subject = "Test subject"
    doc.add_paragraph("Hello, this is a benign DOCX.")
    doc.add_paragraph("Second paragraph of benign content.")
    doc.save(str(path))
    return path


def make_malicious_macro_docx():
    """A DOCX with a VBA macro stub injected into the ZIP."""
    # First create a normal DOCX.
    base_path = FIXTURES_DIR / "_base_macro.docx"
    doc = Document()
    doc.core_properties.title = "Malicious Macro DOCX"
    doc.add_paragraph("This DOCX contains a VBA macro.")
    doc.save(str(base_path))

    # Now copy the ZIP and inject a fake word/vbaProject.xml.
    path = FIXTURES_DIR / "malicious_macro.docx"
    with zipfile.ZipFile(base_path, "r") as src, zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as dst:
        for item in src.infolist():
            dst.writestr(item, src.read(item.filename))
        # Inject the macro file.
        dst.writestr(
            "word/vbaProject.xml",
            "<?xml version='1.0'?>"
            "<vbaProject><module name='Module1'>"
            "Sub AutoOpen()\n"
            "  MsgBox \"Hello from VBA\"\n"
            "End Sub"
            "</module></vbaProject>",
        )
    base_path.unlink()
    return path


def make_malicious_ole_docx():
    """A DOCX with an embedded OLE object (fake PE)."""
    base_path = FIXTURES_DIR / "_base_ole.docx"
    doc = Document()
    doc.core_properties.title = "Malicious OLE DOCX"
    doc.add_paragraph("This DOCX contains an embedded OLE object.")
    doc.save(str(base_path))

    path = FIXTURES_DIR / "malicious_ole.docx"
    with zipfile.ZipFile(base_path, "r") as src, zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as dst:
        for item in src.infolist():
            dst.writestr(item, src.read(item.filename))
        # Inject an embedded OLE object with PE signature.
        # MZ + dummy DOS header.
        pe_bytes = b"MZ\x90\x00\x03\x00\x00\x00" + b"\x00" * 56 + b"PE\x00\x00"
        dst.writestr("word/embeddings/oleObject1.bin", pe_bytes)
    base_path.unlink()
    return path


def make_malicious_link_docx():
    """A DOCX with an external phishing hyperlink (micros0ft homograph)."""
    path = FIXTURES_DIR / "malicious_link.docx"
    doc = Document()
    doc.core_properties.title = "Malicious Link DOCX"
    doc.add_paragraph("This DOCX contains a suspicious external hyperlink:")

    # Add a paragraph with an external hyperlink.
    para = doc.add_paragraph()
    run = para.add_run("Click here to verify your account")
    # python-docx doesn't directly expose external hyperlinks, so we
    # post-process the XML.
    doc.save(str(path))

    # Post-process: add a hyperlink relationship and wrap the run.
    import re
    from docx.opc.packuri import PackURI
    from docx.opc.part import Part
    from docx.opc.constants import CONTENT_TYPE as CT

    # Simpler approach: just unzip, edit document.xml.rels to add an
    # external hyperlink relationship. The parser will detect it.
    base_path2 = FIXTURES_DIR / "_base_link.docx"
    Path(path).rename(base_path2)

    with zipfile.ZipFile(base_path2, "r") as src, zipfile.ZipFile(path, "w", zipfile.ZIP_DEFLATED) as dst:
        for item in src.infolist():
            data = src.read(item.filename)
            if item.filename == "word/_rels/document.xml.rels":
                # Inject a new external-link relationship.
                text = data.decode("utf-8")
                new_rel = (
                    '<Relationship Id="rIdEvil" '
                    'Type="http://schemas.openxmlformats.org/officeDocument/2006/relationships/hyperlink" '
                    'Target="https://micros0ft.com/account/verify" '
                    'TargetMode="External"/>'
                )
                text = text.replace("</Relationships>", new_rel + "</Relationships>")
                data = text.encode("utf-8")
            dst.writestr(item, data)
    base_path2.unlink()
    return path


def main():
    paths = [
        make_benign_docx(),
        make_malicious_macro_docx(),
        make_malicious_ole_docx(),
        make_malicious_link_docx(),
    ]
    for p in paths:
        print(f"  wrote {p} ({p.stat().st_size} bytes)")


if __name__ == "__main__":
    main()