Allow dumping decompiled bytecode to file for both raw afp/bsi pairs and TXP2 containers.

bemani/format/afp/__init__.py

@@ -1,5 +1,23 @@
 from .geo import Shape, DrawParams
-from .swf import SWF, NamedTagReference
+from .swf import (
+    SWF,
+    NamedTagReference,
+    Frame,
+    Tag,
+    TagPointer,
+    AP2ShapeTag,
+    AP2DefineFontTag,
+    AP2TextChar,
+    AP2TextLine,
+    AP2DefineMorphShapeTag,
+    AP2DefineButtonTag,
+    AP2DefineTextTag,
+    AP2DoActionTag,
+    AP2PlaceObjectTag,
+    AP2RemoveObjectTag,
+    AP2DefineSpriteTag,
+    AP2DefineEditTextTag,
+)
 from .container import TXP2File, PMAN, Texture, TextureRegion, Unknown1, Unknown2
 from .render import AFPRenderer
 from .types import Matrix, Color, Point, Rectangle, AP2Tag, AP2Action, AP2Object, AP2Pointer
@@ -9,6 +27,21 @@ __all__ = [
     'Shape',
     'DrawParams',
     'SWF',
+    'Frame',
+    'Tag',
+    'TagPointer',
+    'AP2ShapeTag',
+    'AP2DefineFontTag',
+    'AP2TextChar',
+    'AP2TextLine',
+    'AP2DefineMorphShapeTag',
+    'AP2DefineButtonTag',
+    'AP2DefineTextTag',
+    'AP2DoActionTag',
+    'AP2PlaceObjectTag',
+    'AP2RemoveObjectTag',
+    'AP2DefineSpriteTag',
+    'AP2DefineEditTextTag',
     'NamedTagReference',
     'TXP2File',
     'PMAN',

bemani/format/afp/swf.py

@@ -177,7 +177,7 @@ class AP2TextLine:
         }
 
 
-class AP2DefineMorphShape(Tag):
+class AP2DefineMorphShapeTag(Tag):
     def __init__(self, id: int) -> None:
         # TODO: I need to figure out what morph shapes actually DO, and take the
         # values that I parsed out store them here...
@@ -189,7 +189,7 @@ class AP2DefineMorphShape(Tag):
         }
 
 
-class AP2DefineButton(Tag):
+class AP2DefineButtonTag(Tag):
     def __init__(self, id: int) -> None:
         # TODO: I need to figure out what buttons actually DO, and take the
         # values that I parsed out store them here...
@@ -326,7 +326,7 @@ class AP2RemoveObjectTag(Tag):
 
 
 class AP2DefineSpriteTag(Tag):
-    def __init__(self, id: int, tags: List[Tag], frames: List[Frame], references: Dict[int, str]) -> None:
+    def __init__(self, id: int, tags: List[Tag], frames: List[Frame], references: Dict[str, int]) -> None:
         super().__init__(id)
 
         # The list of tags that this sprite consists of. Sprites are, much like vanilla
@@ -428,7 +428,7 @@ class SWF(TrackedCoverage, VerboseOutput):
         self.frames: List[Frame] = []
 
         # Reference LUT for mapping object reference IDs to names a used in bytecode.
-        self.references: Dict[int, str] = {}
+        self.references: Dict[str, int] = {}
 
         # SWF string table. This is used for faster lookup of strings as well as
         # tracking which strings in the table have been parsed correctly.
@@ -934,7 +934,18 @@ class SWF(TrackedCoverage, VerboseOutput):
 
         return ByteCode(bytecode_name, actions, offset_ptr)
 
-    def __parse_tag(self, ap2_version: int, afp_version: int, ap2data: bytes, tagid: int, size: int, dataoffset: int, tag_frame: int, prefix: str = "") -> Tag:
+    def __parse_tag(
+        self,
+        ap2_version: int,
+        afp_version: int,
+        ap2data: bytes,
+        tagid: int,
+        size: int,
+        dataoffset: int,
+        tag_parent_sprite: Optional[int],
+        tag_frame: int,
+        prefix: str = "",
+    ) -> Tag:
         if tagid == AP2Tag.AP2_SHAPE:
             if size != 4:
                 raise Exception(f"Invalid shape size {size}")
