bemaniutils/bemani/utils/trafficgen.py

import argparse
import sys
from typing import Any, Dict, Optional
import yaml

from bemani.client import ClientProtocol, BaseClient
from bemani.client.iidx import (
    IIDXTricoroClient,
    IIDXSpadaClient,
    IIDXPendualClient,
    IIDXCopulaClient,
    IIDXSinobuzClient,
    IIDXCannonBallersClient,
    IIDXRootageClient,
)
from bemani.client.jubeat import (
    JubeatSaucerClient,
    JubeatSaucerFulfillClient,
    JubeatPropClient,
    JubeatQubellClient,
    JubeatClanClient,
)
from bemani.client.popn import (
    PopnMusicTuneStreetClient,
    PopnMusicFantasiaClient,
    PopnMusicSunnyParkClient,
    PopnMusicLapistoriaClient,
    PopnMusicEclaleClient,
    PopnMusicUsaNekoClient,
    PopnMusicPeaceClient,
)
from bemani.client.ddr import (
    DDRX2Client,
    DDRX3Client,
    DDR2013Client,
    DDR2014Client,
    DDRAceClient,
)
from bemani.client.sdvx import (
    SoundVoltexBoothClient,
    SoundVoltexInfiniteInfectionClient,
    SoundVoltexGravityWarsS1Client,
    SoundVoltexGravityWarsS2Client,
    SoundVoltexHeavenlyHavenClient,
)
from bemani.client.museca import (
    Museca1Client,
    Museca1PlusClient,
)
from bemani.client.reflec import (
    ReflecBeat,
    ReflecBeatLimelight,
    ReflecBeatColette,
    ReflecBeatGroovinUpper,
    ReflecBeatVolzza,
    ReflecBeatVolzza2,
)
from bemani.client.bishi import TheStarBishiBashiClient
from bemani.client.mga.mga import MetalGearArcadeClient


def get_client(proto: ClientProtocol, pcbid: str, game: str, config: Dict[str, Any]) -> BaseClient:
    if game == 'pnm-tune-street':
        return PopnMusicTuneStreetClient(
            proto,
            pcbid,
            config,
        )
    if game == 'pnm-fantasia':
        return PopnMusicFantasiaClient(
            proto,
            pcbid,
            config,
        )
    if game == 'pnm-sunny-park':
        return PopnMusicSunnyParkClient(
            proto,
            pcbid,
            config,
        )
    if game == 'pnm-lapistoria':
        return PopnMusicLapistoriaClient(
            proto,
            pcbid,
            config,
        )
    if game == 'pnm-eclale':
        return PopnMusicEclaleClient(
            proto,
            pcbid,
            config,
        )
    if game == 'pnm-usaneko':
        return PopnMusicUsaNekoClient(
            proto,
            pcbid,
            config,
        )
    if game == 'pnm-peace':
        return PopnMusicPeaceClient(
            proto,
            pcbid,
            config,
        )
    if game == 'jubeat-saucer':
        return JubeatSaucerClient(
            proto,
            pcbid,
            config,
        )
    if game == 'jubeat-saucer-fulfill':
        return JubeatSaucerFulfillClient(
            proto,
            pcbid,
            config,
        )
    if game == 'jubeat-prop':
        return JubeatPropClient(
            proto,
            pcbid,
            config,
        )
    if game == 'jubeat-qubell':
        return JubeatQubellClient(
            proto,
            pcbid,
            config,
        )
    if game == 'jubeat-clan':
        return JubeatClanClient(
            proto,
            pcbid,
            config,
        )
    if game == 'iidx-rootage':
        return IIDXRootageClient(
            proto,
            pcbid,
            config,
        )
    if game == 'iidx-cannonballers':
        return IIDXCannonBallersClient(
            proto,
            pcbid,
            config,
        )
    if game == 'iidx-sinobuz':
        return IIDXSinobuzClient(
            proto,
            pcbid,
            config,
        )
    if game == 'iidx-copula':
        return IIDXCopulaClient(
            proto,
            pcbid,
            config,
        )
    if game == 'iidx-pendual':
        return IIDXPendualClient(
            proto,
            pcbid,
            config,
        )
    if game == 'iidx-spada':
        return IIDXSpadaClient(
            proto,
            pcbid,
            config,
        )
    if game == 'iidx-tricoro':
        return IIDXTricoroClient(
            proto,
            pcbid,
            config,
        )
    if game == 'bishi':
        return TheStarBishiBashiClient(
            proto,
            pcbid,
            config,
        )
    if game == 'ddr-x2':
        return DDRX2Client(
            proto,
            pcbid,
            config,
        )
    if game == 'ddr-x3':
        return DDRX3Client(
            proto,
            pcbid,
            config,
        )
    if game == 'ddr-2013':
        return DDR2013Client(
            proto,
            pcbid,
            config,
        )
    if game == 'ddr-2014':
        return DDR2014Client(
            proto,
            pcbid,
            config,
        )
    if game == 'ddr-ace':
        return DDRAceClient(
            proto,
            pcbid,
            config,
        )
    if game == 'sdvx-booth':
        return SoundVoltexBoothClient(
            proto,
            pcbid,
            config,
        )
    if game == 'sdvx-infinite-infection':
        return SoundVoltexInfiniteInfectionClient(
            proto,
            pcbid,
            config,
        )
    if game == 'sdvx-gravity-wars-s1':
        return SoundVoltexGravityWarsS1Client(
            proto,
            pcbid,
            config,
        )
    if game == 'sdvx-gravity-wars-s2':
        return SoundVoltexGravityWarsS2Client(
            proto,
            pcbid,
            config,
        )
    if game == 'sdvx-heavenly-haven':
        return SoundVoltexHeavenlyHavenClient(
            proto,
            pcbid,
            config,
        )
    if game == 'museca-1':
        return Museca1Client(
            proto,
            pcbid,
            config,
        )
    if game == 'museca-1+1/2':
        return Museca1PlusClient(
            proto,
            pcbid,
            config,
        )
    if game == 'reflec':
        return ReflecBeat(
            proto,
            pcbid,
            config,
        )
    if game == 'reflec-limelight':
        return ReflecBeatLimelight(
            proto,
            pcbid,
            config,
        )
    if game == 'reflec-colette':
        return ReflecBeatColette(
            proto,
            pcbid,
            config,
        )
    if game == 'reflec-groovin-upper':
        return ReflecBeatGroovinUpper(
            proto,
            pcbid,
            config,
        )
    if game == 'reflec-volzza':
        return ReflecBeatVolzza(
            proto,
            pcbid,
            config,
        )
    if game == 'reflec-volzza2':
        return ReflecBeatVolzza2(
            proto,
            pcbid,
            config,
        )
    if game == 'metal-gear-arcade':
        return MetalGearArcadeClient(
            proto,
            pcbid,
            config,
        )

    raise Exception(f'Unknown game {game}')


