From 3a1ebf3f3e4670937e94f1db06753f423c8a72c7 Mon Sep 17 00:00:00 2001
From: Jennifer Taylor <dragonminded@dragonminded.com>
Date: Tue, 19 Sep 2023 00:08:01 +0000
Subject: [PATCH] Add a TDXT extractor and partial updater.

---
 bemani/format/__init__.py |   2 +
 bemani/format/tdxt.py     |   2 +-
 bemani/utils/tdxtutils.py | 117 ++++++++++++++++++++++++++++++++++++++
 tdxtutils                 |  12 ++++
 verifytyping              |   1 +
 5 files changed, 133 insertions(+), 1 deletion(-)
 create mode 100644 bemani/utils/tdxtutils.py
 create mode 100755 tdxtutils

diff --git a/bemani/format/__init__.py b/bemani/format/__init__.py
index e64ab35..4bae4a4 100644
--- a/bemani/format/__init__.py
+++ b/bemani/format/__init__.py
@@ -1,6 +1,7 @@
 from bemani.format.ifs import IFS
 from bemani.format.arc import ARC
 from bemani.format.twodx import TwoDX
+from bemani.format.tdxt import TDXT
 from bemani.format.iidxchart import IIDXChart
 from bemani.format.iidxmusicdb import IIDXMusicDB, IIDXSong
 
@@ -9,6 +10,7 @@ __all__ = [
     "IFS",
     "ARC",
     "TwoDX",
+    "TDXT",
     "IIDXChart",
     "IIDXMusicDB",
     "IIDXSong",
diff --git a/bemani/format/tdxt.py b/bemani/format/tdxt.py
index ebacf8c..59f5d7c 100644
--- a/bemani/format/tdxt.py
+++ b/bemani/format/tdxt.py
@@ -78,7 +78,7 @@ class TDXT:
             f"{endian}4sIIIHHIII",
             raw_data[0:32],
         )
-        if raw_length != len(raw_data):
+        if (raw_length + 64) != len(raw_data):
             raise Exception("Invalid texture length!")
 
         # I have only ever observed the following values across two different games.
diff --git a/bemani/utils/tdxtutils.py b/bemani/utils/tdxtutils.py
new file mode 100644
index 0000000..eca8415
--- /dev/null
+++ b/bemani/utils/tdxtutils.py
@@ -0,0 +1,117 @@
+#! /usr/bin/env python3
+import argparse
+import io
+import json
+import math
+import os
+import os.path
+import sys
+import textwrap
+from PIL import Image, ImageDraw
+from typing import Any, Dict, List, Optional, Tuple, TypeVar
+
+from bemani.format import TDXT
+
+
+def extract_texture(
+    fname: str,
+    output_fname: Optional[str],
+) -> int:
+    with open(fname, "rb") as bfp:
+        tdxt = TDXT.fromBytes(bfp.read())
+
+    if output_fname is None:
+        output_fname = os.path.splitext(os.path.abspath(fname))[0] + ".png"
+
+    if not output_fname.lower().endswith(".png"):
+        raise Exception("Invalid output file format!")
+
+    # Actually place the files down.
+    output_dir = os.path.dirname(os.path.abspath(output_fname))
+    os.makedirs(output_dir, exist_ok=True)
+
+    print(f"Extracting texture from {os.path.abspath(fname)} to {os.path.abspath(output_fname)}")
+    with open(output_fname, "wb") as bfp:
+        tdxt.img.save(bfp, format="PNG")
+
+    return 0
+
+
+def update_texture(
+    fname: str,
+    input_fname: str,
+) -> int:
+    with open(fname, "rb") as bfp:
+        tdxt = TDXT.fromBytes(bfp.read())
+
+    if not input_fname.lower().endswith(".png"):
+        raise Exception("Invalid output file format!")
+
+    with open(input_fname, "rb") as bfp:
+        img = Image.open(io.BytesIO(bfp.read()))
+
+    tdxt.img = img
+
+    print(f"Updating texture in {os.path.abspath(fname)} from {os.path.abspath(input_fname)}")
+    with open(fname, "wb") as bfp:
+        bfp.write(tdxt.toBytes())
+
+    return 0
+
+
+def main() -> int:
+    parser = argparse.ArgumentParser(
+        description="Konami TDXT graphic file unpacker/repacker."
+    )
+    subparsers = parser.add_subparsers(help="Action to take", dest="action")
+
+    unpack_parser = subparsers.add_parser(
+        "unpack",
+        help="Unpack texture data into a PNG file.",
+    )
+    unpack_parser.add_argument(
+        "infile",
+        metavar="INFILE",
+        help="The TDXT container to unpack the texture from.",
+    )
+    unpack_parser.add_argument(
+        "outfile",
+        metavar="OUTFILE",
+        nargs="?",
+        default=None,
+        help="The PNG file to unpack the texture to.",
+    )
+
+    update_parser = subparsers.add_parser(
+        "update",
+        help="Update texture data from a PNG file.",
+    )
+    update_parser.add_argument(
+        "outfile",
+        metavar="OUTFILE",
+        help="The TDXT container to update the texture to, must already exist.",
+    )
+    update_parser.add_argument(
+        "infile",
+        metavar="INFILE",
+        help="The PNG file to update the texture from.",
+    )
+
+    args = parser.parse_args()
+
+    if args.action == "unpack":
+        return extract_texture(
+            args.infile,
+            args.outfile,
+        )
+    elif args.action == "update":
+        return update_texture(
+            args.outfile,
+            args.infile,
+        )
+    else:
+        raise Exception(f"Invalid action {args.action}!")
+
+
+if __name__ == "__main__":
+    sys.exit(main())
diff --git a/tdxtutils b/tdxtutils
new file mode 100755
index 0000000..49ddf3a
--- /dev/null
+++ b/tdxtutils
@@ -0,0 +1,12 @@
+#! /usr/bin/env python3
+if __name__ == "__main__":
+    import os
+    path = os.path.abspath(os.path.dirname(__file__))
+    name = os.path.basename(__file__)
+
+    import sys
+    sys.path.append(path)
+    os.environ["SQLALCHEMY_SILENCE_UBER_WARNING"] = "1"
+
+    import runpy
+    runpy.run_module(f"bemani.utils.{name}", run_name="__main__")
diff --git a/verifytyping b/verifytyping
index 6ba1f8a..2fbe9c0 100755
--- a/verifytyping
+++ b/verifytyping
@@ -22,6 +22,7 @@ declare -a arr=(
     "scheduler"
     "services"
     "struct"
+    "tdxtutils"
     "trafficgen"
     "twodxutils"
 )