@@ -961,7 +972,7 @@ class SWF(TrackedCoverage, VerboseOutput):
                 self.add_coverage(dataoffset + 4, 4)
 
             self.vprint(f"{prefix}    Tag ID: {sprite_id}")
-            tags, frames, references = self.__parse_tags(ap2_version, afp_version, ap2data, subtags_offset, tag_frame, prefix="      " + prefix)
+            tags, frames, references = self.__parse_tags(ap2_version, afp_version, ap2data, subtags_offset, sprite_id, prefix="      " + prefix)
 
             return AP2DefineSpriteTag(sprite_id, tags, frames, references)
         elif tagid == AP2Tag.AP2_DEFINE_FONT:
@@ -997,7 +1008,7 @@ class SWF(TrackedCoverage, VerboseOutput):
             return AP2DefineFontTag(font_id, fontname, xml_prefix, heights, text_indexes)
         elif tagid == AP2Tag.AP2_DO_ACTION:
             datachunk = ap2data[dataoffset:(dataoffset + size)]
-            bytecode = self.__parse_bytecode(f"on_enter_frame_{tag_frame}", datachunk, prefix=prefix)
+            bytecode = self.__parse_bytecode(f"on_enter_{f'sprite_{tag_parent_sprite}' if tag_parent_sprite is not None else 'main'}_frame_{tag_frame}", datachunk, prefix=prefix)
             self.add_coverage(dataoffset, size)
 
             return AP2DoActionTag(bytecode)
@@ -1593,7 +1604,7 @@ class SWF(TrackedCoverage, VerboseOutput):
 
                     self.vprint(f"{prefix}        Floats: {fv1} {fv2} {fv3} {fv4} {fv5} {fv6} {fv7} {fv8}")
 
-            return AP2DefineMorphShape(define_shape_id)
+            return AP2DefineMorphShapeTag(define_shape_id)
         elif tagid == AP2Tag.AP2_DEFINE_BUTTON:
             flags, button_id, source_tags_count, bytecode_count = struct.unpack("<HHHH", ap2data[dataoffset:(dataoffset + 8)])
             self.add_coverage(dataoffset, 8)
@@ -1757,12 +1768,20 @@ class SWF(TrackedCoverage, VerboseOutput):
                 self.vprint(f"{prefix}      Offset: {hex(loc)}, Sound Unk1: {unk1}, Source Tag ID: {sound_source_tag}")
                 raise Exception("TODO: Need to examine this section further if I find data with it!")
 
-            return AP2DefineButton(button_id)
+            return AP2DefineButtonTag(button_id)
         else:
             self.vprint(f"Unknown tag {hex(tagid)} with data {ap2data[dataoffset:(dataoffset + size)]!r}")
             raise Exception(f"Unimplemented tag {hex(tagid)}!")
 
