Implement our own unifi websocket connection

main.py

@@ -10,17 +10,19 @@ import json
 import os
 import sys
 import asyncio
-import aiohttp
-import RPi.GPIO as GPIO
+
+try:
+    import RPi.GPIO as GPIO
+    IS_PI = True
+except ModuleNotFoundError:
+    IS_PI = False
+
 import time
 from signal import *
-from aiohttp import ClientSession, CookieJar
 
+import unifi
 import settings
 
-from pyunifiprotect.unifi_protect_server import UpvServer
-from pyunifiprotect.exceptions import NvrError
-
 RELAY_ON = False
 RELAY_OFF = True
 
@@ -29,7 +31,7 @@ allow_watchdog = False
 cooldown_time = time.time()
 
 def set_relay(pin, state):
-    GPIO.output(pin, state)
+    if IS_PI: GPIO.output(pin, state)
     logging.info('Set relay on pin %s to %s', pin, 'ON' if state == RELAY_ON else 'OFF')
 
 def pulse_relay(pin):
@@ -37,7 +39,7 @@ def pulse_relay(pin):
     time.sleep(0.5)  # atomic
     set_relay(pin, RELAY_OFF)
 
-def ring_bell(mac):
+def ring_bell(camera):
     global allow_watchdog, cooldown_time
 
     if not allow_watchdog and not DEBUG and not NO_WATCHDOG:
@@ -51,71 +53,34 @@ def ring_bell(mac):
     cooldown_time = time.time()
 
     try:
-        doorbell = settings.DOORBELLS[mac]
+        doorbell = settings.DOORBELLS[camera]
         pulse_relay(doorbell['gpio'])
     except KeyError:
-        logging.error('Doorbell %s not found!', mac)
-
-def subscriber(updates):
-    logging.debug('Subscription: updates=%s', updates)
-
-    for _, data in updates.items():
-        if data['event_type'] == 'ring' and data['event_ring_on']:
-            logging.info('%s: %s is ringing!', data['mac'], data['name'])
-            ring_bell(data['mac'])
+        logging.error('Doorbell %s not found!', camera)
 
 def feed_watchdog():
     with open('/dev/watchdog', 'w') as wdt:
         wdt.write('1')
 
-async def ws_listener():
-    session = ClientSession(cookie_jar=CookieJar(unsafe=True))
 
-    unifiprotect = UpvServer(
-        session,
-        settings.UFP_ADDRESS,
-        settings.UFP_PORT,
-        settings.UFP_USERNAME,
-        settings.UFP_PASSWORD,
-    )
+async def process_message(msg):
+    if msg.get('type', '') != 'ring':
+        return
 
-    await unifiprotect.update()
+    logging.info('Ring message: %s', msg)
 
-    unsub = unifiprotect.subscribe_websocket(subscriber)
-
-    logging.info('Connecting to websocket.')
-    await asyncio.sleep(2)
+    ring_bell(msg['camera'])
 
+async def main():
     while True:
         try:
-            updates = await unifiprotect.update()
-            logging.debug('')
-            logging.debug('Updates: %s', json.dumps(updates, indent=4))
-        except NvrError:
-            logging.error('Error updating connection. Reconnecting...')
-            break
-
-        active_ws = await unifiprotect.check_ws()
-        if not active_ws:
-            logging.error('Websocket unactive. Reconnecting...')
-            break
-
-        if allow_watchdog:
-            feed_watchdog()
-
-        await asyncio.sleep(1)
-
-    await session.close()
-    unsub()
-
-async def connect():
-    while True:
-        try:
-            await ws_listener()
-        except NvrError as e:
+            async for msg in unifi.connect():
+                await process_message(msg)
+        except BaseException as e:
             logging.error('Error connecting to Unifi Protect: %s. Trying again...', str(e))
             await asyncio.sleep(3)
 
+
 def disable_relays_on_exit(*args):
     logging.info('Exiting, disabling relays...')
     for _, doorbell in settings.DOORBELLS.items():
@@ -124,11 +89,11 @@ def disable_relays_on_exit(*args):
     os._exit(0)
 
 def init():
-    GPIO.setmode(GPIO.BCM)
-    GPIO.setwarnings(False)
+    if IS_PI: GPIO.setmode(GPIO.BCM)
+    if IS_PI: GPIO.setwarnings(False)
 
     for _, doorbell in settings.DOORBELLS.items():
-        GPIO.setup(doorbell['gpio'], GPIO.OUT)
+        if IS_PI: GPIO.setup(doorbell['gpio'], GPIO.OUT)
         set_relay(doorbell['gpio'], RELAY_OFF)
         #pulse_relay(doorbell['gpio'])
         time.sleep(1)
@@ -145,5 +110,5 @@ if __name__ == '__main__':
     init()
 
     loop = asyncio.get_event_loop()
-    loop.run_until_complete(connect())
+    loop.run_until_complete(main())
     loop.close()

requirements.txt

@@ -1 +1,9 @@
-pyunifiprotect==1.1.0
+aiohttp==3.8.1
+aiosignal==1.2.0
+async-timeout==4.0.2
+attrs==21.4.0
+charset-normalizer==2.0.12
+frozenlist==1.3.0
+idna==3.3
+multidict==6.0.2
+yarl==1.7.2

unifi.py

@@ -0,0 +1,41 @@
+import asyncio
+import aiohttp
+import zlib
+import struct
+import json
+
+import settings
+
+HEADER_LENGTH = 8
+
+async def connect():
+    data = dict(
+        username=settings.UFP_USERNAME,
+        password=settings.UFP_PASSWORD,
+        rememberMe=True,
+    )
+
+    async with aiohttp.ClientSession() as session:
+        async with session.post(settings.UFP_ADDRESS + '/api/auth/login', json=data, ssl=False) as resp:
+            cookie = resp.cookies['TOKEN']
+
+        headers = {'cookie': cookie.key + '=' + cookie.value}
+        async with session.ws_connect(settings.UFP_ADDRESS + '/proxy/protect/ws/updates', headers=headers, ssl=False) as ws:
+            async for msg in ws:
+                packet_type, payload_format, deflated, unknown, payload_size = struct.unpack('!bbbbi', msg.data[0:HEADER_LENGTH])
+                action_start = HEADER_LENGTH
+                action_packet = zlib.decompress(msg.data[action_start:])
+                data_start = payload_size + 2*HEADER_LENGTH
+                data_packet = zlib.decompress(msg.data[data_start:])
+
+                yield json.loads(data_packet.decode())
+
+
+async def test():
+    async for msg in connect():
+        print(msg)
+
+if __name__ == '__main__':
+    loop = asyncio.get_event_loop()
+    loop.run_until_complete(test())
+    loop.close()