From c17ab2114074b86a3fc9505c17391016a99007c9 Mon Sep 17 00:00:00 2001 From: Robert Martin Date: Thu, 30 Apr 2020 20:44:07 +0200 Subject: [PATCH] merged run_amiibo_cli with run_controller_cli; some cleanups --- joycontrol/command_line_interface.py | 3 +- joycontrol/controller.py | 11 +++ joycontrol/controller_state.py | 7 +- joycontrol/protocol.py | 1 - run_amiibo_cli.py | 104 ------------------------ run_controller_cli.py | 115 +++++++++++++++++---------- 6 files changed, 87 insertions(+), 154 deletions(-) delete mode 100644 run_amiibo_cli.py diff --git a/joycontrol/command_line_interface.py b/joycontrol/command_line_interface.py index 8c4cad6..526726d 100644 --- a/joycontrol/command_line_interface.py +++ b/joycontrol/command_line_interface.py @@ -1,5 +1,6 @@ import inspect import logging +import shlex from aioconsole import ainput @@ -122,7 +123,7 @@ class ControllerCLI: buttons_to_push = [] for command in user_input.split('&&'): - cmd, *args = command.split() + cmd, *args = shlex.split(command) if cmd == 'exit': return diff --git a/joycontrol/controller.py b/joycontrol/controller.py index 33da3ae..3faecb4 100644 --- a/joycontrol/controller.py +++ b/joycontrol/controller.py @@ -18,3 +18,14 @@ class Controller(enum.Enum): return 'Pro Controller' else: raise NotImplementedError() + + @staticmethod + def from_arg(arg): + if arg == 'JOYCON_R': + return Controller.JOYCON_R + elif arg == 'JOYCON_L': + return Controller.JOYCON_L + elif arg == 'PRO_CONTROLLER': + return Controller.PRO_CONTROLLER + else: + raise ValueError(f'Unknown controller "{arg}".') diff --git a/joycontrol/controller_state.py b/joycontrol/controller_state.py index aee833e..997ed6b 100644 --- a/joycontrol/controller_state.py +++ b/joycontrol/controller_state.py @@ -48,6 +48,9 @@ class ControllerState: def get_flash_memory(self): return self._spi_flash + def set_nfc(self, nfc_content): + self._nfc_content = nfc_content + async def send(self): """ Invokes protocol.send_controller_state(). Returns after the controller state was send. @@ -199,10 +202,6 @@ async def button_push(controller_state, *buttons, sec=0.1): await controller_state.send() -async def set_nfc(controller_state, nfc_content): - controller_state._nfc_content = nfc_content - - class _StickCalibration: def __init__(self, h_center, v_center, h_max_above_center, v_max_above_center, h_max_below_center, v_max_below_center): self.h_center = h_center diff --git a/joycontrol/protocol.py b/joycontrol/protocol.py index 97b2537..cecf0b3 100644 --- a/joycontrol/protocol.py +++ b/joycontrol/protocol.py @@ -260,7 +260,6 @@ class ControllerProtocol(BaseProtocol): else: logging.info(f'Unknown MCU sub command {sub_command}') - async def _reply_to_sub_command(self, report): # classify sub command try: diff --git a/run_amiibo_cli.py b/run_amiibo_cli.py deleted file mode 100644 index 57a00e0..0000000 --- a/run_amiibo_cli.py +++ /dev/null @@ -1,104 +0,0 @@ -import argparse -import asyncio -import logging -import os -from contextlib import contextmanager - -from joycontrol import logging_default as log -from joycontrol.command_line_interface import ControllerCLI -from joycontrol.controller_state import ControllerState, button_push, set_nfc -from joycontrol.protocol import controller_protocol_factory, Controller -from joycontrol.server import create_hid_server - -logger = logging.getLogger(__name__) - - -async def _main(controller, reconnect_bt_addr=None, capture_file=None, spi_flash=None, device_id=None, amiibo=None): - factory = controller_protocol_factory(controller, spi_flash=spi_flash) - ctl_psm, itr_psm = 17, 19 - transport, protocol = await create_hid_server(factory, - reconnect_bt_addr=reconnect_bt_addr, - ctl_psm=ctl_psm, itr_psm=itr_psm, capture_file=capture_file, device_id=device_id) - - controller_state = protocol.get_controller_state() - if amiibo: - await set_nfc(controller_state, amiibo.read()) - - await controller_state.connect() - - async def amiibo(filename): - with open(filename, "rb") as amiibo_file: - content = amiibo_file.read() - await set_nfc(controller_state, content) - - async def remove_amiibo(): - await controller_state.set_nfc(None) - - cli = ControllerCLI(controller_state) - cli.add_command('amiibo', amiibo) - cli.add_command('remove_amiibo', remove_amiibo) - await cli.run() - - logger.info('Stopping communication...') - await transport.close() - - -if __name__ == '__main__': - # check if root - if not os.geteuid() == 0: - raise PermissionError('Script must be run as root!') - - # setup logging - log.configure() - - parser = argparse.ArgumentParser() - #parser.add_argument('controller', help='JOYCON_R, JOYCON_L or PRO_CONTROLLER') - parser.add_argument('-l', '--log') - parser.add_argument('-d', '--device_id') - parser.add_argument('--spi_flash') - parser.add_argument('-r', '--reconnect_bt_addr', type=str, default=None, - help='The Switch console bluetooth address, for reconnecting as an already paired controller') - parser.add_argument('-a', '--amiibo', type=argparse.FileType('rb'), default=None, - help='The amiibo dump file') - args = parser.parse_args() - - """ - if args.controller == 'JOYCON_R': - controller = Controller.JOYCON_R - elif args.controller == 'JOYCON_L': - controller = Controller.JOYCON_L - elif args.controller == 'PRO_CONTROLLER': - controller = Controller.PRO_CONTROLLER - else: - raise ValueError(f'Unknown controller "{args.controller}".') - """ - controller = Controller.PRO_CONTROLLER - - spi_flash = None - if args.spi_flash: - with open(args.spi_flash, 'rb') as spi_flash_file: - spi_flash = spi_flash_file.read() - - # creates file if arg is given - @contextmanager - def get_output(path=None): - """ - Opens file if path is given - """ - if path is not None: - file = open(path, 'wb') - yield file - file.close() - else: - yield None - - with get_output(args.log) as capture_file: - loop = asyncio.get_event_loop() - loop.run_until_complete(_main( - controller, - reconnect_bt_addr=args.reconnect_bt_addr, - capture_file=capture_file, - spi_flash=spi_flash, - device_id=args.device_id, - amiibo=args.amiibo - )) diff --git a/run_controller_cli.py b/run_controller_cli.py index 6f517b2..9204b6e 100644 --- a/run_controller_cli.py +++ b/run_controller_cli.py @@ -4,7 +4,6 @@ import argparse import asyncio import logging import os -from contextlib import contextmanager from aioconsole import ainput @@ -133,31 +132,79 @@ async def test_controller_buttons(controller_state: ControllerState): await button_push(controller_state, 'home') -async def _main(controller, reconnect_bt_addr=None, capture_file=None, spi_flash=None, device_id=None): - factory = controller_protocol_factory(controller, spi_flash=spi_flash) - ctl_psm, itr_psm = 17, 19 - transport, protocol = await create_hid_server(factory, reconnect_bt_addr=reconnect_bt_addr, ctl_psm=ctl_psm, - itr_psm=itr_psm, capture_file=capture_file, device_id=device_id) +async def set_amiibo(controller_state, file_path): + """ + Sets nfc content of the controller state to contents of the given file. + :param controller_state: Emulated controller state + :param file_path: Path to amiibo dump file + """ + loop = asyncio.get_event_loop() - controller_state = protocol.get_controller_state() + with open(file_path, 'rb') as amiibo_file: + content = await loop.run_in_executor(None, amiibo_file.read) + controller_state.set_nfc(content) - # Create command line interface and add some extra commands - cli = ControllerCLI(controller_state) - # Wrap the script so we can pass the controller state. The doc string will be printed when calling 'help' - async def _run_test_controller_buttons(): - """ - test_buttons - Navigates to the "Test Controller Buttons" menu and presses all buttons. - """ - await test_controller_buttons(controller_state) +async def _main(args): + # parse the spi flash + spi_flash = None + if args.spi_flash: + with open(args.spi_flash, 'rb') as spi_flash_file: + spi_flash = FlashMemory(spi_flash_file.read()) - # add the script from above - cli.add_command('test_buttons', _run_test_controller_buttons) + # Get controller name to emulate from arguments + controller = Controller.from_arg(args.controller) - await cli.run() + with utils.get_output(path=args.log, default=None) as capture_file: + factory = controller_protocol_factory(controller, spi_flash=spi_flash) + ctl_psm, itr_psm = 17, 19 + transport, protocol = await create_hid_server(factory, reconnect_bt_addr=args.reconnect_bt_addr, + ctl_psm=ctl_psm, + itr_psm=itr_psm, capture_file=capture_file, + device_id=args.device_id) - logger.info('Stopping communication...') - await transport.close() + controller_state = protocol.get_controller_state() + + # Create command line interface and add some extra commands + cli = ControllerCLI(controller_state) + + # Wrap the script so we can pass the controller state. The doc string will be printed when calling 'help' + async def _run_test_controller_buttons(): + """ + test_buttons - Navigates to the "Test Controller Buttons" menu and presses all buttons. + """ + await test_controller_buttons(controller_state) + + # add the script from above + cli.add_command('test_buttons', _run_test_controller_buttons) + + # Create amiibo command + async def amiibo(*args): + """ + amiibo - Sets amiibo content + + Usage: + amiibo Set controller state NFC content to file + amiibo remove Remove NFC content from controller state + """ + if controller_state.get_controller() == Controller.JOYCON_L: + raise ValueError('NFC content cannot be set for JOYCON_L') + elif not args: + raise ValueError('"amiibo" command requires amiibo dump file path as argument!') + elif args[0] == 'remove': + controller_state.set_nfc(None) + print('Removed nfc content.') + else: + await set_amiibo(controller_state, args[0]) + + # add the script from above + cli.add_command('amiibo', amiibo) + + try: + await cli.run() + finally: + logger.info('Stopping communication...') + await transport.close() if __name__ == '__main__': @@ -178,27 +225,7 @@ if __name__ == '__main__': help='The Switch console Bluetooth address, for reconnecting as an already paired controller') args = parser.parse_args() - if args.controller == 'JOYCON_R': - controller = Controller.JOYCON_R - elif args.controller == 'JOYCON_L': - controller = Controller.JOYCON_L - elif args.controller == 'PRO_CONTROLLER': - controller = Controller.PRO_CONTROLLER - else: - raise ValueError(f'Unknown controller "{args.controller}".') - - spi_flash = None - if args.spi_flash: - with open(args.spi_flash, 'rb') as spi_flash_file: - spi_flash = FlashMemory(spi_flash_file.read()) - - with utils.get_output(path=args.log, default=None) as capture_file: - loop = asyncio.get_event_loop() - loop.run_until_complete( - _main(controller, - reconnect_bt_addr=args.reconnect_bt_addr, - capture_file=capture_file, - spi_flash=spi_flash, - device_id=args.device_id - ) - ) + loop = asyncio.get_event_loop() + loop.run_until_complete( + _main(args) + )