diff --git a/pm3py/hw.py b/pm3py/hw.py index 5d6952a..519554b 100644 --- a/pm3py/hw.py +++ b/pm3py/hw.py @@ -40,16 +40,81 @@ class HWCommands: resp = await self._t.send_ng(Cmd.STATUS, payload, timeout=2.0 + speed_test_timeout / 1000) return {"status": resp.status} - async def led(self, led: int, action: str = "on", brightness: int = 100, - speed: int = 500, count: int = 5) -> dict: - """Control LEDs. led: bitmask (1=A,2=B,4=C,8=D). action: on/off/toggle/pwm/pulse/fade/blink.""" - actions = {"off": 0, "on": 1, "toggle": 2, "pwm": 3, "pulse": 4, "fade": 5, "blink": 6} - if action not in actions: - raise ValueError(f"Unknown action: {action}. Use one of: {list(actions.keys())}") - payload = struct.pack(" dict: + """Control LEDs on the Proxmark3. + + Args: + led: Color name(s), letter(s), or bitmask int. + Names: "green", "red", "orange", "blue", "all" + Letters: "a", "b", "c", "d" + Comma-separated: "green,red" + Bitmask: 1=A, 2=B, 4=C, 8=D + + Action kwargs (pick one): + on: Turn LED on + off: Turn LED off + toggle: Toggle LED state + brightness: Set PWM brightness 0-100 (PWM LEDs only) + pulse: Pulse effect (PWM LEDs only) + fade: Fade effect (PWM LEDs only) + blink: Blink effect (all LEDs) + + Effect params: + speed: Cycle time in ms (default 500) + count: Repeat count, 0=infinite (default 5) + """ + # Resolve LED mask + if isinstance(led, int): + led_mask = led & 0x0F + else: + led_mask = self._parse_led_str(led) + + # Resolve action + if brightness >= 0: + action_code = 3 + brightness_val = min(max(brightness, 0), 100) + elif off: + action_code = 0 + brightness_val = 0 + elif toggle: + action_code = 2 + brightness_val = 0 + elif pulse: + action_code = 4 + brightness_val = 0 + elif fade: + action_code = 5 + brightness_val = 0 + elif blink: + action_code = 6 + brightness_val = 0 + elif on or True: # default to "on" + action_code = 1 + brightness_val = 0 + + payload = struct.pack(" int: + """Parse LED string into bitmask. Accepts names, letters, or 'all'.""" + LED_MAP = { + "a": 0x01, "b": 0x02, "c": 0x04, "d": 0x08, + "green": 0x02, "orange": 0x01, "red": 0x04, "blue": 0x08, + "all": 0x0F, + } + mask = 0 + for token in led.lower().replace(" ", "").split(","): + if token in LED_MAP: + mask |= LED_MAP[token] + else: + raise ValueError(f"Unknown LED: {token!r}. Use: {', '.join(LED_MAP.keys())}") + return mask + async def dbg(self, level: int | None = None) -> dict: """Get or set debug level. None=get, 0-4=set.""" if level is None: diff --git a/tests/test_hw.py b/tests/test_hw.py index 3220dfc..e14f571 100644 --- a/tests/test_hw.py +++ b/tests/test_hw.py @@ -39,9 +39,36 @@ def test_status(): result = asyncio.get_event_loop().run_until_complete(hw.status()) assert result["status"] == 0 -def test_led(): +def test_led_on_by_name(): transport = make_mock_transport() hw = HWCommands(transport) transport.send_ng.return_value = make_ng_response(Cmd.LED_CONTROL, 0, b"") - result = asyncio.get_event_loop().run_until_complete(hw.led(1, "on")) + result = asyncio.get_event_loop().run_until_complete(hw.led("green", on=True)) transport.send_ng.assert_called_once() + +def test_led_pulse(): + transport = make_mock_transport() + hw = HWCommands(transport) + transport.send_ng.return_value = make_ng_response(Cmd.LED_CONTROL, 0, b"") + result = asyncio.get_event_loop().run_until_complete( + hw.led("red", pulse=True, speed=300, count=3)) + transport.send_ng.assert_called_once() + # Verify the payload + call_args = transport.send_ng.call_args + payload = call_args[0][1] # second positional arg + led_mask, action, brightness, speed, count = struct.unpack("