bemaniutils/bemani/client/popn/lapistoria.py

import random
import time
from typing import Any, Dict, List, Optional

from bemani.client.base import BaseClient
from bemani.protocol import Node


class PopnMusicLapistoriaClient(BaseClient):
    NAME = "ＴＥＳＴ"

    def verify_pcb22_boot(self) -> None:
        call = self.call_node()

        # Construct node
        pcb22 = Node.void("pcb22")
        call.add_child(pcb22)
        pcb22.set_attribute("method", "boot")

        # Swap with server
        resp = self.exchange("", call)

        # Verify that response is correct
        self.assert_path(resp, "response/pcb22/@status")

    def verify_info22_common(self) -> None:
        call = self.call_node()

        # Construct node
        info22 = Node.void("info22")
        call.add_child(info22)
        info22.set_attribute("loc_id", "JP-1")
        info22.set_attribute("method", "common")

        # Swap with server
        resp = self.exchange("", call)

        # Verify that response is correct
        self.assert_path(resp, "response/info22")

        for name in [
            "phase",
            "story",
        ]:
            node = resp.child("info22").child(name)

            if node is None:
                raise Exception(f"Missing node '{name}' in response!")
            if node.data_type != "void":
                raise Exception(f"Node '{name}' has wrong data type!")

    def verify_player22_read(
        self, ref_id: str, msg_type: str
    ) -> Optional[Dict[str, Any]]:
        call = self.call_node()

        # Construct node
        player22 = Node.void("player22")
        call.add_child(player22)
        player22.set_attribute("method", "read")

        player22.add_child(Node.string("ref_id", value=ref_id))
        player22.add_child(Node.string("shop_name", ""))
        player22.add_child(Node.s8("pref", 51))

        # Swap with server
        resp = self.exchange("", call)

        if msg_type == "new":
            # Verify that response is correct
            self.assert_path(resp, "response/player22/@status")

            status = int(resp.child("player22").attribute("status"))
            if status != 109:
                raise Exception(
                    f"Reference ID '{ref_id}' returned invalid status '{status}'"
                )

            # No score data
            return None
        elif msg_type == "query":
            # Verify that the response is correct
            self.assert_path(resp, "response/player22/account/name")
            self.assert_path(resp, "response/player22/account/g_pm_id")
            self.assert_path(resp, "response/player22/account/my_best")
            self.assert_path(resp, "response/player22/account/latest_music")
            self.assert_path(resp, "response/player22/netvs")
            self.assert_path(resp, "response/player22/config")
            self.assert_path(resp, "response/player22/option")
            self.assert_path(resp, "response/player22/info")
            self.assert_path(resp, "response/player22/custom_cate")
            self.assert_path(resp, "response/player22/customize")

            name = resp.child("player22").child("account").child("name").value
            if name != self.NAME:
                raise Exception(f"Invalid name '{name}' returned for Ref ID '{ref_id}'")

            # Extract and return score data
            medals: Dict[int, List[int]] = {}
            scores: Dict[int, List[int]] = {}
            courses: Dict[int, Dict[str, int]] = {}
            for child in resp.child("player22").children:
                if child.name == "music":
                    songid = child.child_value("music_num")
                    chart = child.child_value("sheet_num")
                    medal = child.child_value("clear_type")
                    points = child.child_value("score")

                    if songid not in medals:
                        medals[songid] = [0, 0, 0, 0]
                    medals[songid][chart] = medal
                    if songid not in scores:
                        scores[songid] = [0, 0, 0, 0]
                    scores[songid][chart] = points

                if child.name == "course":
                    courseid = child.child_value("course_id")
                    medal = child.child_value("clear_medal")
                    combo = child.child_value("max_cmbo")
                    stage1 = child.child_value("stage1_score")
                    stage2 = child.child_value("stage2_score")
                    stage3 = child.child_value("stage3_score")
                    stage4 = child.child_value("stage4_score")
                    total = child.child_value("total_score")
                    courses[courseid] = {
                        "id": courseid,
                        "medal": medal,
                        "combo": combo,
                        "stage1": stage1,
                        "stage2": stage2,
                        "stage3": stage3,
                        "stage4": stage4,
                        "total": total,
                    }

            return {"medals": medals, "scores": scores, "courses": courses}

        else:
            raise Exception(f"Unrecognized message type '{msg_type}'")

    def verify_player22_write(self, ref_id: str, scores: List[Dict[str, Any]]) -> None:
        call = self.call_node()

        # Construct node
        player22 = Node.void("player22")
        call.add_child(player22)
        player22.set_attribute("method", "write")
        player22.add_child(Node.string("ref_id", value=ref_id))

        # Add required children
        config = Node.void("config")
        player22.add_child(config)
        config.add_child(Node.s16("chara", value=1543))

        # Add requested scores
        for score in scores:
            stage = Node.void("stage")
            player22.add_child(stage)
            stage.add_child(Node.s16("no", score["id"]))
            stage.add_child(Node.u8("sheet", score["chart"]))
            stage.add_child(Node.u16("clearmedal", score["medal"]))
            stage.add_child(Node.s32("nscore", score["score"]))

        # Swap with server
        resp = self.exchange("", call)
        self.assert_path(resp, "response/player22/@status")

    def verify_player22_write_music(self, ref_id: str, score: Dict[str, Any]) -> None:
        call = self.call_node()

        # Construct node
        player22 = Node.void("player22")
        call.add_child(player22)
        player22.set_attribute("method", "write_music")
        player22.add_child(Node.string("ref_id", ref_id))
        player22.add_child(Node.string("name", self.NAME))
        player22.add_child(Node.u8("stage", 0))
        player22.add_child(Node.s16("music_num", score["id"]))
        player22.add_child(Node.u8("sheet_num", score["chart"]))
        player22.add_child(Node.u8("clearmedal", score["medal"]))
        player22.add_child(Node.s32("score", score["score"]))
        player22.add_child(Node.s16("combo", 0))
        player22.add_child(Node.s16("cool", 0))
        player22.add_child(Node.s16("great", 0))
        player22.add_child(Node.s16("good", 0))
        player22.add_child(Node.s16("bad", 0))

        # Swap with server
        resp = self.exchange("", call)
        self.assert_path(resp, "response/player22/@status")

    def verify_player22_write_course(self, ref_id: str, course: Dict[str, int]) -> None:
        call = self.call_node()

        # Construct node
        player22 = Node.void("player22")
        call.add_child(player22)
        player22.set_attribute("method", "write_course")
        player22.add_child(Node.s16("pref", 51))
        player22.add_child(Node.string("location_id", "JP-1"))
        player22.add_child(Node.string("ref_id", ref_id))
        player22.add_child(Node.string("data_id", ref_id))
        player22.add_child(Node.string("name", self.NAME))
        player22.add_child(Node.s16("chara_num", 1543))
        player22.add_child(Node.s32("play_id", 0))
        player22.add_child(Node.s16("course_id", course["id"]))
        player22.add_child(Node.s16("stage1_music_num", 148))
        player22.add_child(Node.u8("stage1_sheet_num", 1))
        player22.add_child(Node.s16("stage2_music_num", 550))
        player22.add_child(Node.u8("stage2_sheet_num", 1))
        player22.add_child(Node.s16("stage3_music_num", 1113))
        player22.add_child(Node.u8("stage3_sheet_num", 1))
        player22.add_child(Node.s16("stage4_music_num", 341))
        player22.add_child(Node.u8("stage4_sheet_num", 1))
        player22.add_child(Node.u8("norma_type", 2))
        player22.add_child(Node.s32("norma_1_num", 5))
        player22.add_child(Node.s32("norma_2_num", 0))
        player22.add_child(Node.u8("clear_medal", course["medal"]))
        player22.add_child(Node.u8("clear_norma", 2))
        player22.add_child(Node.s32("total_score", course["total"]))
        player22.add_child(Node.s16("max_combo", course["combo"]))

        for stage, music in enumerate([148, 550, 1113, 341]):
            stagenode = Node.void("stage")
            player22.add_child(stagenode)
            stagenode.add_child(Node.u8("stage", stage))
            stagenode.add_child(Node.s16("music_num", music))
            stagenode.add_child(Node.u8("sheet_num", 1))
            stagenode.add_child(Node.s32("score", course[f"stage{stage + 1}"]))

        # Swap with server
        resp = self.exchange("", call)
        self.assert_path(resp, "response/player22/@status")

    def verify_player22_new(self, ref_id: str) -> None:
        call = self.call_node()

        # Construct node
        player22 = Node.void("player22")
        call.add_child(player22)
        player22.set_attribute("method", "new")

        player22.add_child(Node.string("ref_id", ref_id))
        player22.add_child(Node.string("name", self.NAME))
        player22.add_child(Node.string("shop_name", ""))
        player22.add_child(Node.s8("pref", 51))

        # Swap with server
        resp = self.exchange("", call)

        # Verify nodes that cause crashes if they don't exist
        self.assert_path(resp, "response/player22/account")

    def verify(self, cardid: Optional[str]) -> None:
        # Verify boot sequence is okay
        self.verify_services_get(
            expected_services=[
                "pcbtracker",
                "pcbevent",
                "local",
                "message",
                "facility",
                "cardmng",
                "package",
                "posevent",
                "pkglist",
                "dlstatus",
                "eacoin",
                "lobby",
                "ntp",
                "keepalive",
            ]
        )
        paseli_enabled = self.verify_pcbtracker_alive()
        self.verify_message_get()
        self.verify_package_list()
        self.verify_facility_get()
        self.verify_pcbevent_put()
        self.verify_pcb22_boot()
        self.verify_info22_common()

        # Verify card registration and profile lookup
        if cardid is not None:
            card = cardid
        else:
            card = self.random_card()
            print(f"Generated random card ID {card} for use.")

        if cardid is None:
            self.verify_cardmng_inquire(
                card, msg_type="unregistered", paseli_enabled=paseli_enabled
            )
            ref_id = self.verify_cardmng_getrefid(card)
            if len(ref_id) != 16:
                raise Exception(
                    f"Invalid refid '{ref_id}' returned when registering card"
                )
            if ref_id != self.verify_cardmng_inquire(
                card, msg_type="new", paseli_enabled=paseli_enabled
            ):
                raise Exception(f"Invalid refid '{ref_id}' returned when querying card")
            self.verify_player22_read(ref_id, msg_type="new")
            self.verify_player22_new(ref_id)
        else:
            print("Skipping new card checks for existing card")
            ref_id = self.verify_cardmng_inquire(
                card, msg_type="query", paseli_enabled=paseli_enabled
            )

        # Verify pin handling and return card handling
        self.verify_cardmng_authpass(ref_id, correct=True)
        self.verify_cardmng_authpass(ref_id, correct=False)
        if ref_id != self.verify_cardmng_inquire(
            card, msg_type="query", paseli_enabled=paseli_enabled
        ):
            raise Exception(f"Invalid refid '{ref_id}' returned when querying card")

        if cardid is None:
            # Verify score handling
            scores = self.verify_player22_read(ref_id, msg_type="query")
            if scores is None:
                raise Exception("Expected to get scores back, didn't get anything!")
            for medal in scores["medals"]:
                for i in range(4):
                    if medal[i] != 0:
                        raise Exception("Got nonzero medals count on a new card!")
            for score in scores["scores"]:
                for i in range(4):
                    if score[i] != 0:
                        raise Exception("Got nonzero scores count on a new card!")
            for _ in scores["courses"]:
                raise Exception("Got nonzero courses count on a new card!")

            for phase in [1, 2]:
                if phase == 1:
                    dummyscores = [
                        # An okay score on a chart
                        {
                            "id": 987,
                            "chart": 2,
                            "medal": 5,
                            "score": 76543,
                        },
                        # A good score on an easier chart of the same song
                        {
                            "id": 987,
                            "chart": 0,
                            "medal": 6,
                            "score": 99999,
                        },
                        # A bad score on a hard chart
                        {
                            "id": 741,
                            "chart": 3,
                            "medal": 2,
                            "score": 45000,
                        },
                        # A terrible score on an easy chart
                        {
                            "id": 742,
                            "chart": 1,
                            "medal": 2,
                            "score": 1,
                        },
                    ]
                    # Random score to add in
                    songid = random.randint(907, 950)
                    chartid = random.randint(0, 3)
                    score = random.randint(0, 100000)
                    medal = random.randint(1, 11)
                    dummyscores.append(
                        {
                            "id": songid,
                            "chart": chartid,
                            "medal": medal,
                            "score": score,
                        }
                    )
                if phase == 2:
                    dummyscores = [
                        # A better score on the same chart
                        {
                            "id": 987,
                            "chart": 2,
                            "medal": 5,
                            "score": 98765,
                        },
                        # A worse score on another same chart
                        {
                            "id": 987,
                            "chart": 0,
                            "medal": 3,
                            "score": 12345,
                            "expected_score": 99999,
                            "expected_medal": 6,
                        },
                    ]

                for dummyscore in dummyscores:
                    self.verify_player22_write_music(ref_id, dummyscore)
                self.verify_player22_write(ref_id, dummyscores)
                scores = self.verify_player22_read(ref_id, msg_type="query")
                for score in dummyscores:
                    newscore = scores["scores"][score["id"]][score["chart"]]
                    newmedal = scores["medals"][score["id"]][score["chart"]]

                    if "expected_score" in score:
                        expected_score = score["expected_score"]
                    else:
                        expected_score = score["score"]
                    if "expected_medal" in score:
                        expected_medal = score["expected_medal"]
                    else:
                        expected_medal = score["medal"]

                    if newscore != expected_score:
                        raise Exception(
                            f'Expected a score of \'{expected_score}\' for song \'{score["id"]}\' chart \'{score["chart"]}\' but got score \'{newscore}\''
                        )
                    if newmedal != expected_medal:
                        raise Exception(
                            f'Expected a medal of \'{expected_medal}\' for song \'{score["id"]}\' chart \'{score["chart"]}\' but got medal \'{newmedal}\''
                        )

                # Sleep so we don't end up putting in score history on the same second
                time.sleep(1)

            # Write a random course so we know we can retrieve them.
            course = {
                "id": random.randint(1, 100),
                "medal": 2,
                "combo": random.randint(10, 100),
                "stage1": random.randint(70000, 100000),
                "stage2": random.randint(70000, 100000),
                "stage3": random.randint(70000, 100000),
                "stage4": random.randint(70000, 100000),
            }
            course["total"] = sum(course[f"stage{i + 1}"] for i in range(4))
            self.verify_player22_write_course(ref_id, course)

            # Now, grab the profile one more time and see that it is there.
            scores = self.verify_player22_read(ref_id, msg_type="query")
            if len(scores["courses"]) != 1:
                raise Exception("Did not get a course back after saving!")
            if course["id"] not in scores["courses"]:
                raise Exception("Did not get expected course back after saving!")
            for key in [
                "medal",
                "combo",
                "stage1",
                "stage2",
                "stage3",
                "stage4",
                "total",
            ]:
                if course[key] != scores["courses"][course["id"]][key]:
                    raise Exception(
                        f'Expected a {key} of \'{course[key]}\' but got \'{scores["courses"][course["id"]][key]}\''
                    )
        else:
            print("Skipping score checks for existing card")

        # Verify paseli handling
        if paseli_enabled:
            print("PASELI enabled for this PCBID, executing PASELI checks")
        else:
            print("PASELI disabled for this PCBID, skipping PASELI checks")
            return

        sessid, balance = self.verify_eacoin_checkin(card)
        if balance == 0:
            print("Skipping PASELI consume check because card has 0 balance")
        else:
            self.verify_eacoin_consume(sessid, balance, random.randint(0, balance))
        self.verify_eacoin_checkout(sessid)