def mainloop(address: str, port: int, configfile: str, action: str, game: str, cardid: Optional[str], verbose: bool) -> None:
    games = {
        'pnm-tune-street': {
            'name': "Pop'n Music Tune Street",
            'model': "K39:J:B:A:2010122200",
            'old_profile_model': "J39:J:A:A",
            'avs': None,
        },
        'pnm-fantasia': {
            'name': "Pop'n Music Fantasia",
            'model': "L39:J:B:A:2012091900",
            'old_profile_model': "K39:J:A:A",
            'avs': "2.13.6 r4921",
        },
        'pnm-sunny-park': {
            'name': "Pop'n Music Sunny Park",
            'model': "M39:J:B:A:2014061900",
            'old_profile_model': "L39:J:A:A",
            'avs': "2.15.8 r6631",
        },
        'pnm-lapistoria': {
            'name': "Pop'n Music Lapistoria",
            'model': "M39:J:B:A:2015081900",
            'old_profile_model': "M39:J:B:A",
            'avs': "2.15.8 r6631",
        },
        'pnm-eclale': {
            'name': "Pop'n Music Eclale",
            'model': "M39:J:B:A:2016100500",
            'old_profile_model': "M39:J:B:A",
            'avs': "2.15.8 r6631",
        },
        'pnm-usaneko': {
            'name': "Pop'n Music Usagi to Neko to Shounen no Yume",
            'model': "M39:J:B:A:2018101500",
            'old_profile_model': "M39:J:B:A",
            'avs': "2.15.8 r6631",
        },
        'pnm-peace': {
            'name': "Pop'n Music peace",
            'model': "M39:J:B:A:2020092800",
            'old_profile_model': "M39:J:B:A",
            'avs': "2.15.8 r6631",
        },
        'jubeat-saucer': {
            'name': "Jubeat Saucer",
            'model': "L44:J:A:A:2014012802",
            'avs': "2.15.8 r6631",
        },
        'jubeat-saucer-fulfill': {
            'name': "Jubeat Saucer Fulfill",
            'model': "L44:J:B:A:2014111800",
            'avs': "2.15.8 r6631",
        },
        'jubeat-prop': {
            'name': "Jubeat Prop",
            'model': "L44:J:B:A:2016031700",
            'avs': "2.15.8 r6631",
        },
        'jubeat-qubell': {
            'name': "Jubeat Qubell",
            'model': "L44:J:D:A:2016111400",
            'avs': "2.15.8 r6631",
        },
        'jubeat-clan': {
            'name': "Jubeat Clan",
            'model': "L44:J:E:A:2018070901",
            'avs': "2.17.3 r8311",
        },
        'iidx-rootage': {
            'name': "Beatmania IIDX ROOTAGE",
            'model': "LDJ:J:A:A:2019090200",
            'avs': "2.17.0 r7883"
        },
        'iidx-cannonballers': {
            'name': "Beatmania IIDX CANNON BALLERS",
            'model': "LDJ:J:A:A:2018091900",
            'avs': "2.17.0 r7883"
        },
        'iidx-sinobuz': {
            'name': "Beatmania IIDX SINOBUZ",
            'model': "LDJ:J:A:A:2017082800",
            'avs': "2.16.1 r6901",
        },
        'iidx-copula': {
            'name': "Beatmania IIDX copula",
            'model': "LDJ:J:A:A:2016083100",
            'avs': "2.16.1 r6901",
        },
        'iidx-pendual': {
            'name': "Beatmania IIDX PENDUAL",
            'model': "LDJ:A:A:A:2015080500",
            'avs': "2.16.1 r6901",
        },
        'iidx-spada': {
            'name': "Beatmania IIDX SPADA",
            'model': "LDJ:A:A:A:2014071600",
            'avs': "2.16.1 r6901",
        },
        'iidx-tricoro': {
            'name': "Beatmania IIDX Tricoro",
            'model': "LDJ:J:A:A:2013090900",
            'avs': "2.15.8 r6631",
        },
        'bishi': {
            'name': "The★BishiBashi",
            'model': "IBB:A:A:A:2009092900",
            'avs': None,
        },
        'ddr-x2': {
            'name': "DanceDanceRevolution X2",
            'model': "JDX:J:A:A:2010111000",
            'avs': None,
        },
        'ddr-x3': {
            'name': "DanceDanceRevolution X3 VS 2ndMIX",
            'model': "KDX:J:A:A:2012112600",
            'avs': "2.13.6 r4921",
        },
        'ddr-2013': {
            'name': "DanceDanceRevolution (2013)",
            'model': "MDX:J:A:A:2014032700",
            'avs': "2.15.8 r6631",
        },
        'ddr-2014': {
            'name': "DanceDanceRevolution (2014)",
            'model': "MDX:A:A:A:2015122100",
            'avs': "2.15.8 r6631",
        },
        'ddr-ace': {
            'name': "DanceDanceRevolution A",
            'model': "MDX:U:D:A:2017121400",
            'avs': "2.15.8 r6631",
        },
        'sdvx-booth': {
            'name': "SOUND VOLTEX BOOTH",
            'model': "KFC:J:A:A:2013052900",
            'avs': "2.15.8 r6631",
        },
        'sdvx-infinite-infection': {
            'name': "SOUND VOLTEX II -infinite infection-",
            'model': "KFC:J:A:A:2014102200",
            'avs': "2.15.8 r6631",
        },
        'sdvx-gravity-wars-s1': {
            'name': "SOUND VOLTEX III GRAVITY WARS Season 1",
            'model': "KFC:J:A:A:2015111602",
            'avs': "2.15.8 r6631",
        },
        'sdvx-gravity-wars-s2': {
            'name': "SOUND VOLTEX III GRAVITY WARS Season 2",
            'model': "KFC:J:A:A:2016121900",
            'avs': "2.15.8 r6631",
        },
        'sdvx-heavenly-haven': {
            'name': "SOUND VOLTEX IV HEAVENLY HAVEN",
            'model': "KFC:J:A:A:2019020600",
            'avs': "2.15.8 r6631",
        },
        'museca-1': {
            'name': "MÚSECA",
            'model': "PIX:J:A:A:2016071300",
            'avs': "2.17.0 r7883",
        },
        'museca-1+1/2': {
            'name': "MÚSECA 1+1/2",
            'model': "PIX:J:A:A:2017042600",
            'avs': "2.17.0 r7883",
        },
        'reflec': {
            'name': "REFLEC BEAT",
            'model': "KBR:A:A:A:2011112300",
            'avs': None,
        },
        'reflec-limelight': {
            'name': "REFLEC BEAT limelight",
            'model': "LBR:A:A:A:2012082900",
            'avs': "2.13.6 r4921",
        },
        'reflec-colette': {
            'name': "REFLEC BEAT colette",
            'model': "MBR:J:A:A:2014011600",
            'avs': "2.15.8 r6631",
        },
        'reflec-groovin-upper': {
            'name': "REFLEC BEAT groovin'!! Upper",
            'model': "MBR:J:A:A:2015102100",
            'avs': "2.15.8 r6631",
        },
        'reflec-volzza': {
            'name': "REFLEC BEAT VOLZZA",
            'model': "MBR:J:A:A:2016030200",
            'avs': "2.15.8 r6631",
        },
        'reflec-volzza2': {
            'name': "REFLEC BEAT VOLZZA 2",
            'model': "MBR:J:A:A:2016100400",
            'avs': "2.15.8 r6631",
        },
        'metal-gear-arcade': {
            'name': "Metal Gear Arcade",
            'model': "I36:J:A:A:2011092900",
            'avs': None,
        },
    }
    if action == 'list':
        for game in sorted([game for game in games]):
            print(f'{game} - {games[game]["name"]}')
        sys.exit(0)
    if action == 'game':
        if game not in games:
            print(f'Unknown game {game}')
            sys.exit(2)

        config = yaml.safe_load(open(configfile))

        print(f'Emulating {games[game]["name"]}')
        emu = get_client(
            ClientProtocol(
                address,
                port,
                config['core']['encryption'],
                config['core']['compression'],
                verbose,
            ),
            config['core']['pcbid'],
            game,
            games[game],
        )

        emu.verify(cardid)


