feat(clew): extract Ren'Py ASTs

src/clew/BUILD.bazel

@@ -1,8 +1,16 @@
-load("@rules_python//python:defs.bzl", "py_binary")
+load("@rules_python//python:defs.bzl", "py_binary", "py_library")
 
 py_binary(
     name = "clew",
     srcs = ["clew.py"],
     main = "clew.py",
-    deps = ["//third_party/python:rpycdec"],
+    deps = [
+        ":extract_renpy_ast",
+        "//third_party/python:rpycdec",
+    ],
+)
+
+py_library(
+    name = "extract_renpy_ast",
+    srcs = ["extract_renpy_ast.py"],
 )

src/clew/clew.py

@@ -4,12 +4,17 @@ import argparse
 import logging
 import os
 import shutil
+import subprocess
 import sys
 from collections import deque
+from collections.abc import Callable
+from concurrent.futures import ThreadPoolExecutor, as_completed
 from pathlib import Path
 
 from rpycdec.rpa import extract_rpa
+from src.clew import extract_renpy_ast as extract_renpy_ast_script
 
+EXTRACTOR = Path(extract_renpy_ast_script.__file__).resolve()
 LOG = logging.getLogger(__name__)
 
 def parse_args(argv: list[str]) -> argparse.Namespace:
@@ -34,6 +39,11 @@ def parse_args(argv: list[str]) -> argparse.Namespace:
         action="store_true",
         help="Replace Clew-managed output directories.",
     )
+    parser.add_argument(
+        "--no-extract-renpy-ast",
+        action="store_true",
+        help="Skip Ren'Py AST extraction.",
+    )
     return parser.parse_args(argv)
 
 
@@ -60,7 +70,7 @@ def extract_archive(archive: Path, destination: Path) -> None:
         extract_rpa(handle, output_dir=str(destination))
 
 
-def extract_game_files(game_root: Path, output_path: Path) -> None:
+def extract_game_files(game_root: Path, output_path: Path, overwrite: bool) -> None:
     queue = deque(game_root.iterdir())
 
     while queue:
@@ -71,22 +81,112 @@ def extract_game_files(game_root: Path, output_path: Path) -> None:
             destination.mkdir(parents=True, exist_ok=True)
             queue.extend(item.iterdir())
         elif item.is_file() and item.suffix.lower() == ".rpa":
-            extract_archive(item, destination.with_suffix(""))
+            archive_output = destination.with_suffix("")
+            if archive_output.is_dir() and not overwrite:
+                LOG.info(f"skipping existing archive extraction: {archive_output}")
+                continue
+            extract_archive(item, archive_output)
         elif item.is_file():
             destination.parent.mkdir(parents=True, exist_ok=True)
             shutil.copy2(item, destination)
 
 
+def find(root: Path, predicate: Callable[[Path], bool]) -> list[Path]:
+    matches = []
+    for current, _directories, files in os.walk(root):
+        current_path = Path(current)
+        if predicate(current_path):
+            matches.append(current_path)
+        for file_name in files:
+            file_path = current_path / file_name
+            if predicate(file_path):
+                matches.append(file_path)
+    return sorted(matches)
+
+def is_game_root(path: Path) -> bool:
+    return (
+        (path / "game").is_dir()
+        and (path / "renpy").is_dir()
+        and (path / "lib" / "linux-x86_64" / "python").is_file())
+
+
+def is_renpy_script(path: Path) -> bool:
+    return path.is_file() and path.suffix.lower() in {".rpyc", ".rpymc"}
+
+
+def run_extract_renpy_ast(
+    bundled_python: Path,
+    game_root: Path,
+    compiled_script: Path,
+    output: Path,
+) -> None:
+    output.parent.mkdir(parents=True, exist_ok=True)
+    command = [
+        "steam-run",
+        str(bundled_python),
+        str(EXTRACTOR),
+        "--game-root",
+        str(game_root),
+        "--input",
+        str(compiled_script),
+        "--output",
+        str(output),
+    ]
+    subprocess.run(command, check=True)
+
+
+def extract_game_ast(extracted_root: Path, ast_root: Path, overwrite: bool) -> None:
+    game_roots = find(extracted_root, is_game_root)
+    if not game_roots:
+        LOG.warning(f"no Ren'Py game roots found under {extracted_root}")
+        return
+
+    tasks = [
+        (game_root, script_path)
+        for game_root in game_roots
+        for script_path in find(game_root / "game", is_renpy_script)
+    ]
+
+    with ThreadPoolExecutor() as executor:
+        futures = {}
+        for game_root, script_path in tasks:
+            output_path = (
+                ast_root
+                / game_root.relative_to(extracted_root)
+                / script_path.relative_to(game_root).with_suffix(script_path.suffix + ".json")
+            )
+            if output_path.exists() and not overwrite:
+                LOG.info(f"skipping existing AST extraction: {output_path}")
+                continue
+
+            future = executor.submit(
+                run_extract_renpy_ast,
+                game_root / "lib" / "linux-x86_64" / "python",
+                game_root,
+                script_path,
+                output_path,
+            )
+            futures[future] = script_path
+
+        for future in as_completed(futures):
+            future.result()
+            LOG.info(f"finished AST extraction: {futures[future]}")
+
+
 def run(args: argparse.Namespace) -> None:
     game_root = args.game_root.resolve()
     output_root = resolve_output_root(args.output_root)
 
     if not game_root.is_dir():
