From b8d2853f3988fff6fb52c88b9b57cf607a847a38 Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Wed, 9 Aug 2023 20:00:54 +0200 Subject: [PATCH 1/9] Use consistent naming for command units --- software/script/chameleon_cli_main.py | 60 +++++++++++++-------------- software/script/chameleon_cli_unit.py | 46 ++++++++++---------- 2 files changed, 53 insertions(+), 53 deletions(-) diff --git a/software/script/chameleon_cli_main.py b/software/script/chameleon_cli_main.py index 83fba87..082ad1d 100755 --- a/software/script/chameleon_cli_main.py +++ b/software/script/chameleon_cli_main.py @@ -37,9 +37,9 @@ BANNER = f""" """ -def new_uint(unit_clz, help_msg): +def new_unit(unit_clz, help_msg): """ - new a uint dict object + Create a new unit dict object :param unit_clz: unit implement class :param help_msg: unit usage :return: a dict... @@ -58,65 +58,65 @@ class ChameleonCLI: def __init__(self): self.cmd_maps = { 'hw': { - 'connect': new_uint(chameleon_cli_unit.HWConnect, "Connect to chameleon by serial port"), + 'connect': new_unit(chameleon_cli_unit.HWConnect, "Connect to chameleon by serial port"), 'chipid': { - 'get': new_uint(chameleon_cli_unit.HWChipIdGet, "Get device chipset ID"), + 'get': new_unit(chameleon_cli_unit.HWChipIdGet, "Get device chipset ID"), 'help': "Device chipsed ID get" }, 'address': { - 'get': new_uint(chameleon_cli_unit.HWAddressGet, "Get device address (used with Bluetooth)"), + 'get': new_unit(chameleon_cli_unit.HWAddressGet, "Get device address (used with Bluetooth)"), 'help': "Device address get" }, 'mode': { - 'set': new_uint(chameleon_cli_unit.HWModeSet, "Change device mode to tag reader or tag emulator"), - 'get': new_uint(chameleon_cli_unit.HWModeGet, "Get current device mode"), + 'set': new_unit(chameleon_cli_unit.HWModeSet, "Change device mode to tag reader or tag emulator"), + 'get': new_unit(chameleon_cli_unit.HWModeGet, "Get current device mode"), 'help': "Device mode get/set" }, 'slot': { - 'change': new_uint(chameleon_cli_unit.HWSlotSet, "Set emulation tag slot activated."), - 'type': new_uint(chameleon_cli_unit.HWSlotTagType, "Set emulation tag type"), - 'init': new_uint(chameleon_cli_unit.HWSlotDataDefault, "Set emulation tag data to default"), - 'enable': new_uint(chameleon_cli_unit.HWSlotEnableSet, "Set emulation tag slot enable or disable"), + 'change': new_unit(chameleon_cli_unit.HWSlotSet, "Set emulation tag slot activated."), + 'type': new_unit(chameleon_cli_unit.HWSlotTagType, "Set emulation tag type"), + 'init': new_unit(chameleon_cli_unit.HWSlotDataDefault, "Set emulation tag data to default"), + 'enable': new_unit(chameleon_cli_unit.HWSlotEnableSet, "Set emulation tag slot enable or disable"), 'nick': { - 'set': new_uint(chameleon_cli_unit.HWSlotNickSet, "Set tag nick name for slot"), - 'get': new_uint(chameleon_cli_unit.HWSlotNickGet, "Get tag nick name for slot"), + 'set': new_unit(chameleon_cli_unit.HWSlotNickSet, "Set tag nick name for slot"), + 'get': new_unit(chameleon_cli_unit.HWSlotNickGet, "Get tag nick name for slot"), 'help': "Get/Set tag nick name for slot", }, - 'update': new_uint(chameleon_cli_unit.HWSlotUpdate, "Update config & data to device flash"), - 'openall': new_uint(chameleon_cli_unit.HWSlotOpenAll, "Open all slot and set to default data"), + 'update': new_unit(chameleon_cli_unit.HWSlotUpdate, "Update config & data to device flash"), + 'openall': new_unit(chameleon_cli_unit.HWSlotOpenAll, "Open all slot and set to default data"), 'help': "Emulation tag slot.", }, - 'dfu': new_uint(chameleon_cli_unit.HWDFU, "Restart application to bootloader mode(Not yet implement dfu)."), + 'dfu': new_unit(chameleon_cli_unit.HWDFU, "Restart application to bootloader mode(Not yet implement dfu)."), 'help': "hardware controller", }, 'hf': { '14a': { - 'scan': new_uint(chameleon_cli_unit.HF14AScan, "Scan 14a tag, and print basic information"), - 'info': new_uint(chameleon_cli_unit.HF14AInfo, "Scan 14a tag, and print detail information"), + 'scan': new_unit(chameleon_cli_unit.HF14AScan, "Scan 14a tag, and print basic information"), + 'info': new_unit(chameleon_cli_unit.HF14AInfo, "Scan 14a tag, and print detail information"), 'help': "ISO14443-a tag read/write/info...", }, 'mf': { - 'nested': new_uint(chameleon_cli_unit.HFMFNested, "Mifare Classic nested recover key"), - 'darkside': new_uint(chameleon_cli_unit.HFMFDarkside, "Mifare Classic darkside recover key"), - 'rdbl': new_uint(chameleon_cli_unit.HFMFRDBL, "MiFARE Classic read one block"), - 'wrbl': new_uint(chameleon_cli_unit.HFMFWRBL, "MiFARE Classic write one block"), + 'nested': new_unit(chameleon_cli_unit.HFMFNested, "Mifare Classic nested recover key"), + 'darkside': new_unit(chameleon_cli_unit.HFMFDarkside, "Mifare Classic darkside recover key"), + 'rdbl': new_unit(chameleon_cli_unit.HFMFRDBL, "MiFARE Classic read one block"), + 'wrbl': new_unit(chameleon_cli_unit.HFMFWRBL, "MiFARE Classic write one block"), 'detection': { - 'enable': new_uint(chameleon_cli_unit.HFMFDetectionEnable, "Detection enable"), - 'count': new_uint(chameleon_cli_unit.HFMFDetectionLogCount, "Detection log count"), - 'decrypt': new_uint(chameleon_cli_unit.HFMFDetectionDecrypt, "Download log and decrypt keys"), + 'enable': new_unit(chameleon_cli_unit.HFMFDetectionEnable, "Detection enable"), + 'count': new_unit(chameleon_cli_unit.HFMFDetectionLogCount, "Detection log count"), + 'decrypt': new_unit(chameleon_cli_unit.HFMFDetectionDecrypt, "Download log and decrypt keys"), 'help': "Mifare Classic detection log" }, - 'sim': new_uint(chameleon_cli_unit.HFMFSim, "Simulation a mifare classic card"), - 'eload': new_uint(chameleon_cli_unit.HFMFELoad, "Load data to emulator memory"), + 'sim': new_unit(chameleon_cli_unit.HFMFSim, "Simulation a mifare classic card"), + 'eload': new_unit(chameleon_cli_unit.HFMFELoad, "Load data to emulator memory"), 'help': "Mifare Classic mini/1/2/4, attack/read/write" }, 'help': "high frequency tag/reader", }, 'lf': { 'em': { - 'read': new_uint(chameleon_cli_unit.LFEMRead, "Scan em410x tag and print id"), - 'write': new_uint(chameleon_cli_unit.LFEMWriteT55xx, "Write em410x id to t55xx"), - 'sim': new_uint(chameleon_cli_unit.LFEMSim, "Simulation a em410x id card"), + 'read': new_unit(chameleon_cli_unit.LFEMRead, "Scan em410x tag and print id"), + 'write': new_unit(chameleon_cli_unit.LFEMWriteT55xx, "Write em410x id to t55xx"), + 'sim': new_unit(chameleon_cli_unit.LFEMSim, "Simulation a em410x id card"), 'help': "EM410x read/write/emulator", }, 'help': "low frequency tag/reader", diff --git a/software/script/chameleon_cli_unit.py b/software/script/chameleon_cli_unit.py index 4cb0ded..4de0681 100644 --- a/software/script/chameleon_cli_unit.py +++ b/software/script/chameleon_cli_unit.py @@ -162,7 +162,7 @@ class DeviceRequiredUnit(BaseCLIUnit): raise NotImplementedError("Please implement this") -class ReaderRequiredUint(DeviceRequiredUnit): +class ReaderRequiredUnit(DeviceRequiredUnit): """ Make sure of device enter to reader mode. """ @@ -171,7 +171,7 @@ class ReaderRequiredUint(DeviceRequiredUnit): raise NotImplementedError("Please implement this") def before_exec(self, args: argparse.Namespace): - if super(ReaderRequiredUint, self).before_exec(args): + if super(ReaderRequiredUnit, self).before_exec(args): ret = self.cmd_standard.is_reader_device_mode() if ret: return True @@ -254,7 +254,7 @@ class HWAddressGet(DeviceRequiredUnit): print(f' - Device address: ' + self.cmd_positive.get_device_address()) -class HF14AScan(ReaderRequiredUint): +class HF14AScan(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: pass @@ -275,7 +275,7 @@ class HF14AScan(ReaderRequiredUint): return self.scan() -class HF14AInfo(ReaderRequiredUint): +class HF14AInfo(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: pass @@ -305,7 +305,7 @@ class HF14AInfo(ReaderRequiredUint): self.info() -class HFMFNested(ReaderRequiredUint): +class HFMFNested(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: type_choices = ['A', 'B', 'a', 'b'] @@ -410,7 +410,7 @@ class HFMFNested(ReaderRequiredUint): return -class HFMFDarkside(ReaderRequiredUint): +class HFMFDarkside(ReaderRequiredUnit): def __init__(self): super().__init__() @@ -476,7 +476,7 @@ class HFMFDarkside(ReaderRequiredUint): return -class BaseMF1AuthOpera(ReaderRequiredUint): +class BaseMF1AuthOpera(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: type_choices = ['A', 'B', 'a', 'b'] @@ -727,7 +727,7 @@ class HFMFSim(DeviceRequiredUnit): print(" - Set anti-collision resources success") -class LFEMRead(ReaderRequiredUint): +class LFEMRead(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: return None @@ -738,7 +738,7 @@ class LFEMRead(ReaderRequiredUint): print(f" - EM410x ID(10H): {colorama.Fore.GREEN}{id_hex}{colorama.Style.RESET_ALL}") -class LFEMCardRequiredUint(DeviceRequiredUnit): +class LFEMCardRequiredUnit(DeviceRequiredUnit): @staticmethod def add_card_arg(parser: ArgumentParserNoExit): @@ -746,7 +746,7 @@ class LFEMCardRequiredUint(DeviceRequiredUnit): return parser def before_exec(self, args: argparse.Namespace): - if super(LFEMCardRequiredUint, self).before_exec(args): + if super(LFEMCardRequiredUnit, self).before_exec(args): if not re.match(r"^[a-fA-F0-9]{10}$", args.id): raise ArgsParserError("ID must include 10 HEX symbols") return True @@ -759,15 +759,15 @@ class LFEMCardRequiredUint(DeviceRequiredUnit): raise NotImplementedError("Please implement this") -class LFEMWriteT55xx(LFEMCardRequiredUint, ReaderRequiredUint): +class LFEMWriteT55xx(LFEMCardRequiredUnit, ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() return self.add_card_arg(parser) def before_exec(self, args: argparse.Namespace): - b1 = super(LFEMCardRequiredUint, self).before_exec(args) - b2 = super(ReaderRequiredUint, self).before_exec(args) + b1 = super(LFEMCardRequiredUnit, self).before_exec(args) + b2 = super(ReaderRequiredUnit, self).before_exec(args) return b1 and b2 # lf em write --id 4400999559 @@ -778,7 +778,7 @@ class LFEMWriteT55xx(LFEMCardRequiredUint, ReaderRequiredUint): print(f" - EM410x ID(10H): {id_hex} write done.") -class SlotIndexRequireUint(DeviceRequiredUnit): +class SlotIndexRequireUnit(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: raise NotImplementedError() @@ -793,7 +793,7 @@ class SlotIndexRequireUint(DeviceRequiredUnit): help="Slot index", metavar="number", choices=slot_choices) return parser -class SenseTypeRequireUint(DeviceRequiredUnit): +class SenseTypeRequireUnit(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: raise NotImplementedError() @@ -809,7 +809,7 @@ class SenseTypeRequireUint(DeviceRequiredUnit): return parser -class HWSlotSet(SlotIndexRequireUint): +class HWSlotSet(SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -822,7 +822,7 @@ class HWSlotSet(SlotIndexRequireUint): print(f" - Set slot {slot_index} activated success.") -class TagTypeRequiredUint(DeviceRequiredUnit): +class TagTypeRequiredUnit(DeviceRequiredUnit): @staticmethod def add_type_args(parser: ArgumentParserNoExit): @@ -843,7 +843,7 @@ class TagTypeRequiredUint(DeviceRequiredUnit): raise NotImplementedError() -class HWSlotTagType(TagTypeRequiredUint, SlotIndexRequireUint): +class HWSlotTagType(TagTypeRequiredUnit, SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -859,7 +859,7 @@ class HWSlotTagType(TagTypeRequiredUint, SlotIndexRequireUint): print(f' - Set slot tag type success.') -class HWSlotDataDefault(TagTypeRequiredUint, SlotIndexRequireUint): +class HWSlotDataDefault(TagTypeRequiredUnit, SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -876,7 +876,7 @@ class HWSlotDataDefault(TagTypeRequiredUint, SlotIndexRequireUint): print(f' - Set slot tag data init success.') -class HWSlotEnableSet(SlotIndexRequireUint): +class HWSlotEnableSet(SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() self.add_slot_args(parser) @@ -891,7 +891,7 @@ class HWSlotEnableSet(SlotIndexRequireUint): print(f' - Set slot {slot_num} {"enable" if enable else "disable"} success.') -class LFEMSim(LFEMCardRequiredUint): +class LFEMSim(LFEMCardRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -905,7 +905,7 @@ class LFEMSim(LFEMCardRequiredUint): print(f' - Set em410x tag id success.') -class HWSlotNickSet(SlotIndexRequireUint, SenseTypeRequireUint): +class HWSlotNickSet(SlotIndexRequireUnit, SenseTypeRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() self.add_slot_args(parser) @@ -924,7 +924,7 @@ class HWSlotNickSet(SlotIndexRequireUint, SenseTypeRequireUint): print(f' - Set tag nick name for slot {slot_num} success.') -class HWSlotNickGet(SlotIndexRequireUint, SenseTypeRequireUint): +class HWSlotNickGet(SlotIndexRequireUnit, SenseTypeRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() self.add_slot_args(parser) From 3259de2f25287003509d96ecfd8190775bd0c85b Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Wed, 9 Aug 2023 23:37:52 +0200 Subject: [PATCH 2/9] Replace the PositiveChameleonCMD class with an exception decorator --- software/script/chameleon_cli_unit.py | 86 +++++++------- software/script/chameleon_cmd.py | 162 +++++--------------------- software/script/chameleon_utils.py | 31 +++++ 3 files changed, 98 insertions(+), 181 deletions(-) create mode 100644 software/script/chameleon_utils.py diff --git a/software/script/chameleon_cli_unit.py b/software/script/chameleon_cli_unit.py index 4de0681..f276772 100644 --- a/software/script/chameleon_cli_unit.py +++ b/software/script/chameleon_cli_unit.py @@ -59,12 +59,8 @@ class BaseCLIUnit: self._device_com = com @property - def cmd_positive(self) -> chameleon_cmd.BaseChameleonCMD: - return chameleon_cmd.PositiveChameleonCMD(self.device_com) - - @property - def cmd_standard(self) -> chameleon_cmd.BaseChameleonCMD: - return chameleon_cmd.BaseChameleonCMD(self.device_com) + def cmd(self) -> chameleon_cmd.ChameleonCMD: + return chameleon_cmd.ChameleonCMD(self.device_com) def args_parser(self) -> ArgumentParserNoExit or None: """ @@ -172,7 +168,7 @@ class ReaderRequiredUnit(DeviceRequiredUnit): def before_exec(self, args: argparse.Namespace): if super(ReaderRequiredUnit, self).before_exec(args): - ret = self.cmd_standard.is_reader_device_mode() + ret = self.cmd.is_reader_device_mode() if ret: return True else: @@ -221,10 +217,10 @@ class HWModeSet(DeviceRequiredUnit): def on_exec(self, args: argparse.Namespace): if args.mode == 'reader' or args.mode == 'r': - self.cmd_standard.set_reader_device_mode(True) + self.cmd.set_reader_device_mode(True) print("Switch to { Tag Reader } mode successfully.") else: - self.cmd_standard.set_reader_device_mode(False) + self.cmd.set_reader_device_mode(False) print("Switch to { Tag Emulator } mode successfully.") @@ -233,7 +229,7 @@ class HWModeGet(DeviceRequiredUnit): pass def on_exec(self, args: argparse.Namespace): - print(f"- Device Mode ( Tag {'Reader' if self.cmd_standard.is_reader_device_mode() else 'Emulator'} )") + print(f"- Device Mode ( Tag {'Reader' if self.cmd.is_reader_device_mode() else 'Emulator'} )") class HWChipIdGet(DeviceRequiredUnit): @@ -242,7 +238,7 @@ class HWChipIdGet(DeviceRequiredUnit): return None def on_exec(self, args: argparse.Namespace): - print(f' - Device chip ID: ' + self.cmd_positive.get_device_chip_id()) + print(f' - Device chip ID: ' + self.cmd.get_device_chip_id()) class HWAddressGet(DeviceRequiredUnit): @@ -251,7 +247,7 @@ class HWAddressGet(DeviceRequiredUnit): return None def on_exec(self, args: argparse.Namespace): - print(f' - Device address: ' + self.cmd_positive.get_device_address()) + print(f' - Device address: ' + self.cmd.get_device_address()) class HF14AScan(ReaderRequiredUnit): @@ -259,7 +255,7 @@ class HF14AScan(ReaderRequiredUnit): pass def scan(self): - resp: chameleon_com.Response = self.cmd_standard.scan_tag_14a() + resp: chameleon_com.Response = self.cmd.scan_tag_14a() if resp.status == chameleon_status.Device.HF_TAG_OK: info = chameleon_cstruct.parse_14a_scan_tag_result(resp.data) print(f"- UID Size: {info['uid_size']}") @@ -282,11 +278,11 @@ class HF14AInfo(ReaderRequiredUnit): def info(self): # detect mf1 support - resp = self.cmd_positive.detect_mf1_support() + resp = self.cmd.detect_mf1_support() if resp.status == chameleon_status.Device.HF_TAG_OK: # detect prng print("- Mifare Classic technology") - resp = self.cmd_standard.detect_mf1_nt_level() + resp = self.cmd.detect_mf1_nt_level() if resp.status == 0x00: prng_level = "Weak" elif resp.status == 0x24: @@ -336,8 +332,8 @@ class HFMFNested(ReaderRequiredUnit): :return: """ # acquire - dist_resp = self.cmd_positive.detect_nt_distance(block_known, type_known, key_known) - nt_resp = self.cmd_positive.acquire_nested(block_known, type_known, key_known, block_target, type_target) + dist_resp = self.cmd.detect_nt_distance(block_known, type_known, key_known) + nt_resp = self.cmd.acquire_nested(block_known, type_known, key_known, block_target, type_target) # parse dist_obj = chameleon_cstruct.parse_nt_distance_detect_result(dist_resp.data) nt_obj = chameleon_cstruct.parse_nested_nt_acquire_group(nt_resp.data) @@ -371,7 +367,7 @@ class HFMFNested(ReaderRequiredUnit): print(f" - [{len(key_list)} candidate keys found ]") for key in key_list: key_bytes = bytearray.fromhex(key) - ret = self.cmd_standard.auth_mf1_key(block_target, type_target, key_bytes) + ret = self.cmd.auth_mf1_key(block_target, type_target, key_bytes) if ret.status == chameleon_status.Device.HF_TAG_OK: return key else: @@ -429,7 +425,7 @@ class HFMFDarkside(ReaderRequiredUnit): first_recover = True retry_count = 0 while retry_count < 0xFF: - darkside_resp = self.cmd_positive.acquire_darkside(block_target, type_target, first_recover, 15) + darkside_resp = self.cmd.acquire_darkside(block_target, type_target, first_recover, 15) first_recover = False # not first run. darkside_obj = chameleon_cstruct.parse_darkside_acquire_result(darkside_resp.data) self.darkside_list.append(darkside_obj) @@ -462,7 +458,7 @@ class HFMFDarkside(ReaderRequiredUnit): # auth key for key in key_list: key_bytes = bytearray.fromhex(key) - auth_ret = self.cmd_positive.auth_mf1_key(block_target, type_target, key_bytes) + auth_ret = self.cmd.auth_mf1_key(block_target, type_target, key_bytes) if auth_ret.status == chameleon_status.Device.HF_TAG_OK: return key return None @@ -509,7 +505,7 @@ class HFMFRDBL(BaseMF1AuthOpera): # hf mf rdbl -b 2 -t A -k FFFFFFFFFFFF def on_exec(self, args: argparse.Namespace): param = self.get_param(args) - resp = self.cmd_positive.read_mf1_block(param.block, param.type, param.key) + resp = self.cmd.read_mf1_block(param.block, param.type, param.key) print(f" - Data: {resp.data.hex()}") @@ -527,7 +523,7 @@ class HFMFWRBL(BaseMF1AuthOpera): if not re.match(r"^[a-fA-F0-9]{32}$", args.data): raise ArgsParserError("Data must include 32 HEX symbols") param.data = bytearray.fromhex(args.data) - resp = self.cmd_standard.write_mf1_block(param.block, param.type, param.key, param.data) + resp = self.cmd.write_mf1_block(param.block, param.type, param.key, param.data) if resp.status == chameleon_status.Device.HF_TAG_OK: print(f" - {colorama.Fore.GREEN}Write done.{colorama.Style.RESET_ALL}") else: @@ -545,7 +541,7 @@ class HFMFDetectionEnable(DeviceRequiredUnit): # hf mf detection enable -e 1 def on_exec(self, args: argparse.Namespace): enable = True if args.enable == 1 else False - self.cmd_positive.set_mf1_detection_enable(enable) + self.cmd.set_mf1_detection_enable(enable) print(f" - Set mf1 detection { 'enable' if enable else 'disable'}.") @@ -556,7 +552,7 @@ class HFMFDetectionLogCount(DeviceRequiredUnit): # hf mf detection count def on_exec(self, args: argparse.Namespace): - data_bytes = self.cmd_standard.get_mf1_detection_count().data + data_bytes = self.cmd.get_mf1_detection_count().data count = int.from_bytes(data_bytes, "little", signed=False) print(f" - MF1 detection log count = {count}") @@ -604,13 +600,13 @@ class HFMFDetectionDecrypt(DeviceRequiredUnit): def on_exec(self, args: argparse.Namespace): buffer = bytearray() index = 0 - count = int.from_bytes(self.cmd_standard.get_mf1_detection_count().data, "little", signed=False) + count = int.from_bytes(self.cmd.get_mf1_detection_count().data, "little", signed=False) if count == 0: print(" - No detection log to download") return print(f" - MF1 detection log count = {count}, start download", end="") while index < count: - tmp = self.cmd_positive.get_mf1_detection_log(index).data + tmp = self.cmd.get_mf1_detection_log(index).data recv_count = int(len(tmp) / HFMFDetectionDecrypt.detection_log_size) index += recv_count buffer.extend(tmp) @@ -684,7 +680,7 @@ class HFMFELoad(DeviceRequiredUnit): block_data = buffer[index: index + 16] index += 16 # load to device - self.cmd_positive.set_mf1_block_data(block, block_data) + self.cmd.set_mf1_block_data(block, block_data) print('.', end='') block += 1 print("\n - Load success") @@ -723,7 +719,7 @@ class HFMFSim(DeviceRequiredUnit): else: raise Exception("UID must be hex") - self.cmd_positive.set_mf1_anti_collision_res(sak, atqa, uid) + self.cmd.set_mf1_anti_collision_res(sak, atqa, uid) print(" - Set anti-collision resources success") @@ -733,7 +729,7 @@ class LFEMRead(ReaderRequiredUnit): return None def on_exec(self, args: argparse.Namespace): - resp = self.cmd_positive.read_em_410x() + resp = self.cmd.read_em_410x() id_hex = resp.data.hex() print(f" - EM410x ID(10H): {colorama.Fore.GREEN}{id_hex}{colorama.Style.RESET_ALL}") @@ -774,7 +770,7 @@ class LFEMWriteT55xx(LFEMCardRequiredUnit, ReaderRequiredUnit): def on_exec(self, args: argparse.Namespace): id_hex = args.id id_bytes = bytearray.fromhex(id_hex) - self.cmd_positive.write_em_410x_to_t55xx(id_bytes) + self.cmd.write_em_410x_to_t55xx(id_bytes) print(f" - EM410x ID(10H): {id_hex} write done.") @@ -818,7 +814,7 @@ class HWSlotSet(SlotIndexRequireUnit): # hw slot change -s 1 def on_exec(self, args: argparse.Namespace): slot_index = args.slot - self.cmd_positive.set_slot_activated(slot_index) + self.cmd.set_slot_activated(slot_index) print(f" - Set slot {slot_index} activated success.") @@ -855,7 +851,7 @@ class HWSlotTagType(TagTypeRequiredUnit, SlotIndexRequireUnit): def on_exec(self, args: argparse.Namespace): tag_type = args.type slot_index = args.slot - self.cmd_positive.set_slot_tag_type(slot_index, tag_type) + self.cmd.set_slot_tag_type(slot_index, tag_type) print(f' - Set slot tag type success.') @@ -872,7 +868,7 @@ class HWSlotDataDefault(TagTypeRequiredUnit, SlotIndexRequireUnit): def on_exec(self, args: argparse.Namespace): tag_type = args.type slot_num = args.slot - self.cmd_positive.set_slot_data_default(slot_num, tag_type) + self.cmd.set_slot_data_default(slot_num, tag_type) print(f' - Set slot tag data init success.') @@ -887,7 +883,7 @@ class HWSlotEnableSet(SlotIndexRequireUnit): def on_exec(self, args: argparse.Namespace): slot_num = args.slot enable = args.enable - self.cmd_positive.set_slot_enable(slot_num, enable) + self.cmd.set_slot_enable(slot_num, enable) print(f' - Set slot {slot_num} {"enable" if enable else "disable"} success.') @@ -901,7 +897,7 @@ class LFEMSim(LFEMCardRequiredUnit): def on_exec(self, args: argparse.Namespace): id_hex = args.id id_bytes = bytearray.fromhex(id_hex) - self.cmd_positive.set_em140x_sim_id(id_bytes) + self.cmd.set_em140x_sim_id(id_bytes) print(f' - Set em410x tag id success.') @@ -920,7 +916,7 @@ class HWSlotNickSet(SlotIndexRequireUnit, SenseTypeRequireUnit): name: str = args.name if len(name.encode(encoding="gbk")) > 32: raise ValueError("Your tag nick name too long.") - self.cmd_positive.set_slot_tag_nick_name(slot_num, sense_type, name) + self.cmd.set_slot_tag_nick_name(slot_num, sense_type, name) print(f' - Set tag nick name for slot {slot_num} success.') @@ -935,7 +931,7 @@ class HWSlotNickGet(SlotIndexRequireUnit, SenseTypeRequireUnit): def on_exec(self, args: argparse.Namespace): slot_num = args.slot sense_type = args.sense_type - res = self.cmd_positive.get_slot_tag_nick_name(slot_num, sense_type) + res = self.cmd.get_slot_tag_nick_name(slot_num, sense_type) print(f' - Get tag nick name for slot {slot_num}: {res.data.decode(encoding="gbk")}') @@ -946,7 +942,7 @@ class HWSlotUpdate(DeviceRequiredUnit): # hw slot update def on_exec(self, args: argparse.Namespace): - self.cmd_positive.update_slot_data_config() + self.cmd.update_slot_data_config() print(f' - Update config and data from device memory to flash success.') @@ -965,17 +961,17 @@ class HWSlotOpenAll(DeviceRequiredUnit): for slot in range(1,9): print(f' Slot{slot} setting...') # first to set tag type - self.cmd_positive.set_slot_tag_type(slot, hf_type) - self.cmd_positive.set_slot_tag_type(slot, lf_type) + self.cmd.set_slot_tag_type(slot, hf_type) + self.cmd.set_slot_tag_type(slot, lf_type) # to init default data - self.cmd_positive.set_slot_data_default(slot, hf_type) - self.cmd_positive.set_slot_data_default(slot, lf_type) + self.cmd.set_slot_data_default(slot, hf_type) + self.cmd.set_slot_data_default(slot, lf_type) # finally, we can enable this slot. - self.cmd_positive.set_slot_enable(slot, True) + self.cmd.set_slot_enable(slot, True) print(f' Open slot{slot} finish') # update config and save to flash - self.cmd_positive.update_slot_data_config() + self.cmd.update_slot_data_config() print(f' - Open all slot and set data to default success.') @@ -987,7 +983,7 @@ class HWDFU(DeviceRequiredUnit): # hw dfu def on_exec(self, args: argparse.Namespace): print("Application restarting...") - self.cmd_standard.enter_dfu_mode() + self.cmd.enter_dfu_mode() # 理论上，上面的指令执行完成后，dfu模式会进入，然后USB会重启， # 我们判断是否成功进入USB，只需要判断USB是否变成DFU设备的VID和PID即可， # 同时我们记得确认设备的信息，一致时才是同一个设备。 diff --git a/software/script/chameleon_cmd.py b/software/script/chameleon_cmd.py index 0f793d5..474cd34 100644 --- a/software/script/chameleon_cmd.py +++ b/software/script/chameleon_cmd.py @@ -2,6 +2,7 @@ import enum import chameleon_com import chameleon_status +from chameleon_utils import NegativeResponseError, expect_response DATA_CMD_GET_APP_VERSION = 1000 DATA_CMD_CHANGE_MODE = 1001 @@ -77,7 +78,7 @@ class TagSpecificType(enum.IntEnum): return enum_list -class BaseChameleonCMD: +class ChameleonCMD: """ Chameleon cmd function """ @@ -126,6 +127,7 @@ class BaseChameleonCMD: """ return self.device.send_cmd_sync(DATA_CMD_CHANGE_MODE, 0x00, 0x0001 if reader_mode else 0x0000) + @expect_response(chameleon_status.Device.HF_TAG_OK) def scan_tag_14a(self): """ 扫描场内的14a标签 @@ -154,6 +156,7 @@ class BaseChameleonCMD: """ return self.device.send_cmd_sync(DATA_CMD_MF1_DARKSIDE_DETECT, 0x00, None, timeout=20) + @expect_response(chameleon_status.Device.HF_TAG_OK) def detect_nt_distance(self, block_known, type_known, key_known): """ 检测卡片的随机数距离 @@ -165,6 +168,7 @@ class BaseChameleonCMD: data.extend(key_known) return self.device.send_cmd_sync(DATA_CMD_MF1_NT_DIST_DETECT, 0x00, data) + @expect_response(chameleon_status.Device.HF_TAG_OK) def acquire_nested(self, block_known, type_known, key_known, block_target, type_target): """ 采集Nested解密需要的关键NT参数 @@ -178,6 +182,7 @@ class BaseChameleonCMD: data.append(block_target) return self.device.send_cmd_sync(DATA_CMD_MF1_NESTED_ACQUIRE, 0x00, data) + @expect_response(chameleon_status.Device.HF_TAG_OK) def acquire_darkside(self, block_target, type_target, first_recover: int or bool, sync_max): """ 采集Darkside解密需要的关键参数 @@ -196,6 +201,10 @@ class BaseChameleonCMD: data.append(sync_max) return self.device.send_cmd_sync(DATA_CMD_MF1_DARKSIDE_ACQUIRE, 0x00, data, timeout=sync_max + 5) + @expect_response([ + chameleon_status.Device.HF_TAG_OK, + chameleon_status.Device.MF_ERRAUTH, + ]) def auth_mf1_key(self, block, type_value, key): """ 验证mf1秘钥，只验证单个扇区的指定类型的秘钥 @@ -210,6 +219,7 @@ class BaseChameleonCMD: data.extend(key) return self.device.send_cmd_sync(DATA_CMD_MF1_CHECK_ONE_KEY_BLOCK, 0x00, data) + @expect_response(chameleon_status.Device.HF_TAG_OK) def read_mf1_block(self, block, type_value, key): """ 读取mf1单块 @@ -224,6 +234,7 @@ class BaseChameleonCMD: data.extend(key) return self.device.send_cmd_sync(DATA_CMD_MF1_READ_ONE_BLOCK, 0x00, data) + @expect_response(chameleon_status.Device.HF_TAG_OK) def write_mf1_block(self, block, type_value, key, block_data): """ 写入mf1单块 @@ -240,6 +251,7 @@ class BaseChameleonCMD: data.extend(block_data) return self.device.send_cmd_sync(DATA_CMD_MF1_WRITE_ONE_BLOCK, 0x00, data) + @expect_response(chameleon_status.Device.LF_TAG_OK) def read_em_410x(self): """ 读取EM410X的卡号 @@ -247,6 +259,7 @@ class BaseChameleonCMD: """ return self.device.send_cmd_sync(DATA_CMD_SCAN_EM410X_TAG, 0x00, None) + @expect_response(chameleon_status.Device.LF_TAG_OK) def write_em_410x_to_t55xx(self, id_bytes: bytearray): """ 写入EM410X卡号到T55XX中 @@ -267,6 +280,7 @@ class BaseChameleonCMD: data.extend(key) return self.device.send_cmd_sync(DATA_CMD_WRITE_EM410X_TO_T5577, 0x00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_slot_activated(self, slot_index): """ 设置当前激活使用的卡槽 @@ -279,6 +293,7 @@ class BaseChameleonCMD: data.append(slot_index - 1) return self.device.send_cmd_sync(DATA_CMD_SET_SLOT_ACTIVATED, 0x00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_slot_tag_type(self, slot_index: int, tag_type: TagSpecificType): """ 设置当前卡槽的模拟卡的标签类型 @@ -294,6 +309,7 @@ class BaseChameleonCMD: data.append(tag_type) return self.device.send_cmd_sync(DATA_CMD_SET_SLOT_TAG_TYPE, 0x00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_slot_data_default(self, slot_index: int, tag_type: TagSpecificType): """ 设置指定卡槽的模拟卡的数据为缺省数据 @@ -309,6 +325,7 @@ class BaseChameleonCMD: data.append(tag_type) return self.device.send_cmd_sync(DATA_CMD_SET_SLOT_DATA_DEFAULT, 0x00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_slot_enable(self, slot_index: int, enable: bool): """ 设置指定的卡槽是否使能 @@ -323,6 +340,7 @@ class BaseChameleonCMD: data.append(0x01 if enable else 0x00) return self.device.send_cmd_sync(DATA_CMD_SET_SLOT_ENABLE, 0X00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_em140x_sim_id(self, id_bytes: bytearray): """ 设置EM410x模拟的卡号 @@ -333,6 +351,7 @@ class BaseChameleonCMD: raise ValueError("The id bytes length must equal 5") return self.device.send_cmd_sync(DATA_CMD_SET_EM410X_EMU_ID, 0x00, id_bytes) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_mf1_detection_enable(self, enable: bool): """ 设置是否使能当前卡槽的侦测 @@ -350,6 +369,7 @@ class BaseChameleonCMD: """ return self.device.send_cmd_sync(DATA_CMD_GET_MF1_DETECTION_COUNT, 0x00, None) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def get_mf1_detection_log(self, index: int): """ 从指定的index位置开始获取侦测日志 @@ -360,6 +380,7 @@ class BaseChameleonCMD: data.extend(index.to_bytes(4, "big", signed=False)) return self.device.send_cmd_sync(DATA_CMD_GET_MF1_DETECTION_RESULT, 0x00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_mf1_block_data(self, block_start: int, block_data: bytearray): """ 设置MF1的模拟卡的块数据 @@ -372,6 +393,7 @@ class BaseChameleonCMD: data.extend(block_data) return self.device.send_cmd_sync(DATA_CMD_LOAD_MF1_BLOCK_DATA, 0x00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_mf1_anti_collision_res(self, sak: bytearray, atqa: bytearray, uid: bytearray): """ 设置MF1的模拟卡的防冲撞资源信息 @@ -386,6 +408,7 @@ class BaseChameleonCMD: data.extend(uid) return self.device.send_cmd_sync(DATA_CMD_SET_MF1_ANTI_COLLISION_RES, 0X00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def set_slot_tag_nick_name(self, slot: int, sense_type: int, name: str): """ 设置MF1的模拟卡的防冲撞资源信息 @@ -399,6 +422,7 @@ class BaseChameleonCMD: data.extend(name.encode(encoding="gbk")) return self.device.send_cmd_sync(DATA_CMD_SET_SLOT_TAG_NICK, 0x00, data) + @expect_response(chameleon_status.Device.STATUS_DEVICE_SUCCESS) def get_slot_tag_nick_name(self, slot: int, sense_type: int): """ 设置MF1的模拟卡的防冲撞资源信息 @@ -426,145 +450,11 @@ class BaseChameleonCMD: return self.device.send_cmd_auto(DATA_CMD_ENTER_BOOTLOADER, 0x00, close=True) -class NegativeResponseError(Exception): - """ - Not positive response - """ - - -class PositiveChameleonCMD(BaseChameleonCMD): - """ - 子类重写基础指令交互实现类，针对每个指令进行单独封装结果处理 - 如果结果是成功状态，那么就返回对应的数据，否则直接抛出异常 - """ - - @staticmethod - def check_status(status_ret, status_except): - """ - 检查状态码，如果在接受为成功的 - :param status_ret: 执行指令之后返回的状态码 - :param status_except: 可以认为是执行成功的状态码 - :return: - """ - if isinstance(status_except, int): - status_except = [status_except] - if status_ret not in status_except: - if status_ret in chameleon_status.Device and status_ret in chameleon_status.message: - raise NegativeResponseError(chameleon_status.message[status_ret]) - else: - raise NegativeResponseError(f"Not positive response and unknown status {status_ret}") - return - - def scan_tag_14a(self): - ret = super(PositiveChameleonCMD, self).scan_tag_14a() - self.check_status(ret.status, chameleon_status.Device.HF_TAG_OK) - return ret - - def detect_nt_distance(self, block_known, type_known, key_known): - ret = super(PositiveChameleonCMD, self).detect_nt_distance(block_known, type_known, key_known) - self.check_status(ret.status, chameleon_status.Device.HF_TAG_OK) - return ret - - def acquire_nested(self, block_known, type_known, key_known, block_target, type_target): - ret = super(PositiveChameleonCMD, self).acquire_nested( - block_known, type_known, key_known, block_target, type_target) - self.check_status(ret.status, chameleon_status.Device.HF_TAG_OK) - return ret - - def acquire_darkside(self, block_target, type_target, first_recover: int or bool, sync_max): - ret = super(PositiveChameleonCMD, self).acquire_darkside(block_target, type_target, first_recover, sync_max) - self.check_status(ret.status, chameleon_status.Device.HF_TAG_OK) - return ret - - def auth_mf1_key(self, block, type_value, key): - ret = super(PositiveChameleonCMD, self).auth_mf1_key(block, type_value, key) - self.check_status(ret.status, [ - chameleon_status.Device.HF_TAG_OK, - chameleon_status.Device.MF_ERRAUTH, - ]) - return ret - - def read_mf1_block(self, block, type_value, key): - ret = super(PositiveChameleonCMD, self).read_mf1_block(block, type_value, key) - self.check_status(ret.status, chameleon_status.Device.HF_TAG_OK) - return ret - - def write_mf1_block(self, block, type_value, key, block_data): - ret = super(PositiveChameleonCMD, self).write_mf1_block(block, type_value, key, block_data) - self.check_status(ret.status, chameleon_status.Device.HF_TAG_OK) - return ret - - def read_em_410x(self): - ret = super(PositiveChameleonCMD, self).read_em_410x() - self.check_status(ret.status, chameleon_status.Device.LF_TAG_OK) - return ret - - def write_em_410x_to_t55xx(self, id_bytes: bytearray): - ret = super(PositiveChameleonCMD, self).write_em_410x_to_t55xx(id_bytes) - self.check_status(ret.status, chameleon_status.Device.LF_TAG_OK) - return ret - - def set_slot_activated(self, slot_index): - ret = super(PositiveChameleonCMD, self).set_slot_activated(slot_index) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_slot_tag_type(self, slot_index: int, tag_type: TagSpecificType): - ret = super(PositiveChameleonCMD, self).set_slot_tag_type(slot_index, tag_type) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_slot_data_default(self, slot_index: int, tag_type: TagSpecificType): - ret = super(PositiveChameleonCMD, self).set_slot_data_default(slot_index, tag_type) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_slot_enable(self, slot_index: int, enable: bool): - ret = super(PositiveChameleonCMD, self).set_slot_enable(slot_index, enable) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_em140x_sim_id(self, id_bytes: bytearray): - ret = super(PositiveChameleonCMD, self).set_em140x_sim_id(id_bytes) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_mf1_detection_enable(self, enable: bool): - ret = super(PositiveChameleonCMD, self).set_mf1_detection_enable(enable) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def get_mf1_detection_log(self, index: int): - ret = super(PositiveChameleonCMD, self).get_mf1_detection_log(index) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_mf1_block_data(self, block_start: int, data: bytearray): - ret = super(PositiveChameleonCMD, self).set_mf1_block_data(block_start, data) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_mf1_anti_collision_res(self, sak: int, atqa: bytearray, uid: bytearray): - ret = super(PositiveChameleonCMD, self).set_mf1_anti_collision_res(sak, atqa, uid) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def set_slot_tag_nick_name(self, slot: int, sense_type: int, name: str): - ret = super(PositiveChameleonCMD, self).set_slot_tag_nick_name(slot, sense_type, name) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - def get_slot_tag_nick_name(self, slot: int, sense_type: int): - ret = super(PositiveChameleonCMD, self).get_slot_tag_nick_name(slot, sense_type) - self.check_status(ret.status, chameleon_status.Device.STATUS_DEVICE_SUCCESS) - return ret - - if __name__ == '__main__': # connect to chameleon dev = chameleon_com.ChameleonCom() dev.open("com19") - cml = BaseChameleonCMD(dev) + cml = ChameleonCMD(dev) ver = cml.get_firmware_version() print(f"Firmware number of application: {ver}") id = cml.get_device_chip_id() diff --git a/software/script/chameleon_utils.py b/software/script/chameleon_utils.py new file mode 100644 index 0000000..219f205 --- /dev/null +++ b/software/script/chameleon_utils.py @@ -0,0 +1,31 @@ +from functools import wraps +from typing import Union + +import chameleon_status + + +class NegativeResponseError(Exception): + """ + Not positive response + """ + +def expect_response(accepted_responses): + if isinstance(accepted_responses, int): + accepted_responses = [accepted_responses] + + def decorator(func): + @wraps(func) + def error_throwing_func(*args, **kwargs): + ret = func(*args, **kwargs) + + if ret.status not in accepted_responses: + if ret.status in chameleon_status.Device and ret.status in chameleon_status.message: + raise NegativeResponseError(chameleon_status.message[ret.status]) + else: + raise NegativeResponseError(f"Not positive response and unknown status {ret.status}") + + return ret + + return error_throwing_func + + return decorator From e9fec48ede4599d82ce064e5f57d0dfc4379b1c8 Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Wed, 9 Aug 2023 23:51:12 +0200 Subject: [PATCH 3/9] Fix docstring for expect_response and rename the associated exception --- software/script/chameleon_cli_main.py | 2 +- software/script/chameleon_cmd.py | 2 +- software/script/chameleon_utils.py | 14 +++++++++----- 3 files changed, 11 insertions(+), 7 deletions(-) diff --git a/software/script/chameleon_cli_main.py b/software/script/chameleon_cli_main.py index 082ad1d..aec9d08 100755 --- a/software/script/chameleon_cli_main.py +++ b/software/script/chameleon_cli_main.py @@ -217,7 +217,7 @@ class ChameleonCLI: continue # start process cmd unit.on_exec(args_parse_result) - except (chameleon_cmd.NegativeResponseError, chameleon_cli_unit.ArgsParserError) as e: + except (chameleon_cmd.UnexpectedResponseError, chameleon_cli_unit.ArgsParserError) as e: print(f"{colorama.Fore.RED}{str(e)}{colorama.Style.RESET_ALL}") except Exception: print(f"CLI exception: {colorama.Fore.RED}{traceback.format_exc()}{colorama.Style.RESET_ALL}") diff --git a/software/script/chameleon_cmd.py b/software/script/chameleon_cmd.py index 474cd34..f7ae20d 100644 --- a/software/script/chameleon_cmd.py +++ b/software/script/chameleon_cmd.py @@ -2,7 +2,7 @@ import enum import chameleon_com import chameleon_status -from chameleon_utils import NegativeResponseError, expect_response +from chameleon_utils import UnexpectedResponseError, expect_response DATA_CMD_GET_APP_VERSION = 1000 DATA_CMD_CHANGE_MODE = 1001 diff --git a/software/script/chameleon_utils.py b/software/script/chameleon_utils.py index 219f205..8cbfee2 100644 --- a/software/script/chameleon_utils.py +++ b/software/script/chameleon_utils.py @@ -4,12 +4,16 @@ from typing import Union import chameleon_status -class NegativeResponseError(Exception): +class UnexpectedResponseError(Exception): """ - Not positive response + Unexpected response exception """ -def expect_response(accepted_responses): +def expect_response(accepted_responses: Union[int, list[int]]): + """ + Decorator for wrapping a Chameleon CMD function to check its response + for expected return codes and throwing an exception otherwise + """ if isinstance(accepted_responses, int): accepted_responses = [accepted_responses] @@ -20,9 +24,9 @@ def expect_response(accepted_responses): if ret.status not in accepted_responses: if ret.status in chameleon_status.Device and ret.status in chameleon_status.message: - raise NegativeResponseError(chameleon_status.message[ret.status]) + raise UnexpectedResponseError(chameleon_status.message[ret.status]) else: - raise NegativeResponseError(f"Not positive response and unknown status {ret.status}") + raise UnexpectedResponseError(f"Unexpected response and unknown status {ret.status}") return ret From e2b5b9b96001e64adb06addd66db04d6fc409059 Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Thu, 10 Aug 2023 00:45:37 +0200 Subject: [PATCH 4/9] Use prompt-toolkit and switch to a different command registration system --- software/script/chameleon_cli_main.py | 244 ++++++++++---------------- software/script/chameleon_cli_unit.py | 58 +++++- software/script/chameleon_utils.py | 51 ++++++ 3 files changed, 191 insertions(+), 162 deletions(-) diff --git a/software/script/chameleon_cli_main.py b/software/script/chameleon_cli_main.py index aec9d08..8e29b23 100755 --- a/software/script/chameleon_cli_main.py +++ b/software/script/chameleon_cli_main.py @@ -8,7 +8,10 @@ import chameleon_com import chameleon_cmd import colorama import chameleon_cli_unit +import chameleon_utils import os +import prompt_toolkit +from prompt_toolkit.formatted_text import ANSI if os.name == 'posix': import readline @@ -37,94 +40,44 @@ BANNER = f""" """ -def new_unit(unit_clz, help_msg): - """ - Create a new unit dict object - :param unit_clz: unit implement class - :param help_msg: unit usage - :return: a dict... - """ - return { - 'unit': unit_clz, - 'help': help_msg, - } - - class ChameleonCLI: """ CLI for chameleon """ def __init__(self): - self.cmd_maps = { - 'hw': { - 'connect': new_unit(chameleon_cli_unit.HWConnect, "Connect to chameleon by serial port"), - 'chipid': { - 'get': new_unit(chameleon_cli_unit.HWChipIdGet, "Get device chipset ID"), - 'help': "Device chipsed ID get" - }, - 'address': { - 'get': new_unit(chameleon_cli_unit.HWAddressGet, "Get device address (used with Bluetooth)"), - 'help': "Device address get" - }, - 'mode': { - 'set': new_unit(chameleon_cli_unit.HWModeSet, "Change device mode to tag reader or tag emulator"), - 'get': new_unit(chameleon_cli_unit.HWModeGet, "Get current device mode"), - 'help': "Device mode get/set" - }, - 'slot': { - 'change': new_unit(chameleon_cli_unit.HWSlotSet, "Set emulation tag slot activated."), - 'type': new_unit(chameleon_cli_unit.HWSlotTagType, "Set emulation tag type"), - 'init': new_unit(chameleon_cli_unit.HWSlotDataDefault, "Set emulation tag data to default"), - 'enable': new_unit(chameleon_cli_unit.HWSlotEnableSet, "Set emulation tag slot enable or disable"), - 'nick': { - 'set': new_unit(chameleon_cli_unit.HWSlotNickSet, "Set tag nick name for slot"), - 'get': new_unit(chameleon_cli_unit.HWSlotNickGet, "Get tag nick name for slot"), - 'help': "Get/Set tag nick name for slot", - }, - 'update': new_unit(chameleon_cli_unit.HWSlotUpdate, "Update config & data to device flash"), - 'openall': new_unit(chameleon_cli_unit.HWSlotOpenAll, "Open all slot and set to default data"), - 'help': "Emulation tag slot.", - }, - 'dfu': new_unit(chameleon_cli_unit.HWDFU, "Restart application to bootloader mode(Not yet implement dfu)."), - 'help': "hardware controller", - }, - 'hf': { - '14a': { - 'scan': new_unit(chameleon_cli_unit.HF14AScan, "Scan 14a tag, and print basic information"), - 'info': new_unit(chameleon_cli_unit.HF14AInfo, "Scan 14a tag, and print detail information"), - 'help': "ISO14443-a tag read/write/info...", - }, - 'mf': { - 'nested': new_unit(chameleon_cli_unit.HFMFNested, "Mifare Classic nested recover key"), - 'darkside': new_unit(chameleon_cli_unit.HFMFDarkside, "Mifare Classic darkside recover key"), - 'rdbl': new_unit(chameleon_cli_unit.HFMFRDBL, "MiFARE Classic read one block"), - 'wrbl': new_unit(chameleon_cli_unit.HFMFWRBL, "MiFARE Classic write one block"), - 'detection': { - 'enable': new_unit(chameleon_cli_unit.HFMFDetectionEnable, "Detection enable"), - 'count': new_unit(chameleon_cli_unit.HFMFDetectionLogCount, "Detection log count"), - 'decrypt': new_unit(chameleon_cli_unit.HFMFDetectionDecrypt, "Download log and decrypt keys"), - 'help': "Mifare Classic detection log" - }, - 'sim': new_unit(chameleon_cli_unit.HFMFSim, "Simulation a mifare classic card"), - 'eload': new_unit(chameleon_cli_unit.HFMFELoad, "Load data to emulator memory"), - 'help': "Mifare Classic mini/1/2/4, attack/read/write" - }, - 'help': "high frequency tag/reader", - }, - 'lf': { - 'em': { - 'read': new_unit(chameleon_cli_unit.LFEMRead, "Scan em410x tag and print id"), - 'write': new_unit(chameleon_cli_unit.LFEMWriteT55xx, "Write em410x id to t55xx"), - 'sim': new_unit(chameleon_cli_unit.LFEMSim, "Simulation a em410x id card"), - 'help': "EM410x read/write/emulator", - }, - 'help': "low frequency tag/reader", - } - } + self.session = prompt_toolkit.PromptSession() # new a device communication instance(only communication) self.device_com = chameleon_com.ChameleonCom() + + def get_cmd_node(self, node: chameleon_utils.CLITree, cmdline: list[str]) -> tuple[chameleon_utils.CLITree, list[str]]: + """ + Recursively traverse the command line tree to get to the matching node + + :return: last matching CLITree node, remaining tokens + """ + # No more subcommands to parse, return node + if cmdline == []: + return node, [] + + for child in node.children: + if cmdline[0] == child.name: + return self.get_cmd_node(child, cmdline[1:]) + + # No matching child node + return node, cmdline[:] + + def get_prompt(self): + """ + Retrieve the cli prompt + + :return: current cmd prompt + """ + device_string = f"{colorama.Fore.GREEN}USB" if self.device_com.isOpen( + ) else f"{colorama.Fore.RED}Offline" + status = f"[{device_string}{colorama.Style.RESET_ALL}] chameleon --> " + return status @staticmethod def print_banner(): @@ -134,24 +87,6 @@ class ChameleonCLI: """ print(colorama.Fore.YELLOW + BANNER) - def parse_cli_cmd(self, cmd_str): - """ - parse cmd from str - :param cmd_str: - :return: - """ - cmds = cmd_str.split(" ") - cmd_maps: dict or types.FunctionType = self.cmd_maps - cmd_end = "" - for cmd in cmds: - if cmd in cmd_maps: # CMD found in map, we can continue find next - cmd_maps = cmd_maps[cmd] - cmd_end = cmd - else: # CMD not found - break - cmd_end_position = cmd_str.index(cmd_end) + len(cmd_end) + 1 - return cmd_maps, (cmd_str[:cmd_end_position], cmd_str[cmd_end_position:]) - def startCLI(self): """ start listen input. @@ -161,74 +96,73 @@ class ChameleonCLI: closing = False while True: # wait user input - status = f"{colorama.Fore.GREEN}USB" if self.device_com.isOpen() else f"{colorama.Fore.RED}Offline" - print(f"[{status}{colorama.Style.RESET_ALL}] chameleon --> ", end="") try: - cmd_str = input().strip() + cmd_str = self.session.prompt(ANSI(self.get_prompt())).strip() except EOFError: - print("") closing = True - - if closing or cmd_str == "exit" or cmd_str == "quit" or cmd_str.startswith('q', 0) or cmd_str.startswith('e', 0): + except KeyboardInterrupt: + closing = True + + if cmd_str in ["exit", "quit", "q", "e"] or closing: print("Bye, thank you. ^.^ ") self.device_com.close() sys.exit(996) - - # clear screen - if cmd_str == "clear": - if platform.system() == 'Windows': - os.system("cls") - elif platform.system() == 'Linux': - os.system("clear") - else: - print("No screen clear implement") + elif cmd_str == "clear": + os.system('clear' if os.name == 'posix' else 'cls') + continue + elif cmd_str == "": continue # parse cmd - cmd_map, args_str = self.parse_cli_cmd(cmd_str) - is_exec_map = 'unit' in cmd_map - if is_exec_map: - # new a unit instance - unit_clz = cmd_map['unit'] - if callable(unit_clz): - unit: chameleon_cli_unit.BaseCLIUnit = unit_clz() - else: - raise TypeError("CMD unit is not a 'BaseCLIUnit'") - # set variables of required - unit.device_com = self.device_com - # parse args - args_parse_result = unit.args_parser() - if args_parse_result is not None: - args: argparse.ArgumentParser = args_parse_result - args.prog = args_str[0] - try: - args_parse_result = args.parse_args(args_str[1].split()) - except chameleon_cli_unit.ArgsParserError as e: - args.print_usage() - print(str(e).strip(), end="\n\n") - continue - except chameleon_cli_unit.ParserExitIntercept: - # don't exit process. - continue - # noinspection PyBroadException - try: - # before process cmd, we need to do something... - if not unit.before_exec(args_parse_result): - continue - # start process cmd - unit.on_exec(args_parse_result) - except (chameleon_cmd.UnexpectedResponseError, chameleon_cli_unit.ArgsParserError) as e: - print(f"{colorama.Fore.RED}{str(e)}{colorama.Style.RESET_ALL}") - except Exception: - print(f"CLI exception: {colorama.Fore.RED}{traceback.format_exc()}{colorama.Style.RESET_ALL}") - elif isinstance(cmd_map, dict): + argv = cmd_str.split() + root_cmd = argv[0] + if root_cmd not in chameleon_cli_unit.root_commands: + # No matching command group print("".ljust(18, "-") + "".ljust(10) + "".ljust(30, "-")) - for map_key in cmd_map: - map_item = cmd_map[map_key] - if 'help' in map_item: - cmd_title = f"{colorama.Fore.GREEN}{map_key}{colorama.Style.RESET_ALL}" - help_line = (f" - {cmd_title}".ljust(37)) + f"[ {map_item['help']} ]" - print(help_line) + for cmd_name, cmd_node in chameleon_cli_unit.root_commands.items(): + cmd_title = f"{colorama.Fore.GREEN}{cmd_name}{colorama.Style.RESET_ALL}" + help_line = (f" - {cmd_title}".ljust(37)) + f"[ {cmd_node.helptext} ]" + print(help_line) + continue + + tree_node, arg_list = self.get_cmd_node(chameleon_cli_unit.root_commands[root_cmd], argv[1:]) + + if not tree_node.cls: + # Found tree node is a group without an implementation, print children + print("".ljust(18, "-") + "".ljust(10) + "".ljust(30, "-")) + for child in tree_node.children: + cmd_title = f"{colorama.Fore.GREEN}{child.name}{colorama.Style.RESET_ALL}" + help_line = (f" - {cmd_title}".ljust(37)) + f"[ {child.helptext} ]" + print(help_line) + continue + + unit: chameleon_cli_unit.BaseCLIUnit = tree_node.cls() + unit.device_com = self.device_com + args_parse_result = unit.args_parser() + + if args_parse_result is not None: + args: argparse.ArgumentParser = args_parse_result + args.prog = tree_node.fullname + try: + args_parse_result = args.parse_args(arg_list) + except chameleon_utils.ArgsParserError as e: + args.print_usage() + print(str(e).strip(), end="\n\n") + continue + except chameleon_utils.ParserExitIntercept: + # don't exit process. + continue + try: + # before process cmd, we need to do something... + if not unit.before_exec(args_parse_result): + continue + + # start process cmd + unit.on_exec(args_parse_result) + except (chameleon_utils.UnexpectedResponseError, chameleon_utils.ArgsParserError) as e: + print(f"{colorama.Fore.RED}{str(e)}{colorama.Style.RESET_ALL}") + except Exception: + print(f"CLI exception: {colorama.Fore.RED}{traceback.format_exc()}{colorama.Style.RESET_ALL}") if __name__ == '__main__': diff --git a/software/script/chameleon_cli_unit.py b/software/script/chameleon_cli_unit.py index f276772..110b5fd 100644 --- a/software/script/chameleon_cli_unit.py +++ b/software/script/chameleon_cli_unit.py @@ -12,17 +12,11 @@ import chameleon_com import chameleon_cmd import chameleon_cstruct import chameleon_status +from chameleon_utils import * description_public = "Please enter correct parameters" -class ArgsParserError(Exception): - pass - - -class ParserExitIntercept(Exception): - pass - class ArgumentParserNoExit(argparse.ArgumentParser): """ @@ -180,6 +174,29 @@ class ReaderRequiredUnit(DeviceRequiredUnit): raise NotImplementedError("Please implement this") +hw = CLITree('hw', 'hardware controller') +hw_chipid = hw.subgroup('chipid', 'Device chipsed ID get') +hw_address = hw.subgroup('address', 'Device address get') +hw_mode = hw.subgroup('mode', 'Device mode get/set') +hw_slot = hw.subgroup('slot', 'Emulation tag slot.') +hw_slot_nick = hw_slot.subgroup('nick', 'Get/Set tag nick name for slot') + +hf = CLITree('hf', 'high frequency tag/reader') +hf_14a = hf.subgroup('14a', 'ISO14443-a tag read/write/info...') +hf_mf = hf.subgroup('mf', 'Mifare Classic mini/1/2/4, attack/read/write') +hf_mf_detection = hf.subgroup( + 'detection', 'Mifare Classic detection log') + +lf = CLITree('lf', 'low frequency tag/reader') +lf_em = lf.subgroup('em', 'EM410x read/write/emulator') + +root_commands: dict[str, CLITree] = { + 'hw': hw, + 'hf': hf, + 'lf': lf, +} + +@hw.command('connect', 'Connect to chameleon by serial port') class HWConnect(BaseCLIUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -206,6 +223,7 @@ class HWConnect(BaseCLIUnit): print(f"Chameleon Connect fail: {str(e)}") +@hw_mode.command('set', 'Change device mode to tag reader or tag emulator') class HWModeSet(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -224,6 +242,7 @@ class HWModeSet(DeviceRequiredUnit): print("Switch to { Tag Emulator } mode successfully.") +@hw_mode.command('get', 'Get current device mode') class HWModeGet(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: pass @@ -232,6 +251,7 @@ class HWModeGet(DeviceRequiredUnit): print(f"- Device Mode ( Tag {'Reader' if self.cmd.is_reader_device_mode() else 'Emulator'} )") +@hw_chipid.command('get', 'Get device chipset ID') class HWChipIdGet(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -241,6 +261,7 @@ class HWChipIdGet(DeviceRequiredUnit): print(f' - Device chip ID: ' + self.cmd.get_device_chip_id()) +@hw_address.command('get', 'Get device address (used with Bluetooth)') class HWAddressGet(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -250,6 +271,7 @@ class HWAddressGet(DeviceRequiredUnit): print(f' - Device address: ' + self.cmd.get_device_address()) +@hf_14a.command('scan', 'Scan 14a tag, and print basic information') class HF14AScan(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: pass @@ -271,6 +293,7 @@ class HF14AScan(ReaderRequiredUnit): return self.scan() +@hf_14a.command('info', 'Scan 14a tag, and print detail information') class HF14AInfo(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -301,6 +324,7 @@ class HF14AInfo(ReaderRequiredUnit): self.info() +@hf_mf.command('nested', 'Mifare Classic nested recover key') class HFMFNested(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -406,6 +430,7 @@ class HFMFNested(ReaderRequiredUnit): return +@hf_mf.command('darkside', 'Mifare Classic darkside recover key') class HFMFDarkside(ReaderRequiredUnit): def __init__(self): @@ -500,6 +525,7 @@ class BaseMF1AuthOpera(ReaderRequiredUnit): raise NotImplementedError("Please implement this") +@hf_mf.command('rdbl', 'MiFARE Classic read one block') class HFMFRDBL(BaseMF1AuthOpera): # hf mf rdbl -b 2 -t A -k FFFFFFFFFFFF @@ -509,6 +535,7 @@ class HFMFRDBL(BaseMF1AuthOpera): print(f" - Data: {resp.data.hex()}") +@hf_mf.command('wrbl', 'MiFARE Classic write one block') class HFMFWRBL(BaseMF1AuthOpera): def args_parser(self) -> ArgumentParserNoExit or None: @@ -530,6 +557,7 @@ class HFMFWRBL(BaseMF1AuthOpera): print(f" - {colorama.Fore.RED}Write fail.{colorama.Style.RESET_ALL}") +@hf_mf_detection.command('enable', 'Detection enable') class HFMFDetectionEnable(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -545,6 +573,7 @@ class HFMFDetectionEnable(DeviceRequiredUnit): print(f" - Set mf1 detection { 'enable' if enable else 'disable'}.") +@hf_mf_detection.command('count', 'Detection log count') class HFMFDetectionLogCount(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -557,6 +586,7 @@ class HFMFDetectionLogCount(DeviceRequiredUnit): print(f" - MF1 detection log count = {count}") +@hf_mf_detection.command('decrypt', 'Download log and decrypt keys') class HFMFDetectionDecrypt(DeviceRequiredUnit): detection_log_size = 18 @@ -639,6 +669,7 @@ class HFMFDetectionDecrypt(DeviceRequiredUnit): return +@hf_mf.command('eload', 'Load data to emulator memory') class HFMFELoad(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -686,6 +717,7 @@ class HFMFELoad(DeviceRequiredUnit): print("\n - Load success") +@hf_mf.command('sim', 'Simulation a mifare classic card') class HFMFSim(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -723,6 +755,7 @@ class HFMFSim(DeviceRequiredUnit): print(" - Set anti-collision resources success") +@lf_em.command('read', 'Scan em410x tag and print id') class LFEMRead(ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -755,6 +788,7 @@ class LFEMCardRequiredUnit(DeviceRequiredUnit): raise NotImplementedError("Please implement this") +@lf_em.command('write', 'Write em410x id to t55xx') class LFEMWriteT55xx(LFEMCardRequiredUnit, ReaderRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -805,6 +839,7 @@ class SenseTypeRequireUnit(DeviceRequiredUnit): return parser +@hw_slot.command('change', 'Set emulation tag slot activated.') class HWSlotSet(SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -839,6 +874,7 @@ class TagTypeRequiredUnit(DeviceRequiredUnit): raise NotImplementedError() +@hw_slot.command('type', 'Set emulation tag type') class HWSlotTagType(TagTypeRequiredUnit, SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -855,6 +891,7 @@ class HWSlotTagType(TagTypeRequiredUnit, SlotIndexRequireUnit): print(f' - Set slot tag type success.') +@hw_slot.command('init', 'Set emulation tag data to default') class HWSlotDataDefault(TagTypeRequiredUnit, SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -872,6 +909,7 @@ class HWSlotDataDefault(TagTypeRequiredUnit, SlotIndexRequireUnit): print(f' - Set slot tag data init success.') +@hw_slot.command('enable', 'Set emulation tag slot enable or disable') class HWSlotEnableSet(SlotIndexRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -887,6 +925,7 @@ class HWSlotEnableSet(SlotIndexRequireUnit): print(f' - Set slot {slot_num} {"enable" if enable else "disable"} success.') +@lf_em.command('sim', 'Simulation a em410x id card') class LFEMSim(LFEMCardRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -901,6 +940,7 @@ class LFEMSim(LFEMCardRequiredUnit): print(f' - Set em410x tag id success.') +@hw_slot_nick.command('set', 'Set tag nick name for slot') class HWSlotNickSet(SlotIndexRequireUnit, SenseTypeRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -920,6 +960,7 @@ class HWSlotNickSet(SlotIndexRequireUnit, SenseTypeRequireUnit): print(f' - Set tag nick name for slot {slot_num} success.') +@hw_slot_nick.command('get', 'Get tag nick name for slot') class HWSlotNickGet(SlotIndexRequireUnit, SenseTypeRequireUnit): def args_parser(self) -> ArgumentParserNoExit or None: parser = ArgumentParserNoExit() @@ -935,6 +976,7 @@ class HWSlotNickGet(SlotIndexRequireUnit, SenseTypeRequireUnit): print(f' - Get tag nick name for slot {slot_num}: {res.data.decode(encoding="gbk")}') +@hw_slot.command('update', 'Update config & data to device flash') class HWSlotUpdate(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -946,6 +988,7 @@ class HWSlotUpdate(DeviceRequiredUnit): print(f' - Update config and data from device memory to flash success.') +@hw_slot.command('openall', 'Open all slot and set to default data') class HWSlotOpenAll(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: @@ -975,6 +1018,7 @@ class HWSlotOpenAll(DeviceRequiredUnit): print(f' - Open all slot and set data to default success.') +@hw.command('dfu', 'Restart application to bootloader mode(Not yet implement dfu).') class HWDFU(DeviceRequiredUnit): def args_parser(self) -> ArgumentParserNoExit or None: diff --git a/software/script/chameleon_utils.py b/software/script/chameleon_utils.py index 8cbfee2..99469b6 100644 --- a/software/script/chameleon_utils.py +++ b/software/script/chameleon_utils.py @@ -4,6 +4,14 @@ from typing import Union import chameleon_status +class ArgsParserError(Exception): + pass + + +class ParserExitIntercept(Exception): + pass + + class UnexpectedResponseError(Exception): """ Unexpected response exception @@ -33,3 +41,46 @@ def expect_response(accepted_responses: Union[int, list[int]]): return error_throwing_func return decorator + + +class CLITree: + """ + Class holding a + + :param name: Name of the command (e.g. "set") + :param helptext: Hint displayed for the command + :param fullname: Full name of the command that includes previous commands (e.g. "hw mode set") + :param cls: A BaseCLIUnit instance handling the command + """ + + def __init__(self, name=None, helptext=None, fullname=None, children=None, cls=None) -> None: + self.name: str = name + self.helptext: str = helptext + self.fullname: str = fullname if fullname else name + self.children: list[CLITree] = children if children else list() + self.cls = cls + + def subgroup(self, name, helptext=None): + """ + Create a child command group + + :param name: Name of the command group + :param helptext: Hint displayed for the group + """ + child = CLITree( + name=name, fullname=f'{self.fullname} {name}', helptext=helptext) + self.children.append(child) + return child + + def command(self, name, helptext=None): + """ + Create a child command + + :param name: Name of the command + :param helptext: Hint displayed for the command + """ + def decorator(cls): + self.children.append( + CLITree(name=name, fullname=f'{self.fullname} {name}', helptext=helptext, cls=cls)) + return cls + return decorator From 3789269c371f3fca440d90f80230c125dbf067bf Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Thu, 10 Aug 2023 00:54:26 +0200 Subject: [PATCH 5/9] Add prompt-toolkit to requirements --- software/script/requirements.txt | 1 + 1 file changed, 1 insertion(+) diff --git a/software/script/requirements.txt b/software/script/requirements.txt index 3419e33..981771a 100644 --- a/software/script/requirements.txt +++ b/software/script/requirements.txt @@ -1,2 +1,3 @@ pyserial==3.5 colorama==0.4.6 +prompt-toolkit==3.0.39 From cdf4669d34273904ad9726c250716f2f66d2199b Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Thu, 10 Aug 2023 01:18:27 +0200 Subject: [PATCH 6/9] Fix KeyboardInterrupt --- software/script/chameleon_cli_main.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/software/script/chameleon_cli_main.py b/software/script/chameleon_cli_main.py index 8e29b23..d898479 100755 --- a/software/script/chameleon_cli_main.py +++ b/software/script/chameleon_cli_main.py @@ -103,7 +103,7 @@ class ChameleonCLI: except KeyboardInterrupt: closing = True - if cmd_str in ["exit", "quit", "q", "e"] or closing: + if closing or cmd_str in ["exit", "quit", "q", "e"]: print("Bye, thank you. ^.^ ") self.device_com.close() sys.exit(996) From 0cfdff809fa087788a9dc40709056a8d74af554d Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Thu, 10 Aug 2023 01:20:22 +0200 Subject: [PATCH 7/9] Add command autocompletion --- software/script/chameleon_cli_main.py | 3 +- software/script/chameleon_utils.py | 81 +++++++++++++++++++++++++++ 2 files changed, 83 insertions(+), 1 deletion(-) diff --git a/software/script/chameleon_cli_main.py b/software/script/chameleon_cli_main.py index d898479..5212ac2 100755 --- a/software/script/chameleon_cli_main.py +++ b/software/script/chameleon_cli_main.py @@ -46,7 +46,8 @@ class ChameleonCLI: """ def __init__(self): - self.session = prompt_toolkit.PromptSession() + self.completer = chameleon_utils.CustomNestedCompleter.from_nested_dict(chameleon_cli_unit.root_commands) + self.session = prompt_toolkit.PromptSession(completer=self.completer) # new a device communication instance(only communication) self.device_com = chameleon_com.ChameleonCom() diff --git a/software/script/chameleon_utils.py b/software/script/chameleon_utils.py index 99469b6..20dc0b1 100644 --- a/software/script/chameleon_utils.py +++ b/software/script/chameleon_utils.py @@ -1,5 +1,7 @@ from functools import wraps from typing import Union +from prompt_toolkit.completion import Completer, NestedCompleter, WordCompleter +from prompt_toolkit.document import Document import chameleon_status @@ -84,3 +86,82 @@ class CLITree: CLITree(name=name, fullname=f'{self.fullname} {name}', helptext=helptext, cls=cls)) return cls return decorator + + +class CustomNestedCompleter(NestedCompleter): + """ + Copy of the NestedCompleter class that accepts a CLITree object and + supports meta_dict for descriptions + """ + + def __init__( + self, options, ignore_case: bool = True, meta_dict: dict = {} + ) -> None: + self.options = options + self.ignore_case = ignore_case + self.meta_dict = meta_dict + + def __repr__(self) -> str: + return f"CustomNestedCompleter({self.options!r}, ignore_case={self.ignore_case!r})" + + @classmethod + def from_nested_dict(cls, data): + options = {} + meta_dict = {} + for key, value in data.items(): + if isinstance(value, Completer): + options[key] = value + elif isinstance(value, dict): + options[key] = cls.from_nested_dict(value) + elif isinstance(value, set): + options[key] = cls.from_nested_dict({item: None for item in value}) + elif isinstance(value, CLITree): + options[key] = cls.from_clitree(value) + meta_dict[key] = value.helptext + else: + assert value is None + options[key] = None + + return cls(options, meta_dict=meta_dict) + + @classmethod + def from_clitree(cls, node): + options = {} + meta_dict = {} + + for child_node in node.children: + options[child_node.name] = cls.from_clitree(child_node) + meta_dict[child_node.name] = child_node.helptext + + return cls(options, meta_dict=meta_dict) + + def get_completions(self, document, complete_event): + # Split document. + text = document.text_before_cursor.lstrip() + stripped_len = len(document.text_before_cursor) - len(text) + + # If there is a space, check for the first term, and use a + # subcompleter. + if " " in text: + first_term = text.split()[0] + completer = self.options.get(first_term) + + # If we have a sub completer, use this for the completions. + if completer is not None: + remaining_text = text[len(first_term) :].lstrip() + move_cursor = len(text) - len(remaining_text) + stripped_len + + new_document = Document( + remaining_text, + cursor_position=document.cursor_position - move_cursor, + ) + + yield from completer.get_completions(new_document, complete_event) + + # No space in the input: behave exactly like `WordCompleter`. + else: + completer = WordCompleter( + list(self.options.keys()), ignore_case=self.ignore_case, meta_dict=self.meta_dict + ) + yield from completer.get_completions(document, complete_event) + From f6e22328bb34d9a487d8608fee955f154e3d0384 Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Sat, 19 Aug 2023 13:44:15 +0200 Subject: [PATCH 8/9] Add working argument completion --- software/script/chameleon_cli_unit.py | 23 ----- software/script/chameleon_utils.py | 116 ++++++++++++++++++++++++-- 2 files changed, 111 insertions(+), 28 deletions(-) diff --git a/software/script/chameleon_cli_unit.py b/software/script/chameleon_cli_unit.py index 110b5fd..7d157a1 100644 --- a/software/script/chameleon_cli_unit.py +++ b/software/script/chameleon_cli_unit.py @@ -14,29 +14,6 @@ import chameleon_cstruct import chameleon_status from chameleon_utils import * -description_public = "Please enter correct parameters" - - - -class ArgumentParserNoExit(argparse.ArgumentParser): - """ - If arg ArgumentParser parse error, we can't exit process, - we must raise exception to stop parse - """ - - def __init__(self, **args): - super().__init__(*args) - self.add_help = False - self.description = description_public - - def exit(self, status: int = ..., message: str or None = ...): - if message: - raise ParserExitIntercept(message) - - def error(self, message: str): - args = {'prog': self.prog, 'message': message} - raise ArgsParserError('%(prog)s: error: %(message)s\n' % args) - class BaseCLIUnit: diff --git a/software/script/chameleon_utils.py b/software/script/chameleon_utils.py index 20dc0b1..777bc29 100644 --- a/software/script/chameleon_utils.py +++ b/software/script/chameleon_utils.py @@ -1,6 +1,8 @@ +import argparse from functools import wraps -from typing import Union +from typing import Iterable, Union from prompt_toolkit.completion import Completer, NestedCompleter, WordCompleter +from prompt_toolkit.completion.base import CompleteEvent, Completion from prompt_toolkit.document import Document import chameleon_status @@ -19,6 +21,27 @@ class UnexpectedResponseError(Exception): Unexpected response exception """ + +class ArgumentParserNoExit(argparse.ArgumentParser): + """ + If arg ArgumentParser parse error, we can't exit process, + we must raise exception to stop parse + """ + + def __init__(self, **args): + super().__init__(*args) + self.add_help = False + self.description = "Please enter correct parameters" + + def exit(self, status: int = ..., message: str or None = ...): + if message: + raise ParserExitIntercept(message) + + def error(self, message: str): + args = {'prog': self.prog, 'message': message} + raise ArgsParserError('%(prog)s: error: %(message)s\n' % args) + + def expect_response(accepted_responses: Union[int, list[int]]): """ Decorator for wrapping a Chameleon CMD function to check its response @@ -116,8 +139,13 @@ class CustomNestedCompleter(NestedCompleter): elif isinstance(value, set): options[key] = cls.from_nested_dict({item: None for item in value}) elif isinstance(value, CLITree): - options[key] = cls.from_clitree(value) - meta_dict[key] = value.helptext + if value.cls: + # CLITree is a standalone command + options[key] = ArgparseCompleter(value.cls().args_parser()) + else: + # CLITree is a command group + options[key] = cls.from_clitree(value) + meta_dict[key] = value.helptext else: assert value is None options[key] = None @@ -130,8 +158,13 @@ class CustomNestedCompleter(NestedCompleter): meta_dict = {} for child_node in node.children: - options[child_node.name] = cls.from_clitree(child_node) - meta_dict[child_node.name] = child_node.helptext + if child_node.cls and child_node.cls().args_parser(): + # CLITree is a standalone command with arguments + options[child_node.name] = ArgparseCompleter(child_node.cls().args_parser()) + else: + # CLITree is a command group + options[child_node.name] = cls.from_clitree(child_node) + meta_dict[child_node.name] = child_node.helptext return cls(options, meta_dict=meta_dict) @@ -165,3 +198,76 @@ class CustomNestedCompleter(NestedCompleter): ) yield from completer.get_completions(document, complete_event) + +class ArgparseCompleter(Completer): + """ + Completer instance for autocompletion of ArgumentParser arguments + + :param parser: ArgumentParser instance + """ + + def __init__(self, parser) -> None: + self.parser: ArgumentParserNoExit = parser + + def check_tokens(self, parsed, unparsed): + suggestions = {} + def check_arg(tokens): + return tokens and tokens[0].startswith('-') + + if not parsed and not unparsed: + # No tokens detected, just show all flags + for action in self.parser._actions: + for opt in action.option_strings: + suggestions[opt] = action.help + return [], [], suggestions + + token = unparsed.pop(0) + + for action in self.parser._actions: + if any(opt == token for opt in action.option_strings): + # Argument fully matches the token + parsed.append(token) + + if action.choices: + # Autocomplete with choices + if unparsed: + # Autocomplete values + value = unparsed.pop(0) + for choice in action.choices: + if str(choice).startswith(value): + suggestions[str(choice)] = None + + parsed.append(value) + + if check_arg(unparsed): + parsed, unparsed, suggestions = self.check_tokens(parsed, unparsed) + + else: + # Show all possible values + for choice in action.choices: + suggestions[str(choice)] = None + + break + else: + # No choices, process further arguments + if check_arg(unparsed): + parsed, unparsed, suggestions = self.check_tokens(parsed, unparsed) + break + elif any(opt.startswith(token) for opt in action.option_strings): + for opt in action.option_strings: + if opt.startswith(token): + suggestions[opt] = action.help + + if suggestions: + unparsed.insert(0, token) + + return parsed, unparsed, suggestions + + def get_completions(self, document, complete_event): + text = document.text_before_cursor + word_before_cursor = document.get_word_before_cursor() + + _, _, suggestions = self.check_tokens(list(), text.split()) + + for key, suggestion in suggestions.items(): + yield Completion(key, -len(word_before_cursor), display=key, display_meta=suggestion) From a64be123db59b1c08a27bc83e0803940797787fb Mon Sep 17 00:00:00 2001 From: Szymon Borecki Date: Sat, 19 Aug 2023 13:46:32 +0200 Subject: [PATCH 9/9] Remove the unnecessary readline import and add a shebang --- software/script/chameleon_cli_main.py | 3 +-- 1 file changed, 1 insertion(+), 2 deletions(-) diff --git a/software/script/chameleon_cli_main.py b/software/script/chameleon_cli_main.py index 5212ac2..5106423 100755 --- a/software/script/chameleon_cli_main.py +++ b/software/script/chameleon_cli_main.py @@ -1,3 +1,4 @@ +#!/usr/bin/env python3 import argparse import os import platform @@ -13,8 +14,6 @@ import os import prompt_toolkit from prompt_toolkit.formatted_text import ANSI -if os.name == 'posix': - import readline ULTRA = r""" ╦ ╦╦ ╔╦╗╦═╗╔═╗