From 421c2dc09b28e3de7e7a502e1a4dcd43394a7374 Mon Sep 17 00:00:00 2001
From: beerpsi <beerpsi@duck.com>
Date: Wed, 24 Jul 2024 08:07:18 +0700
Subject: [PATCH] idk

---
 chuni/calc_title_server_enc_key.py | 23 +++----------
 geki/calc_title_server_enc_key.py  | 52 +++++++++++++++---------------
 2 files changed, 30 insertions(+), 45 deletions(-)

diff --git a/chuni/calc_title_server_enc_key.py b/chuni/calc_title_server_enc_key.py
index d0dbca1..8ce8423 100644
--- a/chuni/calc_title_server_enc_key.py
+++ b/chuni/calc_title_server_enc_key.py
@@ -4,11 +4,8 @@ import struct
 import sys
 import re
 
-from Crypto.Hash import SHA1
 import pefile
 
-from Crypto.Protocol.KDF import PBKDF2
-
 def rva2offset(pe: pefile.PE, rva: int):
     for section in pe.sections:
         if section.contains_rva(rva):
@@ -41,16 +38,10 @@ if (pmatch := KEY_PASSWORD_RE.search(exe)) and (smatch := KEY_SALT_RE.search(exe
     soffset = rva2offset(pe, struct.unpack("<I", smatch.group("offset"))[0] - base_address)
     poffset_end = poffset + exe[poffset:].index(0)
 
-    password = exe[poffset:poffset_end].decode("utf-8")
+    password = exe[poffset:poffset_end]
     salt = exe[soffset:soffset + 16]
 
-    key = PBKDF2(
-        password,
-        salt,
-        dkLen=32,
-        count=smatch.group("iterations")[0],
-        hmac_hash_module=SHA1,
-    )
+    key = pbkdf2_hmac("sha1", password, salt, iterations=smatch.group("iterations")[0], dklen=32)
     key = bytes((x % 0x5E) + 0x21 for x in key)
 
     print(f"Key: {key.hex()}")
@@ -71,16 +62,10 @@ if (pmatch := SALT_PASSWORD_RE.search(exe)) and (smatch := SALT_SALT_RE.search(e
     soffset = rva2offset(pe, struct.unpack("<I", smatch.group("offset"))[0] - base_address)
     poffset_end = poffset + exe[poffset:].index(0)
 
-    password = exe[poffset:poffset_end].decode("utf-8")
+    password = exe[poffset:poffset_end]
     salt = exe[soffset:soffset + 16]
 
-    key = PBKDF2(
-        password,
-        salt,
-        dkLen=8,
-        count=smatch.group("iterations")[0],
-        hmac_hash_module=SHA1,
-    )
+    key = pbkdf2_hmac("sha1", password, salt, iterations=smatch.group("iterations")[0], dklen=8)
     print(f"Endpoint salt: {key.hex()}")
 
 if (match := ITER_COUNT_RE.search(exe)):
diff --git a/geki/calc_title_server_enc_key.py b/geki/calc_title_server_enc_key.py
index 6af1136..2681c28 100644
--- a/geki/calc_title_server_enc_key.py
+++ b/geki/calc_title_server_enc_key.py
@@ -1,6 +1,10 @@
+import argparse
 from hashlib import blake2b
-import os
-import sys
+from pathlib import Path
+from typing import cast
+
+from UnityPy import Environment
+from UnityPy.classes import TextAsset
 
 # MU3.Sys.System.keyIvDigestFixed
 KEY_IV_DIGEST_FIXED = bytes([
@@ -9,34 +13,30 @@ KEY_IV_DIGEST_FIXED = bytes([
     57, 72, 120, 103, 17, 124, 18, 64, 15, 225,
     169, 39
 ])
+FILE_MAP = {
+    "noise0": "Key",
+    "noise1": "IV",
+    "noise2": "Endpoint salt",
+}
 
-if len(sys.argv) < 4:
-    print(f"Usage: python {os.path.basename(__file__)} <noise0> <noise1> <noise2>")
-    exit(1)
+parser = argparse.ArgumentParser()
+_ = parser.add_argument("mu3_data_path", type=Path)
+_ = parser.add_argument("key_iv_digest_fixed", type=bytes.fromhex, nargs="?", default=KEY_IV_DIGEST_FIXED)
 
-with open(sys.argv[1], "rb") as f:
-    key = blake2b(
-        f.read(),
-        digest_size=32,
-        key=KEY_IV_DIGEST_FIXED,
-    ).hexdigest()
+args = parser.parse_args()
 
-    print(f"Key: {key}")
+env = Environment()
 
-with open(sys.argv[2], "rb") as f:
-    iv = blake2b(
-        f.read(),
-        digest_size=32,
-        key=KEY_IV_DIGEST_FIXED,
-    ).hexdigest()
+with (cast(Path, args.mu3_data_path) / "resources.assets").open("rb") as f:
+    env.load_file(f)
 
-    print(f"IV: {iv}")
+    for object in env.objects:
+        if object.type.name != "TextAsset":
+            continue
 
-with open(sys.argv[3], "rb") as f:
-    salt = blake2b(
-        f.read(),
-        digest_size=32,
-        key=KEY_IV_DIGEST_FIXED,
-    ).hexdigest()
+        data = cast(TextAsset, object.read())
 
-    print(f"Endpoint salt: {salt}")
+        if data.name in FILE_MAP:
+            key = blake2b(data.m_Script, digest_size=32, key=args.key_iv_digest_fixed).hexdigest()
+
+            print(f"{FILE_MAP[data.name]}: {key}")