Implement our own unifi websocket connection

This commit is contained in:
Tanner Collin 2022-04-01 17:34:37 -06:00
parent 14051edfc5
commit 055b12ee05
3 changed files with 76 additions and 62 deletions

87
main.py
View File

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

View File

@ -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

41
unifi.py Normal file
View File

@ -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()