merged run_amiibo_cli with run_controller_cli; some cleanups

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

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}".')

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

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:

run_amiibo_cli.py

@@ -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
-        ))

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,11 +132,36 @@ 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):
+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()
+
+    with open(file_path, 'rb') as amiibo_file:
+        content = await loop.run_in_executor(None, amiibo_file.read)
+        controller_state.set_nfc(content)
+
+
+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())
+
+    # Get controller name to emulate from arguments
+    controller = Controller.from_arg(args.controller)
+
+    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=reconnect_bt_addr, ctl_psm=ctl_psm,
+        transport, protocol = await create_hid_server(factory, reconnect_bt_addr=args.reconnect_bt_addr,
-                                                  itr_psm=itr_psm, capture_file=capture_file, device_id=device_id)
+                                                      ctl_psm=ctl_psm,
+                                                      itr_psm=itr_psm, capture_file=capture_file,
+                                                      device_id=args.device_id)
 
         controller_state = protocol.get_controller_state()
 
@@ -154,8 +178,31 @@ async def _main(controller, reconnect_bt_addr=None, capture_file=None, spi_flash
         # add the script from above
         cli.add_command('test_buttons', _run_test_controller_buttons)
 
-    await cli.run()
+        # Create amiibo command
+        async def amiibo(*args):
+            """
+            amiibo - Sets amiibo content
 
+            Usage:
+                amiibo <file_name>          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()
 
@@ -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,
+        _main(args)
-                  reconnect_bt_addr=args.reconnect_bt_addr,
-                  capture_file=capture_file,
-                  spi_flash=spi_flash,
-                  device_id=args.device_id
-                  )
     )