def main() -> None:
    parser = argparse.ArgumentParser(description="A utility to generate game-like traffic for testing an eAmusement server.")
    parser.add_argument("-p", "--port", help="Port to talk to. Defaults to 80", type=int, default=80)
    parser.add_argument("-a", "--address", help="Address to talk to. Defaults to 127.0.0.1", type=str, default="127.0.0.1")
    parser.add_argument("-c", "--config", help="Core configuration. Defaults to trafficgen.yaml", type=str, default="trafficgen.yaml")
    parser.add_argument("-g", "--game", help="The game that should be emulated. Should be one of the games returned by --list", type=str, default=None)
    parser.add_argument("-l", "--list", help="List all known games and exit.", action="store_true")
    parser.add_argument("-i", "--cardid", help="Use this card ID instead of a random one.", type=str, default=None)
    parser.add_argument("-v", "--verbose", help="Print packets that are sent/received.", action='store_true', default=False)
    args = parser.parse_args()

    if args.list:
        action = 'list'
        game = None
    elif args.game:
        action = 'game'
        game = args.game
    else:
        print("Unknown action to perform. Please specify --game <game> or --list")
        sys.exit(1)

    game = {
        'pnm-19': 'pnm-tune-street',
        'pnm-20': 'pnm-fantasia',
        'pnm-21': 'pnm-sunny-park',
        'pnm-22': 'pnm-lapistoria',
        'pnm-23': 'pnm-eclale',
        'pnm-24': 'pnm-usaneko',
        'pnm-25': 'pnm-peace',
        'iidx-20': 'iidx-tricoro',
        'iidx-21': 'iidx-spada',
        'iidx-22': 'iidx-pendual',
        'iidx-23': 'iidx-copula',
        'iidx-24': 'iidx-sinobuz',
        'iidx-25': 'iidx-cannonballers',
        'iidx-26': 'iidx-rootage',
        'jubeat-5': 'jubeat-saucer',
        'jubeat-5+': 'jubeat-saucer-fulfill',
        'jubeat-6': 'jubeat-prop',
        'jubeat-7': 'jubeat-qubell',
        'jubeat-8': 'jubeat-clan',
        'ddr-12': 'ddr-x2',
        'ddr-13': 'ddr-x3',
        'ddr-14': 'ddr-2013',
        'ddr-15': 'ddr-2014',
        'ddr-16': 'ddr-ace',
        'sdvx-1': 'sdvx-booth',
        'sdvx-2': 'sdvx-infinite-infection',
        'sdvx-3s1': 'sdvx-gravity-wars-s1',
        'sdvx-3s2': 'sdvx-gravity-wars-s2',
        'sdvx-4': 'sdvx-heavenly-haven',
        'reflec-1': 'reflec',
        'reflec-2': 'reflec-limelight',
        'reflec-3': 'reflec-colette',
        'reflec-4': 'reflec-groovin-upper',
        'reflec-5': 'reflec-volzza',
        'reflec-6': 'reflec-volzza2',
        'mga': 'metal-gear-arcade',
    }.get(game, game)

    mainloop(args.address, args.port, args.config, action, game, args.cardid, args.verbose)


if __name__ == '__main__':
    main()