support SDGS encryption

titles/chuni/index.py

@@ -39,7 +39,7 @@ class ChuniServlet(BaseServlet):
     def __init__(self, core_cfg: CoreConfig, cfg_dir: str) -> None:
         super().__init__(core_cfg, cfg_dir)
         self.game_cfg = ChuniConfig()
-        self.hash_table: Dict[Dict[str, str]] = {}
+        self.hash_table: Dict[str, Dict[str, str]] = {}
         if path.exists(f"{cfg_dir}/{ChuniConstants.CONFIG_NAME}"):
             self.game_cfg.update(
                 yaml.safe_load(open(f"{cfg_dir}/{ChuniConstants.CONFIG_NAME}"))
@@ -89,45 +89,59 @@ class ChuniServlet(BaseServlet):
             )
             self.logger.inited = True
 
+        known_iter_counts = {
+            ChuniConstants.VER_CHUNITHM_CRYSTAL_PLUS: 67,
+            ChuniConstants.VER_CHUNITHM_PARADISE: 44,
+            f"{ChuniConstants.VER_CHUNITHM_PARADISE}_int": 25,
+            ChuniConstants.VER_CHUNITHM_NEW: 54,
+            f"{ChuniConstants.VER_CHUNITHM_NEW}_int": 49,
+            ChuniConstants.VER_CHUNITHM_NEW_PLUS: 25,
+            ChuniConstants.VER_CHUNITHM_SUN: 70,
+            ChuniConstants.VER_CHUNITHM_SUN_PLUS: 36,
+        }
+
         for version, keys in self.game_cfg.crypto.keys.items():
             if len(keys) < 3:
                 continue
 
-            self.hash_table[version] = {}
+            if isinstance(version, int):
+                version_idx = version
+            else:
+                version_idx = int(version.split("_")[0])
+            
+            salt = bytes.fromhex(keys[2])
 
+            if len(keys) >= 4:
+                iter_count = keys[3]
+            elif (iter_count := known_iter_counts.get(version)) is None:
+                self.logger.error(
+                    "Number of iteration rounds for version %s is not known, but it is not specified in the config",
+                    version,
+                )
+                continue
+
+            self.hash_table[version] = {}
             method_list = [
                 method
-                for method in dir(self.versions[version])
+                for method in dir(self.versions[version_idx])
                 if not method.startswith("__")
             ]
+
             for method in method_list:
                 method_fixed = inflection.camelize(method)[6:-7]
-                # number of iterations was changed to 70 in SUN and then to 36
+
-                if version == ChuniConstants.VER_CHUNITHM_SUN_PLUS:
+                # This only applies for CHUNITHM NEW International and later for some reason.
-                    iter_count = 36
+                # CHUNITHM SUPERSTAR (PLUS) did not add "Exp" to the endpoint when hashing.
-                elif version == ChuniConstants.VER_CHUNITHM_SUN:
+                if (
-                    iter_count = 70
+                    isinstance(version, str)
-                elif version == ChuniConstants.VER_CHUNITHM_NEW_PLUS:
+                    and version.endswith("_int")
-                    iter_count = 25
+                    and version_idx >= ChuniConstants.VER_CHUNITHM_NEW
-                elif version == ChuniConstants.VER_CHUNITHM_NEW:
+                ):
-                    iter_count = 54
+                    method_fixed += "C3Exp"
-                elif version == ChuniConstants.VER_CHUNITHM_PARADISE:
-                    iter_count = 44
-                elif version == ChuniConstants.VER_CHUNITHM_CRYSTAL_PLUS:
-                    iter_count = 67
-                else:
-                    disp_version = (
-                        ChuniConstants.VERSION_NAMES[version]
-                        if version in ChuniConstants.VERSION_NAMES
-                        else str(version)
-                    )
-                    
-                    self.logger.warning("Version %s does not support encryption", disp_version)
-                    continue
                     
                 hash = PBKDF2(
                     method_fixed,
-                    bytes.fromhex(keys[2]),
+                    salt,
                     128,
                     count=iter_count,
                     hmac_hash_module=SHA1,
@@ -137,7 +151,7 @@ class ChuniServlet(BaseServlet):
                 self.hash_table[version][hashed_name] = method_fixed
 
                 self.logger.debug(
-                    f"Hashed v{version} method {method_fixed} with {bytes.fromhex(keys[2])} to get {hash.hex()}"
+                    f"Hashed v{version} method {method_fixed} with {salt} to get {hashed_name}"
                 )
 
     @classmethod
@@ -226,31 +240,39 @@ class ChuniServlet(BaseServlet):
 
         if all(c in string.hexdigits for c in endpoint) and len(endpoint) == 32:
             # If we get a 32 character long hex string, it's a hash and we're
-            # doing encrypted. The likelyhood of false positives is low but
+            # doing encrypted. The likelihood of false positives is low but
             # technically not 0
+
+            if game_code == "SDGS":
+                crypto_cfg_key = f"{internal_ver}_int"
+                hash_table_key = f"{internal_ver}_int"
+            else:
+                crypto_cfg_key = internal_ver
+                hash_table_key = internal_ver
+
             if internal_ver < ChuniConstants.VER_CHUNITHM_NEW:
                 endpoint = request.headers.get("User-Agent").split("#")[0]
 
             else:
-                if internal_ver not in self.hash_table:
+                if hash_table_key not in self.hash_table:
                     self.logger.error(
                         f"v{version} does not support encryption or no keys entered"
                     )
                     return Response(zlib.compress(b'{"stat": "0"}'))
 
-                elif endpoint.lower() not in self.hash_table[internal_ver]:
+                elif endpoint.lower() not in self.hash_table[hash_table_key]:
                     self.logger.error(
                         f"No hash found for v{version} endpoint {endpoint}"
                     )
                     return Response(zlib.compress(b'{"stat": "0"}'))
 
-                endpoint = self.hash_table[internal_ver][endpoint.lower()]
+                endpoint = self.hash_table[hash_table_key][endpoint.lower()]
 
             try:
                 crypt = AES.new(
-                    bytes.fromhex(self.game_cfg.crypto.keys[internal_ver][0]),
+                    bytes.fromhex(self.game_cfg.crypto.keys[crypto_cfg_key][0]),
                     AES.MODE_CBC,
-                    bytes.fromhex(self.game_cfg.crypto.keys[internal_ver][1]),
+                    bytes.fromhex(self.game_cfg.crypto.keys[crypto_cfg_key][1]),
                 )
 
                 req_raw = crypt.decrypt(req_raw)