-        LOG.critical(f'game root is not a directory: {game_root}')
+        LOG.critical(f"game root is not a directory: {game_root}")
         raise SystemExit(2)
 
     extracted = get_output_directory(output_root, "extracted")
-    extract_game_files(game_root, extracted)
+    extract_game_files(game_root, extracted, args.overwrite)
+
+    if not args.no_extract_renpy_ast:
+        ast = get_output_directory(output_root, "ast")
+        extract_game_ast(extracted, ast, args.overwrite)
 
 
 if __name__ == "__main__":

src/clew/extract_renpy_ast.py

@@ -0,0 +1,298 @@
+from __future__ import print_function, unicode_literals
+
+import argparse
+import json
+import os
+import struct
+import sys
+import types
+import zlib
+
+try:
+    basestring
+except NameError:
+    basestring = str
+
+try:
+    unicode
+except NameError:
+    unicode = str
+
+try:
+    long
+except NameError:
+    long = int
+
+
+RPYC2_HEADER = b"RENPY RPC2"
+SCHEMA = "clew-renpy-ast-v1"
+
+
+def parse_args(argv):
+    parser = argparse.ArgumentParser(
+        prog="extract-renpy-ast",
+        description="Extract compiled Ren'Py script ASTs as generic JSON.",
+    )
+    parser.add_argument("--game-root", required=True)
+    parser.add_argument("--input", required=True)
+    parser.add_argument("--output", required=True)
+    parser.add_argument("--max-depth", type=int)
+    return parser.parse_args(argv)
+
+
+def read_rpyc_data(path, slot):
+    with open(path, "rb") as handle:
+        header_data = handle.read(1024)
+
+        if header_data[:len(RPYC2_HEADER)] != RPYC2_HEADER:
+            if slot != 1:
+                return None
+            handle.seek(0)
+            return zlib.decompress(handle.read())
+
+        position = len(RPYC2_HEADER)
+        while True:
+            header_slot, start, length = struct.unpack("III", header_data[position:position + 12])
+            if slot == header_slot:
+                break
+            if header_slot == 0:
+                return None
+            position += 12
+
+        handle.seek(start)
+        return zlib.decompress(handle.read(length))
+
+
+def load_compiled_script(path):
+    import renpy.script
+
+    for slot in (2, 1):
+        data = read_rpyc_data(path, slot)
+        if data:
+            metadata, roots = renpy.script.loads(data)
+            return slot, metadata, roots
+
+    raise RuntimeError("no loadable slot found in {0}".format(path))
+
+
+def class_name(value):
+    return value.__class__.__module__ + "." + value.__class__.__name__
+
+
+def collect_slots(value):
+    slots = []
+    try:
+        classes = value.__class__.mro()
+    except AttributeError:
+        classes = [value.__class__]
+
+    for cls in reversed(classes):
+        cls_slots = getattr(cls, "__slots__", [])
+        if isinstance(cls_slots, basestring):
+            cls_slots = [cls_slots]
+        for slot in cls_slots:
+            if slot not in slots:
+                slots.append(slot)
+    return slots
+
+
+def safe_text(value):
+    if isinstance(value, unicode):
+        return value
+    if isinstance(value, str):
+        try:
+            return value.decode("utf-8")
+        except UnicodeDecodeError:
+            return {"type": "bytes", "repr": repr(value)}
+    return value
+
+
+def is_primitive(value):
+    return value is None or isinstance(value, (bool, int, long, float, basestring))
+
+
+def serialize_function(value):
+    code = getattr(value, "func_code", None) or getattr(value, "__code__", None)
+    result = {
+        "type": "function",
+        "name": safe_text(getattr(value, "func_name", None) or getattr(value, "__name__", None)),
+        "module": safe_text(getattr(value, "__module__", None)),
+    }
+    if code is not None:
+        result["filename"] = safe_text(getattr(code, "co_filename", None))
+        result["firstlineno"] = getattr(code, "co_firstlineno", None)
+    return result
+
+
+def serialize_method(value):
+    function = getattr(value, "im_func", None) or getattr(value, "__func__", None)
+    receiver = getattr(value, "im_self", None) or getattr(value, "__self__", None)
+    result = serialize_function(function) if function is not None else {
+        "type": "function",
+        "name": safe_text(getattr(value, "__name__", None)),
+        "module": safe_text(getattr(value, "__module__", None)),
+    }
+    result["type"] = "method"
+    if receiver is not None:
+        result["receiver_class"] = class_name(receiver)
+    return result
+
+
+def serialize_type(value):
+    return {
+        "type": "python-type",
+        "name": safe_text(getattr(value, "__name__", None)),
+        "module": safe_text(getattr(value, "__module__", None)),
+    }
+
+
+class Serializer(object):
+    def __init__(self, max_depth):
+        self.max_depth = max_depth
+        self.next_id = 1
+        self.seen = {}
+        self.active_containers = set()
+
+    def serialize(self, value, depth=0):
+        if is_primitive(value):
+            return safe_text(value)
+
+        if self.max_depth is not None and depth >= self.max_depth:
+            return {
+                "type": "truncated",
+                "class": class_name(value),
+                "repr": repr(value),
+            }
+
+        if isinstance(value, (list, tuple)):
+            object_id = id(value)
+            if object_id in self.active_containers:
+                return {"type": "recursion", "class": class_name(value)}
+            self.active_containers.add(object_id)
+            items = [self.serialize(item, depth + 1) for item in value]
+            self.active_containers.remove(object_id)
+            return {"type": "tuple" if isinstance(value, tuple) else "list", "items": items}
+
+        if isinstance(value, dict):
+            object_id = id(value)
+            if object_id in self.active_containers:
+                return {"type": "recursion", "class": class_name(value)}
+            self.active_containers.add(object_id)
+            items = []
+            for key, item in sorted(value.items(), key=lambda pair: repr(pair[0])):
+                items.append({"key": self.serialize(key, depth + 1), "value": self.serialize(item, depth + 1)})
+            self.active_containers.remove(object_id)
+            return {"type": "dict", "items": items}
+
+        if isinstance(value, types.FunctionType) or isinstance(value, types.BuiltinFunctionType):
+            return serialize_function(value)
+
+        if isinstance(value, types.MethodType):
+            return serialize_method(value)
+
+        if isinstance(value, type):
+            return serialize_type(value)
+
+        identity = id(value)
+        if identity in self.seen:
+            return {"type": "ref", "id": self.seen[identity]}
+
+        value_id = self.next_id
+        self.next_id += 1
+        self.seen[identity] = value_id
+
+        attributes = {}
+        for slot in collect_slots(value):
+            if slot.startswith("__"):
+                continue
+            try:
+                slot_value = getattr(value, slot)
+            except AttributeError:
+                continue
+            except Exception as error:
+                attributes[slot] = {
+                    "type": "attribute-error",
+                    "class": class_name(error),
+                    "repr": repr(error),
+                }
+                continue
+            attributes[slot] = self.serialize(slot_value, depth + 1)
+
+        value_dict = getattr(value, "__dict__", None)
+        if value_dict:
+            for key, item in sorted(value_dict.items(), key=lambda pair: repr(pair[0])):
+                if isinstance(key, basestring):
+                    key_text = safe_text(key)
+                else:
+                    key_text = repr(key)
+                attributes[key_text] = self.serialize(item, depth + 1)
+
+        if not attributes:
+            del self.seen[identity]
+            return {
+                "type": "repr",
+                "class": class_name(value),
+                "repr": repr(value),
+            }
+
+        return {
+            "type": "object",
+            "id": value_id,
+            "class": class_name(value),
+            "attributes": attributes,
+        }
+
+
+def write_json(path, document):
+    parent = os.path.dirname(path)
+    if parent and not os.path.isdir(parent):
+        os.makedirs(parent)
+
+    data = json.dumps(document, separators=(",", ":"), sort_keys=True)
+    if not isinstance(data, bytes):
+        data = data.encode("utf-8")
+
+    with open(path, "wb") as handle:
+        handle.write(data)
+        handle.write(b"\n")
+
+
+def import_renpy(game_root):
+    sys.path.insert(0, game_root)
+
+    import renpy
+    renpy.import_all()
+
+    class Dummy(object):
+        pass
+
+    renpy.game.script = Dummy()
+    renpy.game.script.all_pyexpr = None
+    renpy.game.script.record_pycode = False
+    renpy.game.log = Dummy()
+    renpy.game.log.mutated = {}
+
+    return renpy
+
+
+def main(argv):
+    args = parse_args(argv)
+    renpy = import_renpy(args.game_root)
+    slot, metadata, roots = load_compiled_script(args.input)
+
+    serializer = Serializer(args.max_depth)
+    document = {
+        "schema": SCHEMA,
+        "source": os.path.abspath(args.input),
+        "game_root": os.path.abspath(args.game_root),
+        "slot": slot,
+        "python_version": sys.version,
+        "renpy_version": getattr(renpy, "version", None),
+        "metadata": serializer.serialize(metadata),
+        "roots": serializer.serialize(roots),
+    }
+    write_json(args.output, document)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])