From cc7dfa89de29c378a82cdc4dda3d4a6b7054e245 Mon Sep 17 00:00:00 2001
From: michael <operations@dangerousthings.com>
Date: Mon, 16 Mar 2026 00:20:47 -0700
Subject: [PATCH] feat(pm3py): validate PWM LED capability using device
 capabilities, raise on non-PWM LEDs

---
 pm3py/hw.py      | 56 ++++++++++++++++++++++++++++++++++++++++++++----
 tests/test_hw.py | 38 ++++++++++++++++++++++++++------
 2 files changed, 83 insertions(+), 11 deletions(-)

diff --git a/pm3py/hw.py b/pm3py/hw.py
index 2ca494a..52b41da 100644
--- a/pm3py/hw.py
+++ b/pm3py/hw.py
@@ -8,8 +8,15 @@ from .transport import PM3Transport, PM3Response, PM3Error, ProgressCallback
 class HWCommands:
     """Proxmark3 hardware commands."""
 
+    # PWM-capable LEDs by platform (bitmask)
+    # PM3 Easy: A (green) and B (red) support PWM
+    # RDV4: A (orange) and D (red2) support PWM
+    _PWM_LEDS_EASY = 0x03  # A=0x01 | B=0x02
+    _PWM_LEDS_RDV4 = 0x09  # A=0x01 | D=0x08
+
     def __init__(self, transport: PM3Transport):
         self._t = transport
+        self._is_rdv4: bool | None = None  # cached after first capabilities() call
 
     async def ping(self, length: int = 32) -> dict:
         """Send ping and verify echo. Returns {success, length, data}."""
@@ -97,6 +104,10 @@ class HWCommands:
 
         ACTION_NAMES = {0: "off", 1: "on", 2: "toggle", 3: "pwm", 4: "pulse", 5: "fade", 6: "blink"}
 
