nonstop/management/commands/switch-proxy.py

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