python find_keybox.py --all

Upgrade:
- Scans entire directory where find_keybox_v2.py is located.
USAGE:
python find_keybox_v2.py --all

MSTAR/FindKeyBox/find_keybox_v2.py

@@ -0,0 +1,240 @@
+#!/usr/bin/env python3
+import mmap
+import os
+import sys
+import hashlib
+from typing import List, Iterable, Tuple
+
+PREFIX_HEX = "4D 53 54 41 52 5F 53 45 43 55 52 45 5F 53 54 4F 52 45 5F 46 49 4C 45 5F 4D 41 47 49 43 5F 49 44"
+POSTFIX_HEX_LIST = ["00 00"]
+
+EXTRACT_LEN = 228
+PRINT_HEXVIEW = True
+HEXVIEW_WIDTH = 16
+MAX_HITS = 0
+
+REQUIRE_POSTFIX_AT_EXTRACT_END = True
+
+MAX_ZERO_FRACTION = 0.25
+MAX_ZERO_RUN = 16
+TAIL_START = 0x90
+MIN_TAIL_NONZERO_RATIO = 0.70
+
+SCRIPT_PATH = os.path.realpath(__file__)
+SCRIPT_DIR = os.path.dirname(SCRIPT_PATH)
+SCRIPT_NAME = os.path.basename(SCRIPT_PATH)
+
+
+def clean_hex(s: str) -> bytes:
+    s = s.replace("0x", "").replace("0X", "")
+    s = "".join(s.split())
+    if not s:
+        return b""
+    if len(s) % 2 != 0:
+        raise ValueError
+    return bytes.fromhex(s)
+
+
+def iter_all(mm: mmap.mmap, needle: bytes, start: int = 0) -> Iterable[int]:
+    i = start
+    while True:
+        pos = mm.find(needle, i)
+        if pos == -1:
+            return
+        yield pos
+        i = pos + 1
+
+
+def is_printable_ascii(b: int) -> bool:
+    return 32 <= b <= 126
+
+
+def hexview(data: bytes, base_offset: int = 0, width: int = 16) -> str:
+    lines = []
+    for i in range(0, len(data), width):
+        chunk = data[i:i + width]
+        hex_part = " ".join(f"{x:02X}" for x in chunk)
+        hex_part_padded = hex_part.ljust(width * 3 - 1)
+        ascii_part = "".join(chr(x) if is_printable_ascii(x) else "." for x in chunk)
+        lines.append(f"{base_offset + i:08X}  {hex_part_padded}  |{ascii_part}|")
+    return "\n".join(lines)
+
+
+def best_postfix_window(postfixes: List[bytes]) -> int:
+    return max((len(p) for p in postfixes if p), default=0)
+
+
+def find_postfix_after_extract(mm: mmap.mmap, postfixes: List[bytes], prefix_pos: int, file_size: int) -> Tuple[int, bytes]:
+    start = prefix_pos + EXTRACT_LEN
+    if start >= file_size:
+        return -1, b""
+    window = best_postfix_window(postfixes)
+    end = min(file_size, start + window)
+    for pf in postfixes:
+        if not pf:
+            continue
+        if start + len(pf) > file_size:
+            continue
+        pos = mm.find(pf, start, end)
+        if pos != -1:
+            return pos, pf
+    return -1, b""
+
+
+def max_zero_run(b: bytes) -> int:
+    best = 0
+    cur = 0
+    for x in b:
+        if x == 0:
+            cur += 1
+            if cur > best:
+                best = cur
+        else:
+            cur = 0
+    return best
+
+
+def zero_fraction(b: bytes) -> float:
+    if not b:
+        return 1.0
+    z = sum(1 for x in b if x == 0)
+    return z / len(b)
+
+
+def nonzero_ratio(b: bytes) -> float:
+    if not b:
+        return 0.0
+    nz = sum(1 for x in b if x != 0)
+    return nz / len(b)
+
+
+def passes_filters(block: bytes) -> bool:
+    if zero_fraction(block) > MAX_ZERO_FRACTION:
+        return False
+    if max_zero_run(block) > MAX_ZERO_RUN:
+        return False
+    tail = block[TAIL_START:] if TAIL_START < len(block) else b""
+    if tail and nonzero_ratio(tail) < MIN_TAIL_NONZERO_RATIO:
+        return False
+    return True
+
+
+def safe_prefix_from_filename(path: str) -> str:
+    base = os.path.basename(path)
+    root, _ = os.path.splitext(base)
+    safe = []
+    for ch in root:
+        if ch.isalnum() or ch in "._-":
+            safe.append(ch)
+        else:
+            safe.append("_")
+    s = "".join(safe).strip("._-")
+    return s or "file"
+
+
+def is_skipped_file(path: str) -> bool:
+    base = os.path.basename(path)
+    if base == SCRIPT_NAME:
+        return True
+    if base.endswith("_keybox.bin") or "_keybox_" in base:
+        return True
+    return False
+
+
+def out_name_for_file(prefix: str, hit_idx: int) -> str:
+    if hit_idx == 1:
+        return f"{prefix}_keybox.bin"
+    return f"{prefix}_keybox_{hit_idx}.bin"
+
+
+def extract_from_file(path: str, prefix: bytes, postfixes: List[bytes]) -> int:
+    if not os.path.isfile(path):
+        return 0
+
+    file_size = os.path.getsize(path)
+    if file_size <= 0:
+        return 0
+
+    per_file_prefix = safe_prefix_from_filename(path)
+    saved = 0
+    seen = set()
+
+    with open(path, "rb") as f:
+        mm = mmap.mmap(f.fileno(), 0, access=mmap.ACCESS_READ)
+        try:
+            for ppos in iter_all(mm, prefix, 0):
+                if ppos + EXTRACT_LEN > file_size:
+                    continue
+
+                qpos, _ = find_postfix_after_extract(mm, postfixes, ppos, file_size)
+                if qpos == -1:
+                    continue
+
+                if REQUIRE_POSTFIX_AT_EXTRACT_END and qpos != (ppos + EXTRACT_LEN):
+                    continue
+
+                block = mm[ppos:ppos + EXTRACT_LEN]
+
+                if not passes_filters(block):
+                    continue
+
+                h = hashlib.sha256(block).digest()
+                if h in seen:
+                    continue
+                seen.add(h)
+
+                saved += 1
+                out_name = out_name_for_file(per_file_prefix, saved)
+                with open(out_name, "wb") as out:
+                    out.write(block)
+
+                print(f"[{os.path.basename(path)} hit {saved}] -> {out_name} prefix@0x{ppos:X} postfix@0x{qpos:X}")
+
+                if PRINT_HEXVIEW:
+                    print(hexview(block, base_offset=ppos, width=HEXVIEW_WIDTH))
+                    print()
+
+                if MAX_HITS and saved >= MAX_HITS:
+                    break
+        finally:
+            mm.close()
+
+    return saved
+
+
+def main() -> int:
+    try:
+        prefix = clean_hex(PREFIX_HEX)
+        postfixes = [clean_hex(x) for x in POSTFIX_HEX_LIST]
+    except:
+        return 2
+
+    postfixes = [p for p in postfixes if p]
+    if not prefix or not postfixes or EXTRACT_LEN <= 0:
+        return 2
+
+    if len(sys.argv) == 2 and sys.argv[1] == "--all":
+        total = 0
+        for name in sorted(os.listdir(SCRIPT_DIR)):
+            path = os.path.join(SCRIPT_DIR, name)
+            if not os.path.isfile(path):
+                continue
+            if is_skipped_file(path):
+                continue
+            total += extract_from_file(path, prefix, postfixes)
+        return 0 if total else 1
+
+    if len(sys.argv) == 2:
+        path = sys.argv[1]
+        if not os.path.isfile(path):
+            return 2
+        return 0 if extract_from_file(path, prefix, postfixes) else 1
+
+    print("Usage:")
+    print("  python find_keybox.py <file>")
+    print("  python find_keybox.py --all")
+    return 2
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())