FGO/artemis
3
2
Fork 0
forked from Dniel97/artemis
Code Pull Requests Activity

idz: add socket servers

Browse Source
This commit is contained in:
Hay1tsme 2024-01-09 12:21:31 -05:00
parent 4b9db8be3b
commit f65aa4d60a
4 changed files with 77 additions and 114 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

19
titles/idz/echo.py
View File

@ -4,16 +4,11 @@ import logging
from core.config import CoreConfig
from .config import IDZConfig
class IDZEcho:
def connection_made(self, transport):
self.transport = transport
class IDZEcho(DatagramProtocol):
def __init__(self, cfg: CoreConfig, game_cfg: IDZConfig) -> None:
super().__init__()
self.core_config = cfg
self.game_config = game_cfg
self.logger = logging.getLogger("idz")
def datagramReceived(self, data, addr):
self.logger.debug(
f"Echo from from {addr[0]}:{addr[1]} -> {self.transport.getHost().port} - {data.hex()}"
)
self.transport.write(data, addr)
def datagram_received(self, data, addr):
message = data.decode()
self.logger.debug(f'Received echo from {addr}')
self.transport.sendto(data, addr)

2
titles/idz/handlers/load_server_info.py
View File

@ -91,7 +91,7 @@ class IDZHandlerLoadServerInfo(IDZHandlerBase):
)
struct.pack_into("<I", ret, 0x39C + offset, self.game_cfg.ports.echo)
struct.pack_into("<I", ret, 0x39E + offset, self.game_cfg.ports.echo + 1)
struct.pack_into("<I", ret, 0x39E + offset, self.game_cfg.ports.echo) # TODO: Test
struct.pack_into(f"{len_news}s", ret, 0x03A0 + offset, news_str.encode())
struct.pack_into(f"{len_error}s", ret, 0x0424 + offset, err_str.encode())

41
titles/idz/index.py
View File

@ -8,12 +8,13 @@ from logging.handlers import TimedRotatingFileHandler
from os import path
from typing import Tuple, List, Dict
import importlib
import asyncio
from core.config import CoreConfig
from core.title import BaseServlet
from .config import IDZConfig
from .const import IDZConstants
from .userdb import IDZUserDBFactory, IDZKey
from .userdb import IDZUserDB, IDZKey
from .echo import IDZEcho
@ -135,27 +136,27 @@ class IDZServlet(BaseServlet):
except AttributeError as e:
continue
"""
endpoints.serverFromString(
reactor,
f"tcp:{self.game_cfg.ports.userdb}:interface={self.core_cfg.server.listen_address}",
).listen(
IDZUserDBFactory(self.core_cfg, self.game_cfg, self.rsa_keys, handler_map)
loop = asyncio.get_running_loop()
IDZUserDB(self.core_cfg, self.game_cfg, self.rsa_keys, handler_map)
asyncio.create_task(
loop.create_datagram_endpoint(
lambda: IDZEcho(),
local_addr=(self.core_cfg.server.listen_address, self.game_cfg.ports.echo)
)
)
asyncio.create_task(
loop.create_datagram_endpoint(
lambda: IDZEcho(),
local_addr=(self.core_cfg.server.listen_address, self.game_cfg.ports.match)
)
)
asyncio.create_task(
loop.create_datagram_endpoint(
lambda: IDZEcho(),
local_addr=(self.core_cfg.server.listen_address, self.game_cfg.ports.userdb + 1)
)
)
reactor.listenUDP(
self.game_cfg.ports.echo, IDZEcho(self.core_cfg, self.game_cfg)
)
reactor.listenUDP(
self.game_cfg.ports.echo + 1, IDZEcho(self.core_cfg, self.game_cfg)
)
reactor.listenUDP(
self.game_cfg.ports.match, IDZEcho(self.core_cfg, self.game_cfg)
)
reactor.listenUDP(
self.game_cfg.ports.userdb + 1, IDZEcho(self.core_cfg, self.game_cfg)
)
"""
self.logger.info(f"UserDB Listening on port {self.game_cfg.ports.userdb}")
async def render_GET(self, request: Request) -> bytes:

101
titles/idz/userdb.py
View File

