artemis/core/adb_handlers/lookup.py

from typing import Union

from construct import Int8sl, Int32sl, Padding, Struct

from .base import *


class ADBLookupException(Exception):
    pass


class ADBLookupRequest(ADBBaseRequest):
    def __init__(self, data: bytes) -> None:
        super().__init__(data)
        self.access_code = data[0x20:0x2A].hex()
        company_code, fw_version, self.serial_number = struct.unpack_from(
            "<bbI", data, 0x2A
        )

        try:
            self.company_code = CompanyCodes(company_code)
        except ValueError as e:
            raise ADBLookupException(f"Invalid company code - {e}")

        self.fw_version = ReaderFwVer.from_byte(fw_version)


class ADBLookupResponse(ADBBaseResponse):
    def __init__(
        self,
        user_id: Union[int, None],
        game_id: str = "SXXX",
        store_id: int = 1,
        keychip_id: str = "A69E01A8888",
        code: int = 0x06,
        length: int = 0x30,
        status: int = 1,
    ) -> None:
        super().__init__(code, length, status, game_id, store_id, keychip_id)
        self.user_id = user_id if user_id is not None else -1
        self.portal_reg = PortalRegStatus.NO_REG

    @classmethod
    def from_req(cls, req: ADBHeader, user_id: Union[int, None]) -> "ADBLookupResponse":
        c = cls(user_id, req.game_id, req.store_id, req.keychip_id)
        c.head.protocol_ver = req.protocol_ver
        return c

    def make(self):
        resp_struct = Struct("user_id" / Int32sl, "portal_reg" / Int8sl, Padding(11))

        body = resp_struct.build(
            dict(user_id=self.user_id, portal_reg=self.portal_reg.value)
        )

        self.head.length = HEADER_SIZE + len(body)
        return self.head.make() + body


class ADBLookupExResponse(ADBBaseResponse):
    def __init__(
        self,
        user_id: Union[int, None],
        game_id: str = "SXXX",
        store_id: int = 1,
        keychip_id: str = "A69E01A8888",
        code: int = 0x10,
        length: int = 0x130,
        status: int = 1,
    ) -> None:
        super().__init__(code, length, status, game_id, store_id, keychip_id)
        self.user_id = user_id if user_id is not None else -1
        self.portal_reg = PortalRegStatus.NO_REG
        self.auth_key = [0] * 256

    @classmethod
    def from_req(
        cls, req: ADBHeader, user_id: Union[int, None]
    ) -> "ADBLookupExResponse":
        c = cls(user_id, req.game_id, req.store_id, req.keychip_id)
        c.head.protocol_ver = req.protocol_ver
        return c

    def make(self):
        resp_struct = Struct(
            "user_id" / Int32sl,
            "portal_reg" / Int8sl,
            Padding(3),
            "auth_key" / Int8sl[256],
            "relation1" / Int32sl,
            "relation2" / Int32sl,
        )

        body = resp_struct.build(
            dict(
                user_id=self.user_id,
                portal_reg=self.portal_reg.value,
                auth_key=self.auth_key,
                relation1=-1,
                relation2=-1,
            )
        )

        self.head.length = HEADER_SIZE + len(body)
        return self.head.make() + body