Properly handle snapshots if camera added to bridge

pyhap/hap_handler.py

@@ -16,7 +16,7 @@
 
 import pyhap.tlv as tlv
 from pyhap.util import long_to_bytes
-
+from pyhap.const import CATEGORY_BRIDGE
 from .hap_crypto import hap_hkdf, pad_tls_nonce
 
 SNAPSHOT_TIMEOUT = 10
@@ -703,24 +703,29 @@ def _send_tlv_pairing_response(self, data):
 
     def handle_resource(self):
         """Get a snapshot from the camera."""
-        image_size = json.loads(self.request_body.decode("utf-8"))
+        data = json.loads(self.request_body.decode("utf-8"))
+
+        if self.accessory_handler.accessory.category == CATEGORY_BRIDGE:
+            accessory = self.accessory_handler.accessory.accessories.get(data["aid"])
+            if not accessory:
+                raise ValueError(
+                    "Accessory with aid == {} not found".format(data["aid"])
+                )
+        else:
+            accessory = self.accessory_handler.accessory
+
         loop = asyncio.get_event_loop()
-        if hasattr(self.accessory_handler.accessory, "async_get_snapshot"):
-            coro = self.accessory_handler.accessory.async_get_snapshot(image_size)
-        elif hasattr(self.accessory_handler.accessory, "get_snapshot"):
-            coro = asyncio.wait_for(
-                loop.run_in_executor(
-                    None, self.accessory_handler.accessory.get_snapshot, image_size
-                ),
-                SNAPSHOT_TIMEOUT,
-            )
+        if hasattr(accessory, "async_get_snapshot"):
+            coro = accessory.async_get_snapshot(data)
+        elif hasattr(accessory, "get_snapshot"):
+            coro = loop.run_in_executor(None, accessory.get_snapshot, data)
         else:
             raise ValueError(
                 "Got a request for snapshot, but the Accessory "
                 'does not define a "get_snapshot" or "async_get_snapshot" method'
             )
 
-        task = asyncio.ensure_future(coro)
+        task = asyncio.ensure_future(asyncio.wait_for(coro, SNAPSHOT_TIMEOUT))
         self.send_response(200)
         self.send_header("Content-Type", "image/jpeg")
         self.response.task = task

tests/test_hap_handler.py

@@ -6,7 +6,7 @@
 import pytest
 
 from pyhap import hap_handler
-from pyhap.accessory import Accessory
+from pyhap.accessory import Accessory, Bridge
 import pyhap.tlv as tlv
 
 CLIENT_UUID = UUID("7d0d1ee9-46fe-4a56-a115-69df3f6860c1")
@@ -291,3 +291,18 @@ def test_handle_set_handle_set_characteristics_encrypted(driver):
 
     assert response.status_code == 204
     assert response.body == b""
+
+
+def test_handle_snapshot_encrypted_non_existant_accessory(driver):
+    """Verify an encrypted snapshot with non-existant accessory."""
+    bridge = Bridge(driver, "Test Bridge")
+    driver.add_accessory(bridge)
+
+    handler = hap_handler.HAPServerHandler(driver, "peername")
+    handler.is_encrypted = True
+
+    response = hap_handler.HAPResponse()
+    handler.response = response
+    handler.request_body = b'{"image-height":360,"resource-type":"image","image-width":640,"aid":1411620844}'
+    with pytest.raises(ValueError):
+        handler.handle_resource()

tests/test_hap_protocol.py

@@ -6,7 +6,7 @@
 import pytest
 
 from pyhap import hap_protocol
-from pyhap.accessory import Accessory
+from pyhap.accessory import Accessory, Bridge
 from pyhap.hap_handler import HAPResponse
 
 
@@ -474,3 +474,65 @@ async def _async_get_snapshot(*_):
     assert b"-70402" in writer.call_args_list[0][0][0]
 
     hap_proto.close()
+
+
+@pytest.mark.asyncio
+async def test_camera_snapshot_times_out(driver):
+    """Test camera snapshot times out."""
+    loop = MagicMock()
+    transport = MagicMock()
+    connections = {}
+
+    def _get_snapshot(*_):
+        raise asyncio.TimeoutError("timeout")
+
+    acc = Accessory(driver, "TestAcc")
+    acc.get_snapshot = _get_snapshot
+    driver.add_accessory(acc)
+
+    hap_proto = hap_protocol.HAPServerProtocol(loop, connections, driver)
+    hap_proto.connection_made(transport)
+
+    hap_proto.hap_crypto = MockHAPCrypto()
+    hap_proto.handler.is_encrypted = True
+
+    with patch.object(hap_proto.transport, "write") as writer:
+        hap_proto.data_received(
+            b'POST /resource HTTP/1.1\r\nHost: HASS\\032Bridge\\032BROZ\\0323BF435._hap._tcp.local\r\nContent-Length: 79\r\nContent-Type: application/hap+json\r\n\r\n{"image-height":360,"resource-type":"image","image-width":640,"aid":1411620844}'  # pylint: disable=line-too-long
+        )
+        try:
+            await hap_proto.response.task
+        except Exception:  # pylint: disable=broad-except
+            pass
+        await asyncio.sleep(0)
+
+    assert b"-70402" in writer.call_args_list[0][0][0]
+
+    hap_proto.close()
+
+
+@pytest.mark.asyncio
+async def test_camera_snapshot_missing_accessory(driver):
+    """Test camera snapshot that throws an exception."""
+    loop = MagicMock()
+    transport = MagicMock()
+    connections = {}
+
+    bridge = Bridge(driver, "Test Bridge")
+    driver.add_accessory(bridge)
+
+    hap_proto = hap_protocol.HAPServerProtocol(loop, connections, driver)
+    hap_proto.connection_made(transport)
+
+    hap_proto.hap_crypto = MockHAPCrypto()
+    hap_proto.handler.is_encrypted = True
+
+    with patch.object(hap_proto.transport, "write") as writer:
+        hap_proto.data_received(
+            b'POST /resource HTTP/1.1\r\nHost: HASS\\032Bridge\\032BROZ\\0323BF435._hap._tcp.local\r\nContent-Length: 79\r\nContent-Type: application/hap+json\r\n\r\n{"image-height":360,"resource-type":"image","image-width":640,"aid":1411620844}'  # pylint: disable=line-too-long
+        )
+        await asyncio.sleep(0)
+
+    assert hap_proto.response is None
+    assert b"-70402" in writer.call_args_list[0][0][0]
+    hap_proto.close()