1
0
mirror of synced 2024-12-19 08:45:53 +01:00
artemis/titles/idz/userdb.py

169 lines
5.6 KiB
Python
Raw Normal View History

2024-01-09 18:21:31 +01:00
import logging
2023-04-23 10:38:28 +02:00
from Crypto.Cipher import AES
import struct
from typing import Dict, Optional, List, Type
import random
2024-01-09 18:21:31 +01:00
import asyncio
2023-04-23 10:38:28 +02:00
from core.config import CoreConfig
from .database import IDZData
from .config import IDZConfig
from .const import IDZConstants
from .handlers import IDZHandlerBase
HANDLER_MAP: List[Dict]
2023-04-23 10:38:28 +02:00
class IDZKey:
def __init__(self, n, d, e, hashN: int) -> None:
self.N = n
self.d = d
self.e = e
self.hashN = hashN
2024-01-09 18:21:31 +01:00
class IDZUserDB:
def __init__(
self,
core_cfg: CoreConfig,
game_cfg: IDZConfig,
keys: List[IDZKey],
handlers: List[Dict],
) -> None:
self.logger = logging.getLogger("idz")
2023-04-23 10:38:28 +02:00
self.core_config = core_cfg
self.game_config = game_cfg
self.rsa_keys = keys
self.handlers = handlers
self.static_key = bytes.fromhex(self.game_config.server.aes_key)
self.version = None
self.version_internal = None
self.skip_next = False
2024-01-09 18:21:31 +01:00
def start(self) -> None:
self.logger.info(f"Start on port {self.config.aimedb.port}")
asyncio.create_task(asyncio.start_server(self.connection_cb, self.config.server.listen_address, self.config.aimedb.port))
2023-04-23 10:38:28 +02:00
def append_padding(self, data: bytes):
"""Appends 0s to the end of the data until it's at the correct size"""
length = struct.unpack_from("<H", data, 6)
padding_size = length[0] - len(data)
data += bytes(padding_size)
return data
2024-01-09 18:21:31 +01:00
async def connection_cb(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
self.logger.debug(f"Connection made from {writer.get_extra_info('peername')[0]}")
while True:
try:
base = 0
for i in range(len(self.static_key) - 1):
shift = 8 * i
byte = self.static_key[i]
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")
)
self.logger.debug(f"Send handshake {result.hex()}")
writer.write(result)
await writer.drain()
data: bytes = await reader.read(4096)
if len(data) == 0:
self.logger.debug("Connection closed")
return
2024-01-09 20:42:17 +01:00
await self.data.Received(data, reader, writer)
2024-01-09 18:21:31 +01:00
await writer.drain()
except ConnectionResetError as e:
self.logger.debug("Connection reset, disconnecting")
return
def dataReceived(self, data: bytes, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> None:
2023-04-23 10:38:28 +02:00
self.logger.debug(f"Receive data {data.hex()}")
2024-01-09 18:21:31 +01:00
client_ip = writer.get_extra_info('peername')[0]
2023-04-23 10:38:28 +02:00
crypt = AES.new(self.static_key, AES.MODE_ECB)
try:
data_dec = crypt.decrypt(data)
except Exception as e:
2024-01-09 18:21:31 +01:00
self.logger.error(f"Failed to decrypt UserDB request from {client_ip} because {e} - {data.hex()}")
2023-04-23 10:38:28 +02:00
self.logger.debug(f"Decrypt data {data_dec.hex()}")
magic = struct.unpack_from("<I", data_dec, 0)[0]
if magic == 0xFE78571D:
# Ignore
self.logger.info(f"Userdb serverbox request {data_dec.hex()}")
self.skip_next = True
2024-01-09 18:21:31 +01:00
writer.write(b"\x00")
2023-04-23 10:38:28 +02:00
return
elif magic == 0x01020304:
2023-04-23 10:38:28 +02:00
self.version = int(data_dec[16:19].decode())
if self.version == 110:
self.version_internal = IDZConstants.VER_IDZ_110
elif self.version == 130:
self.version_internal = IDZConstants.VER_IDZ_130
elif self.version == 210:
self.version_internal = IDZConstants.VER_IDZ_210
elif self.version == 230:
self.version_internal = IDZConstants.VER_IDZ_230
else:
2023-08-08 16:17:56 +02:00
self.logger.warning(f"Bad version v{self.version}")
2023-04-23 10:38:28 +02:00
self.version = None
self.version_internal = None
self.logger.debug(
2024-01-09 18:21:31 +01:00
f"Userdb v{self.version} handshake response from {client_ip}"
)
2023-04-23 10:38:28 +02:00
return
2023-04-23 10:38:28 +02:00
elif self.skip_next:
self.skip_next = False
2024-01-09 18:21:31 +01:00
writer.write(b"\x00")
2023-04-23 10:38:28 +02:00
return
elif self.version is None:
# We didn't get a handshake before, and this isn't one now, so we're up the creek
self.logger.info(
2024-01-09 18:21:31 +01:00
f"Bad UserDB request from from {client_ip}"
)
2024-01-09 18:21:31 +01:00
writer.write(b"\x00")
2023-04-23 10:38:28 +02:00
return
cmd = struct.unpack_from("<H", data_dec, 0)[0]
handler_cls: Optional[Type[IDZHandlerBase]] = self.handlers[
self.version_internal
].get(cmd, None)
2023-04-23 10:38:28 +02:00
if handler_cls is None:
2023-08-08 16:17:56 +02:00
self.logger.warning(f"No handler for v{self.version} {hex(cmd)} cmd")
2023-04-23 10:38:28 +02:00
handler_cls = IDZHandlerBase
handler = handler_cls(self.core_config, self.game_config, self.version_internal)
self.logger.info(
2024-01-09 18:21:31 +01:00
f"Userdb v{self.version} {handler.name} request from {client_ip}"
)
2023-04-23 10:38:28 +02:00
response = handler.handle(data_dec)
self.logger.debug(f"Response: {response.hex()}")
crypt = AES.new(self.static_key, AES.MODE_ECB)
response_enc = crypt.encrypt(response)
2024-01-09 18:21:31 +01:00
writer.write(response_enc)