+        # Validate PWM capability for PWM actions (brightness, pulse, fade)
+        if action_code in (3, 4, 5):
+            await self._check_pwm_capable(led_mask, ACTION_NAMES[action_code])
+
         payload = struct.pack("<BBBHH", led_mask, action_code, brightness_val, speed, count)
         await self._t.send_ng(Cmd.LED_CONTROL, payload)
         return {
@@ -123,6 +134,27 @@ class HWCommands:
                 raise ValueError(f"Unknown LED: {token!r}. Use: {', '.join(LED_MAP.keys())}")
         return mask
 
+    async def _check_pwm_capable(self, led_mask: int, action: str) -> None:
+        """Raise if any LEDs in mask don't support PWM."""
+        # Fetch capabilities if not cached
+        if self._is_rdv4 is None:
+            try:
+                await self.capabilities()
+            except PM3Error:
+                return  # can't validate, let it through
+
+        pwm_mask = self._PWM_LEDS_RDV4 if self._is_rdv4 else self._PWM_LEDS_EASY
+        non_pwm = led_mask & ~pwm_mask
+        if non_pwm:
+            platform = "RDV4" if self._is_rdv4 else "PM3 Easy"
+            LED_NAMES = {0x01: "A", 0x02: "B", 0x04: "C", 0x08: "D"}
+            bad = [LED_NAMES[b] for b in (0x01, 0x02, 0x04, 0x08) if non_pwm & b]
+            good = [LED_NAMES[b] for b in (0x01, 0x02, 0x04, 0x08) if pwm_mask & b]
+            raise PM3Error(
+                f"LED(s) {','.join(bad)} do not support {action} on {platform}. "
+                f"PWM-capable LEDs: {','.join(good)}"
+            )
+
     async def dbg(self, level: int | None = None) -> dict:
         """Get or set debug level. None=get, 0-4=set."""
         if level is None:
@@ -162,15 +194,31 @@ class HWCommands:
         version = resp.data[0]
         baudrate = struct.unpack_from("<I", resp.data, 1)[0]
         bigbuf_size = struct.unpack_from("<I", resp.data, 5)[0]
-        # Bit flags start at byte 9
-        flags = 0
+
+        # Parse bit flags (bytes 9+)
+        is_rdv4 = False
+        compiled_with_lf = False
+        compiled_with_iso14443a = False
+        compiled_with_iso15693 = False
+        if len(resp.data) > 12:
+            is_rdv4 = bool(resp.data[12] & 0x02)  # byte 12, bit 1
         if len(resp.data) > 9:
-            flags = struct.unpack_from("<I", resp.data, 9)[0] if len(resp.data) >= 13 else resp.data[9]
+            compiled_with_lf = bool(resp.data[9] & 0x80)  # byte 9, bit 7
+        if len(resp.data) > 10:
+            compiled_with_iso14443a = bool(resp.data[10] & 0x40)  # byte 10, bit 6
+        if len(resp.data) > 11:
+            compiled_with_iso15693 = bool(resp.data[11] & 0x01)  # byte 11, bit 0
+
+        self._is_rdv4 = is_rdv4
+
         return {
             "version": version,
             "baudrate": baudrate,
             "bigbuf_size": bigbuf_size,
-            "raw_flags": flags,
+            "is_rdv4": is_rdv4,
+            "compiled_with_lf": compiled_with_lf,
+            "compiled_with_iso14443a": compiled_with_iso14443a,
+            "compiled_with_iso15693": compiled_with_iso15693,
         }
 
     async def tune(self, on_progress: ProgressCallback = None) -> dict:
diff --git a/tests/test_hw.py b/tests/test_hw.py
index e14f571..9a3e7b2 100644
--- a/tests/test_hw.py
+++ b/tests/test_hw.py
@@ -12,6 +12,15 @@ def make_mock_transport():
     transport = AsyncMock()
     return transport
 
+def _make_caps_response(is_rdv4=False):
+    """Build a capabilities response with is_rdv4 flag."""
+    # version(1) + baudrate(4) + bigbuf_size(4) + 4 bytes of flags
+    flags = bytearray(4)
+    flags[0] = 0x80  # compiled_with_lf
+    if is_rdv4:
+        flags[3] = 0x02  # is_rdv4 at byte 12 bit 1
+    return make_ng_response(Cmd.CAPABILITIES, 0, struct.pack("<BII", 7, 115200, 40000) + bytes(flags))
+
 def test_ping():
     transport = make_mock_transport()
     hw = HWCommands(transport)
@@ -49,15 +58,17 @@ def test_led_on_by_name():
 def test_led_pulse():
     transport = make_mock_transport()
     hw = HWCommands(transport)
+    # pulse on "red" (C=0x04) — not PWM-capable on Easy, so pre-set as RDV4
+    hw._is_rdv4 = True  # skip capabilities fetch; RDV4 has PWM on A,D only... but red=C
+    # Actually red=C is not PWM on either platform. Use "orange" (A) which is PWM on both.
     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
+        hw.led("orange", pulse=True, speed=300, count=3))
+    # Verify the payload (last call to send_ng)
     call_args = transport.send_ng.call_args
-    payload = call_args[0][1]  # second positional arg
+    payload = call_args[0][1]
     led_mask, action, brightness, speed, count = struct.unpack("<BBBHH", payload)
-    assert led_mask == 0x04  # red = C
+    assert led_mask == 0x01  # orange = A
     assert action == 4       # pulse
     assert speed == 300
     assert count == 3
@@ -65,10 +76,23 @@ def test_led_pulse():
 def test_led_brightness():
     transport = make_mock_transport()
     hw = HWCommands(transport)
+    # Pre-cache as Easy — PWM capable on A,B only
+    hw._is_rdv4 = False
     transport.send_ng.return_value = make_ng_response(Cmd.LED_CONTROL, 0, b"")
-    result = asyncio.get_event_loop().run_until_complete(hw.led("all", brightness=50))
+    result = asyncio.get_event_loop().run_until_complete(hw.led("green", brightness=50))
     payload = transport.send_ng.call_args[0][1]
     led_mask, action, bright, _, _ = struct.unpack("<BBBHH", payload)
-    assert led_mask == 0x0F  # all
+    assert led_mask == 0x02  # green = B
     assert action == 3       # pwm
     assert bright == 50
+
+def test_led_pwm_non_capable_raises():
+    """PWM on a non-PWM LED should raise PM3Error."""
+    from pm3py.transport import PM3Error
+    import pytest
+    transport = make_mock_transport()
+    hw = HWCommands(transport)
+    hw._is_rdv4 = False  # Easy: only A,B support PWM
+    with pytest.raises(PM3Error, match="do not support"):
+        asyncio.get_event_loop().run_until_complete(
+            hw.led("blue", brightness=50))  # blue=D, not PWM on Easy