begin move
This commit is contained in:
178
core/aimedb.py
178
core/aimedb.py
@ -2,8 +2,8 @@ from twisted.internet.protocol import Factory, Protocol
|
||||
import logging, coloredlogs
|
||||
from Crypto.Cipher import AES
|
||||
import struct
|
||||
from typing import Dict, Tuple, Callable, Union
|
||||
from typing_extensions import Final
|
||||
from typing import Dict, Tuple, Callable, Union, Optional
|
||||
import asyncio
|
||||
from logging.handlers import TimedRotatingFileHandler
|
||||
|
||||
from core.config import CoreConfig
|
||||
@ -11,15 +11,37 @@ from core.utils import create_sega_auth_key
|
||||
from core.data import Data
|
||||
from .adb_handlers import *
|
||||
|
||||
|
||||
class AimedbProtocol(Protocol):
|
||||
class AimedbServlette():
|
||||
request_list: Dict[int, Tuple[Callable[[bytes, int], Union[ADBBaseResponse, bytes]], int, str]] = {}
|
||||
|
||||
def __init__(self, core_cfg: CoreConfig) -> None:
|
||||
self.logger = logging.getLogger("aimedb")
|
||||
self.config = core_cfg
|
||||
def __init__(self, core_cfg: CoreConfig) -> None:
|
||||
self.config = core_cfg
|
||||
self.data = Data(core_cfg)
|
||||
if core_cfg.aimedb.key == "":
|
||||
|
||||
self.logger = logging.getLogger("aimedb")
|
||||
if not hasattr(self.logger, "initted"):
|
||||
log_fmt_str = "[%(asctime)s] Aimedb | %(levelname)s | %(message)s"
|
||||
log_fmt = logging.Formatter(log_fmt_str)
|
||||
|
||||
fileHandler = TimedRotatingFileHandler(
|
||||
"{0}/{1}.log".format(self.config.server.log_dir, "aimedb"),
|
||||
when="d",
|
||||
backupCount=10,
|
||||
)
|
||||
fileHandler.setFormatter(log_fmt)
|
||||
|
||||
consoleHandler = logging.StreamHandler()
|
||||
consoleHandler.setFormatter(log_fmt)
|
||||
|
||||
self.logger.addHandler(fileHandler)
|
||||
self.logger.addHandler(consoleHandler)
|
||||
|
||||
self.logger.setLevel(self.config.aimedb.loglevel)
|
||||
coloredlogs.install(
|
||||
level=core_cfg.aimedb.loglevel, logger=self.logger, fmt=log_fmt_str
|
||||
)
|
||||
self.logger.initted = True
|
||||
|
||||
if not core_cfg.aimedb.key:
|
||||
self.logger.error("!!!KEY NOT SET!!!")
|
||||
exit(1)
|
||||
|
||||
@ -40,27 +62,30 @@ class AimedbProtocol(Protocol):
|
||||
|
||||
self.register_handler(0x13, 0x14, self.handle_log_ex, 'aime_log_ex')
|
||||
self.register_handler(0x64, 0x65, self.handle_hello, 'hello')
|
||||
self.register_handler(0x66, 0, self.handle_goodbye, 'goodbye')
|
||||
|
||||
|
||||
def register_handler(self, cmd: int, resp:int, handler: Callable[[bytes, int], Union[ADBBaseResponse, bytes]], name: str) -> None:
|
||||
self.request_list[cmd] = (handler, resp, name)
|
||||
|
||||
def start(self) -> None:
|
||||
self.logger.info(f"Start on port {self.config.aimedb.port}")
|
||||
asyncio.create_task(asyncio.start_server(self.dataReceived, self.config.server.listen_address, self.config.aimedb.port))
|
||||
|
||||
async def dataReceived(self, reader: asyncio.StreamReader, writer: asyncio.StreamWriter):
|
||||
self.logger.debug(f"Connection made from {writer.get_extra_info('peername')[0]}")
|
||||
while True:
|
||||
try:
|
||||
data: bytes = await reader.read()
|
||||
if len(data) == 0:
|
||||
self.logger.debug("Connection closed")
|
||||
return
|
||||
await self.process_data(data, reader, writer)
|
||||
await writer.drain()
|
||||
except ConnectionResetError as e:
|
||||
self.logger.warn("Connection reset, disconnecting")
|
||||
return
|
||||
|
||||
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
|
||||
|
||||
def connectionMade(self) -> None:
|
||||
self.logger.debug(f"{self.transport.getPeer().host} Connected")
|
||||
|
||||
def connectionLost(self, reason) -> None:
|
||||
self.logger.debug(
|
||||
f"{self.transport.getPeer().host} Disconnected - {reason.value}"
|
||||
)
|
||||
|
||||
def dataReceived(self, data: bytes) -> None:
|
||||
async def process_data(self, data: bytes, reader: asyncio.StreamReader, writer: asyncio.StreamWriter) -> Optional[bytes]:
|
||||
addr = writer.get_extra_info('peername')[0]
|
||||
cipher = AES.new(self.config.aimedb.key.encode(), AES.MODE_ECB)
|
||||
|
||||
try:
|
||||
@ -68,9 +93,9 @@ class AimedbProtocol(Protocol):
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to decrypt {data.hex()} because {e}")
|
||||
return None
|
||||
return
|
||||
|
||||
self.logger.debug(f"{self.transport.getPeer().host} wrote {decrypted.hex()}")
|
||||
self.logger.debug(f"{addr} wrote {decrypted.hex()}")
|
||||
|
||||
try:
|
||||
head = ADBHeader.from_data(decrypted)
|
||||
@ -79,7 +104,9 @@ class AimedbProtocol(Protocol):
|
||||
self.logger.error(f"Error parsing ADB header: {e}")
|
||||
try:
|
||||
encrypted = cipher.encrypt(ADBBaseResponse().make())
|
||||
self.transport.write(encrypted)
|
||||
writer.write(encrypted)
|
||||
await writer.drain()
|
||||
return
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to encrypt default response because {e}")
|
||||
@ -89,46 +116,51 @@ class AimedbProtocol(Protocol):
|
||||
if head.keychip_id == "ABCD1234567" or head.store_id == 0xfff0:
|
||||
self.logger.warning(f"Request from uninitialized AMLib: {vars(head)}")
|
||||
|
||||
if head.cmd == 0x66:
|
||||
self.logger.info("Goodbye")
|
||||
writer.close()
|
||||
return
|
||||
|
||||
handler, resp_code, name = self.request_list.get(head.cmd, (self.handle_default, None, 'default'))
|
||||
|
||||
if resp_code is None:
|
||||
self.logger.warning(f"No handler for cmd {hex(head.cmd)}")
|
||||
|
||||
elif resp_code > 0:
|
||||
self.logger.info(f"{name} from {head.keychip_id} ({head.game_id}) @ {self.transport.getPeer().host}")
|
||||
self.logger.info(f"{name} from {head.keychip_id} ({head.game_id}) @ {addr}")
|
||||
|
||||
resp = handler(decrypted, resp_code)
|
||||
resp = await handler(decrypted, resp_code)
|
||||
|
||||
if type(resp) == ADBBaseResponse or issubclass(type(resp), ADBBaseResponse):
|
||||
resp_bytes = resp.make()
|
||||
if len(resp_bytes) != resp.head.length:
|
||||
resp_bytes = self.append_padding(resp_bytes)
|
||||
|
||||
elif type(resp) == bytes:
|
||||
resp_bytes = resp
|
||||
|
||||
elif resp is None: # Nothing to send, probably a goodbye
|
||||
self.logger.warn(f"None return by handler for {name}")
|
||||
return
|
||||
|
||||
else:
|
||||
self.logger.error(f"Unsupported type returned by ADB handler for {name}: {type(resp)}")
|
||||
raise TypeError(f"Unsupported type returned by ADB handler for {name}: {type(resp)}")
|
||||
|
||||
try:
|
||||
try:
|
||||
encrypted = cipher.encrypt(resp_bytes)
|
||||
self.logger.debug(f"Response {resp_bytes.hex()}")
|
||||
self.transport.write(encrypted)
|
||||
writer.write(encrypted)
|
||||
|
||||
except Exception as e:
|
||||
self.logger.error(f"Failed to encrypt {resp_bytes.hex()} because {e}")
|
||||
|
||||
def handle_default(self, data: bytes, resp_code: int, length: int = 0x20) -> ADBBaseResponse:
|
||||
|
||||
async def handle_default(self, data: bytes, resp_code: int, length: int = 0x20) -> ADBBaseResponse:
|
||||
req = ADBHeader.from_data(data)
|
||||
return ADBBaseResponse(resp_code, length, 1, req.game_id, req.store_id, req.keychip_id, req.protocol_ver)
|
||||
|
||||
def handle_hello(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
return self.handle_default(data, resp_code)
|
||||
async def handle_hello(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
return await self.handle_default(data, resp_code)
|
||||
|
||||
def handle_campaign(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
async def handle_campaign(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
h = ADBHeader.from_data(data)
|
||||
if h.protocol_ver >= 0x3030:
|
||||
req = h
|
||||
@ -143,26 +175,26 @@ class AimedbProtocol(Protocol):
|
||||
# We don't currently support campaigns
|
||||
return resp
|
||||
|
||||
def handle_lookup(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
async def handle_lookup(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
req = ADBLookupRequest(data)
|
||||
user_id = self.data.card.get_user_id_from_card(req.access_code)
|
||||
is_banned = self.data.card.get_card_banned(req.access_code)
|
||||
is_locked = self.data.card.get_card_locked(req.access_code)
|
||||
|
||||
|
||||
ret = ADBLookupResponse.from_req(req.head, user_id)
|
||||
if is_banned and is_locked:
|
||||
ret.head.status = ADBStatus.BAN_SYS_USER
|
||||
elif is_banned:
|
||||
ret.head.status = ADBStatus.BAN_SYS
|
||||
elif is_locked:
|
||||
ret.head.status = ADBStatus.LOCK_USER
|
||||
ret = ADBLookupResponse.from_req(req.head, user_id)
|
||||
|
||||
self.logger.info(
|
||||
f"access_code {req.access_code} -> user_id {ret.user_id}"
|
||||
)
|
||||
return ret
|
||||
|
||||
def handle_lookup_ex(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
async def handle_lookup_ex(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
req = ADBLookupRequest(data)
|
||||
user_id = self.data.card.get_user_id_from_card(req.access_code)
|
||||
|
||||
@ -191,7 +223,7 @@ class AimedbProtocol(Protocol):
|
||||
|
||||
return ret
|
||||
|
||||
def handle_felica_lookup(self, data: bytes, resp_code: int) -> bytes:
|
||||
async def handle_felica_lookup(self, data: bytes, resp_code: int) -> bytes:
|
||||
"""
|
||||
On official, I think a card has to be registered for this to actually work, but
|
||||
I'm making the executive decision to not implement that and just kick back our
|
||||
@ -207,7 +239,7 @@ class AimedbProtocol(Protocol):
|
||||
)
|
||||
return ADBFelicaLookupResponse.from_req(req.head, ac)
|
||||
|
||||
def handle_felica_register(self, data: bytes, resp_code: int) -> bytes:
|
||||
async def handle_felica_register(self, data: bytes, resp_code: int) -> bytes:
|
||||
"""
|
||||
I've never seen this used.
|
||||
"""
|
||||
@ -239,7 +271,7 @@ class AimedbProtocol(Protocol):
|
||||
|
||||
return ADBFelicaLookupResponse.from_req(req.head, ac)
|
||||
|
||||
def handle_felica_lookup_ex(self, data: bytes, resp_code: int) -> bytes:
|
||||
async def handle_felica_lookup_ex(self, data: bytes, resp_code: int) -> bytes:
|
||||
req = ADBFelicaLookup2Request(data)
|
||||
access_code = self.data.card.to_access_code(req.idm)
|
||||
user_id = self.data.card.get_user_id_from_card(access_code=access_code)
|
||||
@ -263,7 +295,7 @@ class AimedbProtocol(Protocol):
|
||||
|
||||
return resp
|
||||
|
||||
def handle_campaign_clear(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
async def handle_campaign_clear(self, data: bytes, resp_code: int) -> ADBBaseResponse:
|
||||
req = ADBCampaignClearRequest(data)
|
||||
|
||||
resp = ADBCampaignClearResponse.from_req(req.head)
|
||||
@ -271,7 +303,7 @@ class AimedbProtocol(Protocol):
|
||||
# We don't support campaign stuff
|
||||
return resp
|
||||
|
||||
def handle_register(self, data: bytes, resp_code: int) -> bytes:
|
||||
async def handle_register(self, data: bytes, resp_code: int) -> bytes:
|
||||
req = ADBLookupRequest(data)
|
||||
user_id = -1
|
||||
|
||||
@ -305,17 +337,17 @@ class AimedbProtocol(Protocol):
|
||||
return resp
|
||||
|
||||
# TODO: Save these in some capacity, as deemed relevant
|
||||
def handle_status_log(self, data: bytes, resp_code: int) -> bytes:
|
||||
async def handle_status_log(self, data: bytes, resp_code: int) -> bytes:
|
||||
req = ADBStatusLogRequest(data)
|
||||
self.logger.info(f"User {req.aime_id} logged {req.status.name} event")
|
||||
return ADBBaseResponse(resp_code, 0x20, 1, req.head.game_id, req.head.store_id, req.head.keychip_id, req.head.protocol_ver)
|
||||
|
||||
def handle_log(self, data: bytes, resp_code: int) -> bytes:
|
||||
async def handle_log(self, data: bytes, resp_code: int) -> bytes:
|
||||
req = ADBLogRequest(data)
|
||||
self.logger.info(f"User {req.aime_id} logged {req.status.name} event, credit_ct: {req.credit_ct} bet_ct: {req.bet_ct} won_ct: {req.won_ct}")
|
||||
return ADBBaseResponse(resp_code, 0x20, 1, req.head.game_id, req.head.store_id, req.head.keychip_id, req.head.protocol_ver)
|
||||
|
||||
def handle_log_ex(self, data: bytes, resp_code: int) -> bytes:
|
||||
async def handle_log_ex(self, data: bytes, resp_code: int) -> bytes:
|
||||
req = ADBLogExRequest(data)
|
||||
strs = []
|
||||
self.logger.info(f"Recieved {req.num_logs} or {len(req.logs)} logs")
|
||||
@ -324,43 +356,3 @@ class AimedbProtocol(Protocol):
|
||||
self.logger.debug(f"User {req.logs[x].aime_id} logged {req.logs[x].status.name} event, credit_ct: {req.logs[x].credit_ct} bet_ct: {req.logs[x].bet_ct} won_ct: {req.logs[x].won_ct}")
|
||||
return ADBLogExResponse.from_req(req.head)
|
||||
|
||||
def handle_goodbye(self, data: bytes, resp_code: int) -> None:
|
||||
self.logger.info(f"goodbye from {self.transport.getPeer().host}")
|
||||
self.transport.loseConnection()
|
||||
return
|
||||
|
||||
class AimedbFactory(Factory):
|
||||
protocol = AimedbProtocol
|
||||
|
||||
def __init__(self, cfg: CoreConfig) -> None:
|
||||
self.config = cfg
|
||||
log_fmt_str = "[%(asctime)s] Aimedb | %(levelname)s | %(message)s"
|
||||
log_fmt = logging.Formatter(log_fmt_str)
|
||||
self.logger = logging.getLogger("aimedb")
|
||||
|
||||
fileHandler = TimedRotatingFileHandler(
|
||||
"{0}/{1}.log".format(self.config.server.log_dir, "aimedb"),
|
||||
when="d",
|
||||
backupCount=10,
|
||||
)
|
||||
fileHandler.setFormatter(log_fmt)
|
||||
|
||||
consoleHandler = logging.StreamHandler()
|
||||
consoleHandler.setFormatter(log_fmt)
|
||||
|
||||
self.logger.addHandler(fileHandler)
|
||||
self.logger.addHandler(consoleHandler)
|
||||
|
||||
self.logger.setLevel(self.config.aimedb.loglevel)
|
||||
coloredlogs.install(
|
||||
level=cfg.aimedb.loglevel, logger=self.logger, fmt=log_fmt_str
|
||||
)
|
||||
|
||||
if self.config.aimedb.key == "":
|
||||
self.logger.error("Please set 'key' field in your config file.")
|
||||
exit(1)
|
||||
|
||||
self.logger.info(f"Ready on port {self.config.aimedb.port}")
|
||||
|
||||
def buildProtocol(self, addr):
|
||||
return AimedbProtocol(self.config)
|
||||
|
Reference in New Issue
Block a user