diff --git a/speculos/api/apdu.py b/speculos/api/apdu.py index a3f09447..f0a501c3 100644 --- a/speculos/api/apdu.py +++ b/speculos/api/apdu.py @@ -18,17 +18,23 @@ def __init__(self, seph: SeProxyHal): self._seph.apdu_callbacks.append(self.seph_apdu_callback) self.response: Optional[bytes] - def exchange(self, data: bytes) -> Generator[bytes, None, None]: + def exchange(self, data: bytes, tick_timeout: int = 0) -> Generator[bytes, None, None]: # force headers to be sent yield b"" + tick_count_before_exchange = self._seph.socket_helper.ticks_count + with self.endpoint_lock: # Lock for a command/response for one client with self.response_condition: self.response = None self._seph.to_app(data) with self.response_condition: while self.response is None: - self.response_condition.wait() + self.response_condition.wait(0.1) + exchange_tick_count = self._seph.socket_helper.ticks_count - tick_count_before_exchange + + if tick_timeout != 0 and exchange_tick_count > tick_timeout: + yield yield json.dumps({"data": self.response.hex()}).encode() def seph_apdu_callback(self, data: bytes) -> None: @@ -55,5 +61,6 @@ def post(self): except jsonschema.exceptions.ValidationError as e: return {"error": f"{e}"}, 400 - data = bytes.fromhex(args.get("data")) - return Response(stream_with_context(self._bridge.exchange(data)), content_type="application/json") + tick_timeout = args.get("data")["tick_timeout"] + data = bytes.fromhex(args.get("data")["apdu"]) + return Response(stream_with_context(self._bridge.exchange(data, tick_timeout)), content_type="application/json") diff --git a/speculos/api/resources/apdu.schema b/speculos/api/resources/apdu.schema index 220fd196..9e1b94a4 100644 --- a/speculos/api/resources/apdu.schema +++ b/speculos/api/resources/apdu.schema @@ -3,7 +3,10 @@ "type": "object", "properties": { - "data": { "type": "string", "pattern": "^([a-fA-F0-9]{2})+$" } + "data": { + "apdu": { "type": "string", "pattern": "^([a-fA-F0-9]{2})+$" }, + "tick_timeout": { "type": "int"} + } }, "required": [ "data" ], "additionalProperties": false diff --git a/speculos/client.py b/speculos/client.py index 117f84fa..68a56c0d 100644 --- a/speculos/client.py +++ b/speculos/client.py @@ -141,10 +141,14 @@ def get_screenshot(self) -> bytes: check_status_code(response, "/screenshot") return response.content - def _apdu_exchange(self, data: bytes) -> bytes: - with self.session.post(f"{self.api_url}/apdu", json={"data": data.hex()}) as response: - apdu_response = ApduResponse(response) - return apdu_response.receive() + def _apdu_exchange(self, data: bytes, tick_timeout: int = 0) -> bytes: + try: + data_payload = {"apdu": data.hex(), "tick_timeout": tick_timeout} + with self.session.post(f"{self.api_url}/apdu", json={"data": data_payload}) as response: + apdu_response = ApduResponse(response) + return apdu_response.receive() + except: + raise TimeoutError() def _apdu_exchange_nowait(self, data: bytes) -> requests.Response: return self.session.post(f"{self.api_url}/apdu", json={"data": data.hex()}, stream=True) @@ -239,9 +243,9 @@ def __exit__( ) -> None: self.stop() - def apdu_exchange(self, cla: int, ins: int, data: bytes = b"", p1: int = 0, p2: int = 0) -> bytes: + def apdu_exchange(self, cla: int, ins: int, data: bytes = b"", p1: int = 0, p2: int = 0, tick_timeout:int = 0) -> bytes: apdu = bytes([cla, ins, p1, p2, len(data)]) + data - return Api._apdu_exchange(self, apdu) + return Api._apdu_exchange(self, apdu, tick_timeout) @contextmanager def apdu_exchange_nowait(