From dd546dcce2a4ada905371f92014af4c1a5efd7d2 Mon Sep 17 00:00:00 2001
From: Hay1tsme <pitok236@gmail.com>
Date: Sat, 20 Sep 2025 16:07:30 -0400
Subject: [PATCH] idz: pretty up no key message, fix double handshake

---
 titles/idz/index.py  | 30 +++++++++++++++++++++++++++---
 titles/idz/userdb.py | 35 +++++++++++++++++++----------------
 2 files changed, 46 insertions(+), 19 deletions(-)

diff --git a/titles/idz/index.py b/titles/idz/index.py
index 0ff5b8d..a9d6775 100644
--- a/titles/idz/index.py
+++ b/titles/idz/index.py
@@ -22,6 +22,7 @@ class IDZServlet(BaseServlet):
     def __init__(self, core_cfg: CoreConfig, cfg_dir: str) -> None:
         super().__init__(core_cfg, cfg_dir)
         self.game_cfg = IDZConfig()
+        self.rsa_keys: List[IDZKey] = []
         if path.exists(f"{cfg_dir}/{IDZConstants.CONFIG_NAME}"):
             self.game_cfg.update(
                 yaml.safe_load(open(f"{cfg_dir}/{IDZConstants.CONFIG_NAME}"))
@@ -38,8 +39,6 @@ class IDZServlet(BaseServlet):
                 backupCount=10,
             )
 
-            self.rsa_keys: List[IDZKey] = []
-
             fileHandler.setFormatter(log_fmt)
 
             consoleHandler = logging.StreamHandler()
@@ -79,7 +78,32 @@ class IDZServlet(BaseServlet):
             return False
 
         if len(game_cfg.rsa_keys) <= 0 or not game_cfg.server.aes_key:
-            logging.getLogger("idz").error("IDZ: No RSA/AES  keys! IDZ cannot start")
+            logger = logging.getLogger("idz")
+            if not hasattr(logger, "inited"):
+                log_fmt_str = "[%(asctime)s] IDZ | %(levelname)s | %(message)s"
+                log_fmt = logging.Formatter(log_fmt_str)
+                fileHandler = TimedRotatingFileHandler(
+                    "{0}/{1}.log".format(core_cfg.server.log_dir, "idz"),
+                    encoding="utf8",
+                    when="d",
+                    backupCount=10,
+                )
+
+                fileHandler.setFormatter(log_fmt)
+
+                consoleHandler = logging.StreamHandler()
+                consoleHandler.setFormatter(log_fmt)
+
+                logger.addHandler(fileHandler)
+                logger.addHandler(consoleHandler)
+
+                logger.setLevel(game_cfg.server.loglevel)
+                coloredlogs.install(
+                    level=game_cfg.server.loglevel, logger=logger, fmt=log_fmt_str
+                )
+                logger.inited = True
+            
+            logger.error("No RSA/AES  keys! IDZ cannot start")
             return False
 
         return True
diff --git a/titles/idz/userdb.py b/titles/idz/userdb.py
index 089778a..83aeabd 100644
--- a/titles/idz/userdb.py
+++ b/titles/idz/userdb.py
@@ -53,28 +53,31 @@ class IDZUserDB:
     
     async def connection_cb(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
         self.logger.debug(f"Connection made from {writer.get_extra_info('peername')[0]}")
+        sent_handshake = False
         while True:
             try:
-                base = 0
+                if not sent_handshake:
+                    base = 0
 
-                for i in range(len(self.static_key) - 1):
-                    shift = 8 * i
-                    byte = self.static_key[i]
+                    for i in range(len(self.static_key) - 1):
+                        shift = 8 * i
+                        byte = self.static_key[i]
 
-                    base |= byte << shift
+                        base |= byte << shift
 
-                rsa_key = random.choice(self.rsa_keys)
-                key_enc: int = pow(base, rsa_key.e, rsa_key.N)
-                result = (
-                    key_enc.to_bytes(0x40, "little")
-                    + struct.pack("<I", 0x01020304)
-                    + rsa_key.hashN.to_bytes(4, "little")
-                )
+                    rsa_key = random.choice(self.rsa_keys)
+                    key_enc: int = pow(base, rsa_key.e, rsa_key.N)
+                    result = (
+                        key_enc.to_bytes(0x40, "little")
+                        + struct.pack("<I", 0x01020304)
+                        + rsa_key.hashN.to_bytes(4, "little")
+                    )
 
-                self.logger.debug(f"Send handshake {result.hex()}")
+                    self.logger.debug(f"Send handshake {result.hex()}")
 
-                writer.write(result)
-                await writer.drain()
+                    writer.write(result)
+                    await writer.drain()
+                    sent_handshake = True
 
                 data: bytes = await reader.read(4096)
                 if len(data) == 0:
@@ -88,7 +91,7 @@ class IDZUserDB:
                 self.logger.debug("Connection reset, disconnecting")
                 return
 
-    def dataReceived(self, data: bytes, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
+    async def dataReceived(self, data: bytes, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
         self.logger.debug(f"Receive data {data.hex()}")
         client_ip = writer.get_extra_info('peername')[0]
         crypt = AES.new(self.static_key, AES.MODE_ECB)