@ -1,15 +1,9 @@
from twisted.internet.protocol import Factory, Protocol
import logging, coloredlogs
import logging
from Crypto.Cipher import AES
import struct
from typing import Dict, Optional, List, Type
from twisted.web import server, resource
from twisted.internet import reactor, endpoints
from starlette.requests import Request
from routes import Mapper
import random
from os import walk
import importlib
import asyncio
from core.config import CoreConfig
from .database import IDZData
@ -28,7 +22,7 @@ class IDZKey:
self.hashN = hashN
class IDZUserDBProtocol(Protocol):
class IDZUserDB:
def __init__(
self,
core_cfg: CoreConfig,
@ -46,6 +40,10 @@ class IDZUserDBProtocol(Protocol):
self.version_internal = None
self.skip_next = False
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))
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)
@ -53,8 +51,10 @@ class IDZUserDBProtocol(Protocol):
data += bytes(padding_size)
return data
def connectionMade(self) -> None:
self.logger.debug(f"{self.transport.getPeer().host} Connected")
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):
@ -73,22 +73,31 @@ class IDZUserDBProtocol(Protocol):
self.logger.debug(f"Send handshake {result.hex()}")
self.transport.write(result)
writer.write(result)
await writer.drain()
def connectionLost(self, reason) -> None:
self.logger.debug(
f"{self.transport.getPeer().host} Disconnected - {reason.value}"
)
data: bytes = await reader.read(4096)
if len(data) == 0:
self.logger.debug("Connection closed")
return
def dataReceived(self, data: bytes) -> None:
await self.dataReceived(data, reader, writer)
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:
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)
try:
data_dec = crypt.decrypt(data)
except Exception as e:
self.logger.error(f"Failed to decrypt UserDB request from {self.transport.getPeer().host} because {e} - {data.hex()}")
self.logger.error(f"Failed to decrypt UserDB request from {client_ip} because {e} - {data.hex()}")
self.logger.debug(f"Decrypt data {data_dec.hex()}")
@ -99,7 +108,7 @@ class IDZUserDBProtocol(Protocol):
self.logger.info(f"Userdb serverbox request {data_dec.hex()}")
self.skip_next = True
self.transport.write(b"\x00")
writer.write(b"\x00")
return
elif magic == 0x01020304:
@ -119,21 +128,21 @@ class IDZUserDBProtocol(Protocol):
self.version_internal = None
self.logger.debug(
f"Userdb v{self.version} handshake response from {self.transport.getPeer().host}"
f"Userdb v{self.version} handshake response from {client_ip}"
)
return
elif self.skip_next:
self.skip_next = False
self.transport.write(b"\x00")
writer.write(b"\x00")
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(
f"Bad UserDB request from from {self.transport.getPeer().host}"
f"Bad UserDB request from from {client_ip}"
)
self.transport.write(b"\x00")
writer.write(b"\x00")
return
cmd = struct.unpack_from("<H", data_dec, 0)[0]
@ -147,7 +156,7 @@ class IDZUserDBProtocol(Protocol):
handler = handler_cls(self.core_config, self.game_config, self.version_internal)
self.logger.info(
f"Userdb v{self.version} {handler.name} request from {self.transport.getPeer().host}"
f"Userdb v{self.version} {handler.name} request from {client_ip}"
)
response = handler.handle(data_dec)
@ -156,46 +165,4 @@ class IDZUserDBProtocol(Protocol):
crypt = AES.new(self.static_key, AES.MODE_ECB)
response_enc = crypt.encrypt(response)
self.transport.write(response_enc)
class IDZUserDBFactory(Factory):
protocol = IDZUserDBProtocol
def __init__(
self,
cfg: CoreConfig,
game_cfg: IDZConfig,
keys: List[IDZKey],
handlers: List[Dict],
) -> None:
self.core_config = cfg
self.game_config = game_cfg
self.keys = keys
self.handlers = handlers
def buildProtocol(self, addr):
return IDZUserDBProtocol(
self.core_config, self.game_config, self.keys, self.handlers
)
class IDZUserDBWeb(resource.Resource):
def __init__(self, core_cfg: CoreConfig, game_cfg: IDZConfig):
super().__init__()
self.isLeaf = True
self.core_config = core_cfg
self.game_config = game_cfg
self.logger = logging.getLogger("idz")
def render_POST(self, request: Request) -> bytes:
self.logger.info(
f"IDZUserDBWeb POST from {request.getClientAddress().host} to {request.uri} with data {request.content.getvalue()}"
)
return b""
def render_GET(self, request: Request) -> bytes:
self.logger.info(
f"IDZUserDBWeb GET from {request.getClientAddress().host} to {request.uri}"
)
return b""
writer.write(response_enc)