diff --git a/pm3py/core/lf.py b/pm3py/core/lf.py index f1c171b..6fd25c6 100644 --- a/pm3py/core/lf.py +++ b/pm3py/core/lf.py @@ -180,34 +180,18 @@ class LFCommands: on_progress: ProgressCallback = None) -> bytes: """Download raw ADC samples from the device's BigBuf. - Must call read() or search() first to capture samples. - Returns raw bytes (1 byte per sample, unsigned 0-255). + Must call read() or search() first to capture samples. Returns raw bytes (1 byte per + sample, unsigned 0-255). The device streams the data as OLD-format CMD_DOWNLOADED_BIGBUF + frames then a final ACK — handled by the transport's bulk-download reader. """ - buf = bytearray() - remaining = count - offset = start - - while remaining > 0: - chunk = min(remaining, PM3_CMD_DATA_SIZE) - resp = await self._t.send_mix( - Cmd.DOWNLOAD_BIGBUF, offset, chunk, 0, timeout=5.0, - ) - # Device responds with CMD_DOWNLOADED_BIGBUF containing the data - dl_resp = await self._t.read_response(timeout=5.0) - if dl_resp.data: - buf.extend(dl_resp.data) - - if on_progress: - ret = on_progress({ - "type": "download", "downloaded": len(buf), "total": count, - }) - if asyncio.iscoroutine(ret): - await ret - - offset += chunk - remaining -= chunk - - return bytes(buf[:count]) + data = await self._t.download_bulk( + Cmd.DOWNLOAD_BIGBUF, start, count, Cmd.DOWNLOADED_BIGBUF, Cmd.ACK, timeout=5.0, + ) + if on_progress: + ret = on_progress({"type": "download", "downloaded": len(data), "total": count}) + if asyncio.iscoroutine(ret): + await ret + return data async def search(self, on_progress: ProgressCallback = None) -> "LFSearchResult": """Search for LF tags. Captures samples and returns a result with callable next steps. diff --git a/pm3py/core/transport.py b/pm3py/core/transport.py index 8386b56..8aa46e7 100644 --- a/pm3py/core/transport.py +++ b/pm3py/core/transport.py @@ -278,6 +278,38 @@ class PM3Transport: frame = encode_old_frame(cmd, arg0, arg1, arg2, data) await self.send_frame(frame) + async def download_bulk(self, cmd: int, start: int, count: int, data_cmd: int, + ack_cmd: int, timeout: float = 5.0, on_bytes=None) -> bytes: + """Run a bulk BigBuf/EML/SPIFFS download and return the concatenated bytes. + + The firmware answers one download request (a MIX frame) by streaming ``count`` bytes as + a run of *OLD*-format ``data_cmd`` frames (no magic — that's why the plain NG reader + can't see them and the naive download hangs), then a final ``ack_cmd`` frame. We read + with the NG/OLD-agnostic reader until the ACK, or until we've collected ``count`` bytes + (draining the trailing ACK best-effort so it can't poison the next command).""" + async with self._lock: + await self.send_frame(encode_mix_frame(cmd, start, count, 0)) + buf = bytearray() + while True: + try: + resp = await self._read_any_response(timeout=timeout) + except PM3Error: + break # timeout (e.g. ACK lost) — return what we have + if resp.cmd == ack_cmd: + break + if resp.cmd == data_cmd and resp.data: + length = resp.oldarg[1] if resp.oldarg and len(resp.oldarg) > 1 else len(resp.data) + buf.extend(resp.data[:max(0, length)]) + if on_bytes is not None: + on_bytes(len(buf)) + if count and len(buf) >= count: + try: # got all the data; consume the trailing ACK + await self._read_any_response(timeout=1.0) + except PM3Error: + pass + break + return bytes(buf[:count]) if count else bytes(buf) + async def _read_any_response(self, timeout: float = 5.0) -> PM3Response: """Read one frame that may be NG or OLD and return a PM3Response.""" if not self._reader: diff --git a/tests/test_transport.py b/tests/test_transport.py index 9d3a201..6ca476c 100644 --- a/tests/test_transport.py +++ b/tests/test_transport.py @@ -79,3 +79,57 @@ def test_decode_response_frame_error_status(): resp = decode_response_frame(frame) assert resp["status"] == -4 # ETIMEOUT assert resp["data"] == b"" + + +# --- bulk download (BigBuf): one request -> OLD data frames -> ACK ----------------------------- + +import asyncio +from unittest.mock import AsyncMock +from pm3py.core.transport import PM3Transport, PM3Response +from pm3py.core.protocol import PM3_CMD_DATA_SIZE + + +def _run(coro): + return asyncio.get_event_loop().run_until_complete(coro) + + +def _bigbuf_frame(offset, data): + # OLD-format CMD_DOWNLOADED_BIGBUF: oldarg=[offset, len, tracelen], data padded to 512 + return PM3Response(cmd=int(Cmd.DOWNLOADED_BIGBUF), status=0, reason=0, ng=False, + data=data + b"\x00" * (PM3_CMD_DATA_SIZE - len(data)), + oldarg=[offset, len(data), 4242]) + + +def _ack(): + return PM3Response(cmd=int(Cmd.ACK), status=0, reason=0, ng=False, data=b"", oldarg=[1, 0, 0]) + + +def _transport_with(frames): + t = PM3Transport("/dev/null") + t.send_frame = AsyncMock() + t._read_any_response = AsyncMock(side_effect=frames) + return t + + +def test_download_bulk_count_terminated(): + payload = bytes((i * 7) % 256 for i in range(1224)) + t = _transport_with([_bigbuf_frame(0, payload[0:512]), _bigbuf_frame(512, payload[512:1024]), + _bigbuf_frame(1024, payload[1024:1224]), _ack()]) + got = _run(t.download_bulk(Cmd.DOWNLOAD_BIGBUF, 0, 1224, Cmd.DOWNLOADED_BIGBUF, Cmd.ACK)) + assert got == payload # exactly count bytes, reassembled in order + assert t.send_frame.call_count == 1 # ONE request, not one per chunk + + +def test_download_bulk_ack_terminated_early(): + # device has fewer bytes than requested -> stops at the ACK, returns what came + payload = bytes(range(200)) + t = _transport_with([_bigbuf_frame(0, payload), _ack()]) + got = _run(t.download_bulk(Cmd.DOWNLOAD_BIGBUF, 0, 99999, Cmd.DOWNLOADED_BIGBUF, Cmd.ACK)) + assert got == payload + + +def test_download_bulk_uses_oldarg_length_not_padding(): + # the OLD frame is always 512 bytes; only oldarg[1] bytes are valid (rest is padding) + t = _transport_with([_bigbuf_frame(0, b"\xAA\xBB\xCC"), _ack()]) + got = _run(t.download_bulk(Cmd.DOWNLOAD_BIGBUF, 0, 3, Cmd.DOWNLOADED_BIGBUF, Cmd.ACK)) + assert got == b"\xAA\xBB\xCC" # trimmed to oldarg length, not 512 of padding