nonstop/management/commands/switch-proxy.py

   1 # proxy between arduino switch and websockets
   2 # it listens on 0.0.0.0, UDP port 1312, "websocket" port 8765
   3
   4 import asyncio
   5 import json
   6 import logging
   7
   8 import requests
   9 import websockets
  10
  11 from nonstop.app_settings import app_settings
  12
  13
  14 logger = logging.getLogger('switch-proxy')
  15
  16
  17 class PanikSwitchProtocol(asyncio.DatagramProtocol):
  18     UDP_PORT = 1312
  19
  20     def __init__(self):
  21         super().__init__()
  22         self.currently_active = 0
  23         self.update_with_http()
  24         self.websocket_connections = set()
  25
  26     def datagram_received(self, data, addr):
  27         logger.debug('Datagram received: %s', data)
  28         new_active = json.loads(data)['active']
  29         if new_active != self.currently_active:
  30             logger.info('UDP update, %s -> %s', self.currently_active, new_active)
  31         self.currently_active = json.loads(data)['active']
  32
  33     async def websocket_handler(self, websocket, path):
  34         self.websocket_connections.add(websocket)
  35         latest_active = None
  36         while True:
  37             if self.currently_active != latest_active:
  38                 latest_active = self.currently_active
  39                 to_trash = []
  40                 for ws in self.websocket_connections:
  41                     try:
  42                         await ws.send(json.dumps({'active': self.currently_active}))
  43                     except websockets.exceptions.ConnectionClosed:
  44                         to_trash.append(ws)
  45                 for ws in to_trash:
  46                     self.websocket_connections.remove(ws)
  47             await asyncio.sleep(0.1)
  48
  49     def update_with_http(self):
  50         try:
  51             resp = requests.get(app_settings.ON_AIR_SWITCH_URL, timeout=2)
  52             if resp.ok:
  53                 new_active = resp.json().get('active')
  54                 if new_active != self.currently_active:
  55                     logger.info('HTTP-only update, %s -> %s', self.currently_active, new_active)
  56                 self.currently_active = resp.json().get('active')
  57         except (OSError, requests.exceptions.RequestException):
  58             pass
  59
  60
  61 async def main():
  62     loop = asyncio.get_running_loop()
  63     proto = PanikSwitchProtocol()
  64     await websockets.serve(proto.websocket_handler, '0.0.0.0', 8765)
  65     transport, protocol = await loop.create_datagram_endpoint(lambda: proto, local_addr=('0.0.0.0', 1312))
  66     try:
  67         while True:
  68             await asyncio.sleep(300)
  69             proto.update_with_http()
  70     finally:
  71         transport.close()
  72
  73
  74 asyncio.run(main())