-    def __parse_tags(self, ap2_version: int, afp_version: int, ap2data: bytes, tags_base_offset: int, sprite_frame: int, prefix: str = "") -> Tuple[List[Tag], List[Frame], Dict[int, str]]:
+    def __parse_tags(
+        self,
+        ap2_version: int,
+        afp_version: int,
+        ap2data: bytes,
+        tags_base_offset: int,
+        sprite: Optional[int],
+        prefix: str = "",
+    ) -> Tuple[List[Tag], List[Frame], Dict[str, int]]:
         name_reference_flags, name_reference_count, frame_count, tags_count, name_reference_offset, frame_offset, tags_offset = struct.unpack(
             "<HHIIIII",
             ap2data[tags_base_offset:(tags_base_offset + 24)]
@@ -1807,19 +1826,19 @@ class SWF(TrackedCoverage, VerboseOutput):
                 raise Exception(f"Invalid tag size {size} ({hex(size)})")
 
             self.vprint(f"{prefix}  Tag: {hex(tagid)} ({AP2Tag.tag_to_name(tagid)}), Size: {hex(size)}, Offset: {hex(tags_offset + 4)}")
-            tags.append(self.__parse_tag(ap2_version, afp_version, ap2data, tagid, size, tags_offset + 4, tag_to_frame[i] + sprite_frame, prefix=prefix))
+            tags.append(self.__parse_tag(ap2_version, afp_version, ap2data, tagid, size, tags_offset + 4, sprite, tag_to_frame[i], prefix=prefix))
             tags_offset += ((size + 3) & 0xFFFFFFFC) + 4  # Skip past tag header and data, rounding to the nearest 4 bytes.
 
         # Finally, parse place object name references.
         self.vprint(f"{prefix}Number of Object Name References: {name_reference_count}, Flags: {hex(name_reference_flags)}")
-        references: Dict[int, str] = {}
+        references: Dict[str, int] = {}
         for i in range(name_reference_count):
-            index, stringoffset = struct.unpack("<HH", ap2data[name_reference_offset:(name_reference_offset + 4)])
+            frameno, stringoffset = struct.unpack("<HH", ap2data[name_reference_offset:(name_reference_offset + 4)])
             strval = self.__get_string(stringoffset)
             self.add_coverage(name_reference_offset, 4)
-            references[index] = strval
+            references[strval] = frameno
 
-            self.vprint(f"{prefix}  Name Reference: {index}, Name: {strval}")
+            self.vprint(f"{prefix}  Frame Number: {frameno}, Name: {strval}")
             name_reference_offset += 4
 
         return tags, frames, references
@@ -2014,7 +2033,7 @@ class SWF(TrackedCoverage, VerboseOutput):
         # Tag sections
         tags_offset = struct.unpack("<I", data[36:40])[0]
         self.add_coverage(36, 4)
-        self.tags, self.frames, self.references = self.__parse_tags(ap2_data_version, version, data, tags_offset, 0)
+        self.tags, self.frames, self.references = self.__parse_tags(ap2_data_version, version, data, tags_offset, None)
 
         # Imported tags sections
         imported_tags_count = struct.unpack("<h", data[34:36])[0]

bemani/utils/afputils.py

@@ -7,17 +7,67 @@ import os.path
 import sys
 import textwrap
 from PIL import Image, ImageDraw  # type: ignore
-from typing import Any, Dict
+from typing import Any, Dict, List
 
-from bemani.format.afp import TXP2File, Shape, SWF, AFPRenderer, Color
+from bemani.format.afp import TXP2File, Shape, SWF, Frame, Tag, AP2DoActionTag, AP2PlaceObjectTag, AP2DefineSpriteTag, AFPRenderer, Color
 from bemani.format import IFS
 
 
+def write_bytecode(swf: SWF, directory: str, verbose: bool=False) -> None:
+    # Actually place the files down.
+    os.makedirs(directory, exist_ok=True)
+
+    # Buffer for where the decompiled data goes.
+    buff: List[str] = []
+    lut: Dict[str, int] = {}
+
+    def bytecode_from_frames(frames: List[Frame]) -> None:
+        for frame in frames:
+            for tag in frame.imported_tags:
+                if tag.init_bytecode:
+                    buff.append(tag.init_bytecode.decompile(verbose=verbose))
+
+    def bytecode_from_tags(tags: List[Tag]) -> None:
+        for tag in tags:
+            if isinstance(tag, AP2DoActionTag):
+                buff.append(tag.bytecode.decompile(verbose=verbose))
+            elif isinstance(tag, AP2PlaceObjectTag):
+                for _, triggers in tag.triggers.items():
+                    for trigger in triggers:
+                        buff.append(trigger.decompile(verbose=verbose))
+            elif isinstance(tag, AP2DefineSpriteTag):
+                lut.update(tag.references)
+                bytecode_from_frames(tag.frames)
+                bytecode_from_tags(tag.tags)
+
+    lut.update(swf.references)
+    bytecode_from_frames(swf.frames)
+    bytecode_from_tags(swf.tags)
+
+    # If we have references, put them at the top as global defines.
+    if lut:
+        buff = [
+            os.linesep.join([
+                '// Defined string references from SWF container, as used for frame lookups.',
+                'FRAME_LUT = {',
+                *[f"    {name!r}: {frame}," for name, frame in lut.items()],
+                '};',
+            ]),
+            *buff,
+        ]
+
+    # Now, write it out.
+    filename = os.path.join(directory, swf.exported_name) + ".code"
+    print(f"Writing code to {filename}...")
+    with open(filename, "wb") as bfp:
+        bfp.write(f"{os.linesep}{os.linesep}".join(buff).encode('utf-8'))
+
+
 def main() -> int:
     parser = argparse.ArgumentParser(description="Konami AFP graphic file unpacker/repacker")
     subparsers = parser.add_subparsers(help='Action to take', dest='action')
 
-    extract_parser = subparsers.add_parser('extract', help='Extract relevant textures from TXP2 container')
+    extract_parser = subparsers.add_parser('extract', help='Extract relevant file data and textures from a TXP2 container')
     extract_parser.add_argument(
         "file",
         metavar="FILE",
@@ -70,6 +120,12 @@ def main() -> int:
         action="store_true",
         help="Write binary SWF files to disk",
     )
+    extract_parser.add_argument(
+        "-y",
+        "--write-bytecode",
+        action="store_true",
+        help="Write decompiled bytecode files to disk",
+    )
 
     update_parser = subparsers.add_parser('update', help='Update relevant textures in a TXP2 container from a directory')
     update_parser.add_argument(
@@ -138,6 +194,32 @@ def main() -> int:
         help="Display verbuse debugging output",
     )
 
+    decompile_parser = subparsers.add_parser('decompile', help='Decompile bytecode in a raw AFP/BSI file pair previously extracted from an IFS or TXP2 container')
+    decompile_parser.add_argument(
+        "afp",
+        metavar="AFPFILE",
+        help="The AFP file to parse",
+    )
+    decompile_parser.add_argument(
+        "bsi",
+        metavar="BSIFILE",
+        help="The BSI file to parse",
+    )
+    decompile_parser.add_argument(
+        "-d",
+        "--directory",
+        metavar="DIR",
+        default='.',
+        type=str,
+        help="Directory to extract to after decompiling. Defaults to current directory.",
+    )
+    decompile_parser.add_argument(
+        "-v",
+        "--verbose",
+        action="store_true",
+        help="Display verbuse debugging output",
+    )
+
     parsegeo_parser = subparsers.add_parser('parsegeo', help='Parse a raw GEO file previously extracted from an IFS or TXP2 container')
     parsegeo_parser.add_argument(
         "geo",
@@ -390,6 +472,9 @@ def main() -> int:
                     if not announced.get(texturename, False):
                         print(f"Cannot extract sprites from {texturename} because it is not a supported format!")
                         announced[texturename] = True
+        if args.write_bytecode:
+            for swf in afpfile.swfdata:
+                write_bytecode(swf, args.dir, verbose=args.verbose)
 
     if args.action == "update":
         # First, parse the file out
@@ -451,6 +536,16 @@ def main() -> int:
         swf.parse(verbose=args.verbose)
         print(json.dumps(swf.as_dict(decompile_bytecode=args.decompile_bytecode, verbose=args.verbose), sort_keys=True, indent=4))
 
+    if args.action == "decompile":
+        # First, load the AFP and BSI files
+        with open(args.afp, "rb") as bafp:
+            with open(args.bsi, "rb") as bbsi:
+                swf = SWF("<unnamed>", bafp.read(), bbsi.read())
+
+        # Now, decompile it
+        swf.parse(verbose=args.verbose)
+        write_bytecode(swf, args.directory, verbose=args.verbose)
+
     if args.action == "parsegeo":
         # First, load the AFP and BSI files
         with open(args.geo, "rb") as bfp: