gps303/wsgateway.py

   1 """ Websocket Gateway """
   2
   3 from datetime import datetime, timezone
   4 from json import dumps, loads
   5 from logging import getLogger
   6 from socket import socket, AF_INET6, SOCK_STREAM, SOL_SOCKET, SO_REUSEADDR
   7 from time import time
   8 from wsproto import ConnectionType, WSConnection
   9 from wsproto.events import (
  10     AcceptConnection,
  11     CloseConnection,
  12     Message,
  13     Ping,
  14     Request,
  15     TextMessage,
  16 )
  17 from wsproto.utilities import RemoteProtocolError
  18 import zmq
  19
  20 from . import common
  21 from .evstore import initdb, fetch
  22 from .gps303proto import (
  23     GPS_POSITIONING,
  24     WIFI_POSITIONING,
  25     parse_message,
  26 )
  27 from .zmsg import Bcast, topic
  28
  29 log = getLogger("gps303/wsgateway")
  30 htmlfile = None
  31
  32
  33 def backlog(imei, numback):
  34     result = []
  35     for is_incoming, timestamp, packet in fetch(
  36         imei,
  37         ((True, GPS_POSITIONING.PROTO), (False, WIFI_POSITIONING.PROTO)),
  38         numback,
  39     ):
  40         msg = parse_message(packet, is_incoming=is_incoming)
  41         result.append(
  42             {
  43                 "imei": imei,
  44                 "timestamp": str(
  45                     datetime.fromtimestamp(timestamp).astimezone(
  46                         tz=timezone.utc
  47                     )
  48                 ),
  49                 "longitude": msg.longitude,
  50                 "latitude": msg.latitude,
  51             }
  52         )
  53     return result
  54
  55
  56 def try_http(data, fd, e):
  57     global htmlfile
  58     try:
  59         lines = data.decode().split("\r\n")
  60         request = lines[0]
  61         headers = lines[1:]
  62         op, resource, proto = request.split(" ")
  63         log.debug(
  64             "HTTP %s for %s, proto %s from fd %d, headers: %s",
  65             op,
  66             resource,
  67             proto,
  68             fd,
  69             headers,
  70         )
  71         try:
  72             pos = resource.index("?")
  73             resource = resource[:pos]
  74         except ValueError:
  75             pass
  76         if op == "GET":
  77             if htmlfile is None:
  78                 return (
  79                     f"{proto} 500 No data configured\r\n"
  80                     f"Content-Type: text/plain\r\n\r\n"
  81                     f"HTML data not configure on the server\r\n".encode()
  82                 )
  83             elif resource == "/":
  84                 try:
  85                     with open(htmlfile, "rb") as fl:
  86                         htmldata = fl.read()
  87                     length = len(htmldata)
  88                     return (
  89                         f"{proto} 200 Ok\r\n"
  90                         f"Content-Type: text/html; charset=utf-8\r\n"
  91                         f"Content-Length: {len(htmldata):d}\r\n\r\n"
  92                     ).encode("utf-8") + htmldata
  93                 except OSError:
  94                     return (
  95                         f"{proto} 500 File not found\r\n"
  96                         f"Content-Type: text/plain\r\n\r\n"
  97                         f"HTML file could not be opened\r\n".encode()
  98                     )
  99             else:
 100                 return (
 101                     f"{proto} 404 File not found\r\n"
 102                     f"Content-Type: text/plain\r\n\r\n"
 103                     f'We can only serve "/"\r\n'.encode()
 104                 )
 105         else:
 106             return (
 107                 f"{proto} 400 Bad request\r\n"
 108                 "Content-Type: text/plain\r\n\r\n"
 109                 "Bad request\r\n".encode()
 110             )
 111     except ValueError:
 112         log.warning("Unparseable data from fd %d: %s", fd, data)
 113         raise e
 114
 115
 116 class Client:
 117     """Websocket connection to the client"""
 118
 119     def __init__(self, sock, addr):
 120         self.sock = sock
 121         self.addr = addr
 122         self.ws = WSConnection(ConnectionType.SERVER)
 123         self.ws_data = b""
 124         self.ready = False
 125         self.imeis = set()
 126
 127     def close(self):
 128         log.debug("Closing fd %d", self.sock.fileno())
 129         self.sock.close()
 130
 131     def recv(self):
 132         try:
 133             data = self.sock.recv(4096)
 134         except OSError:
 135             log.warning(
 136                 "Reading from fd %d: %s",
 137                 self.sock.fileno(),
 138                 e,
 139             )
 140             self.ws.receive_data(None)
 141             return None
 142         if not data:  # Client has closed connection
 143             log.info(
 144                 "EOF reading from fd %d",
 145                 self.sock.fileno(),
 146             )
 147             self.ws.receive_data(None)
 148             return None
 149         try:
 150             self.ws.receive_data(data)
 151         except RemoteProtocolError as e:
 152             log.debug(
 153                 "Websocket error on fd %d, try plain http (%s)",
 154                 self.sock.fileno(),
 155                 e,
 156             )
 157             self.ws_data = try_http(data, self.sock.fileno(), e)
 158             self.write()  # TODO this is a hack
 159             log.debug("Sending HTTP response to %d", self.sock.fileno())
 160             msgs = None
 161         else:
 162             msgs = []
 163             for event in self.ws.events():
 164                 if isinstance(event, Request):
 165                     log.debug("WebSocket upgrade on fd %d", self.sock.fileno())
 166                     # self.ws_data += self.ws.send(event.response())  # Why not?!
 167                     self.ws_data += self.ws.send(AcceptConnection())
 168                     self.ready = True
 169                 elif isinstance(event, (CloseConnection, Ping)):
 170                     log.debug("%s on fd %d", event, self.sock.fileno())
 171                     self.ws_data += self.ws.send(event.response())
 172                 elif isinstance(event, TextMessage):
 173                     log.debug("%s on fd %d", event, self.sock.fileno())
 174                     msg = loads(event.data)
 175                     msgs.append(msg)
 176                     if msg.get("type", None) == "subscribe":
 177                         self.imeis = set(msg.get("imei", []))
 178                         log.debug(
 179                             "subs list on fd %s is %s",
 180                             self.sock.fileno(),
 181                             self.imeis,
 182                         )
 183                 else:
 184                     log.warning("%s on fd %d", event, self.sock.fileno())
 185         return msgs
 186
 187     def wants(self, imei):
 188         log.debug(
 189             "wants %s? set is %s on fd %d",
 190             imei,
 191             self.imeis,
 192             self.sock.fileno(),
 193         )
 194         return imei in self.imeis
 195
 196     def send(self, message):
 197         if self.ready and message["imei"] in self.imeis:
 198             self.ws_data += self.ws.send(Message(data=dumps(message)))
 199
 200     def write(self):
 201         if self.ws_data:
 202             try:
 203                 sent = self.sock.send(self.ws_data)
 204                 self.ws_data = self.ws_data[sent:]
 205             except OSError as e:
 206                 log.error(
 207                     "Sending to fd %d: %s",
 208                     self.sock.fileno(),
 209                     e,
 210                 )
 211                 self.ws_data = b""
 212         return bool(self.ws_data)
 213
 214
 215 class Clients:
 216     def __init__(self):
 217         self.by_fd = {}
 218
 219     def add(self, clntsock, clntaddr):
 220         fd = clntsock.fileno()
 221         log.info("Start serving fd %d from %s", fd, clntaddr)
 222         self.by_fd[fd] = Client(clntsock, clntaddr)
 223         return fd
 224
 225     def stop(self, fd):
 226         clnt = self.by_fd[fd]
 227         log.info("Stop serving fd %d", clnt.sock.fileno())
 228         clnt.close()
 229         del self.by_fd[fd]
 230
 231     def recv(self, fd):
 232         clnt = self.by_fd[fd]
 233         return clnt.recv()
 234
 235     def send(self, msg):
 236         towrite = set()
 237         for fd, clnt in self.by_fd.items():
 238             if clnt.wants(msg["imei"]):
 239                 clnt.send(msg)
 240                 towrite.add(fd)
 241         return towrite
 242
 243     def write(self, towrite):
 244         waiting = set()
 245         for fd, clnt in [(fd, self.by_fd.get(fd)) for fd in towrite]:
 246             if clnt.write():
 247                 waiting.add(fd)
 248         return waiting
 249
 250     def subs(self):
 251         result = set()
 252         for clnt in self.by_fd.values():
 253             result |= clnt.imeis
 254         return result
 255
 256
 257 def runserver(conf):
 258     global htmlfile
 259
 260     initdb(conf.get("storage", "dbfn"))
 261     htmlfile = conf.get("wsgateway", "htmlfile")
 262     zctx = zmq.Context()
 263     zsub = zctx.socket(zmq.SUB)
 264     zsub.connect(conf.get("collector", "publishurl"))
 265     tcpl = socket(AF_INET6, SOCK_STREAM)
 266     tcpl.setblocking(False)
 267     tcpl.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
 268     tcpl.bind(("", conf.getint("wsgateway", "port")))
 269     tcpl.listen(5)
 270     tcpfd = tcpl.fileno()
 271     poller = zmq.Poller()
 272     poller.register(zsub, flags=zmq.POLLIN)
 273     poller.register(tcpfd, flags=zmq.POLLIN)
 274     clients = Clients()
 275     activesubs = set()
 276     try:
 277         towait = set()
 278         while True:
 279             neededsubs = clients.subs()
 280             for imei in neededsubs - activesubs:
 281                 zsub.setsockopt(
 282                     zmq.SUBSCRIBE,
 283                     topic(GPS_POSITIONING.PROTO, True, imei),
 284                 )
 285                 zsub.setsockopt(
 286                     zmq.SUBSCRIBE,
 287                     topic(WIFI_POSITIONING.PROTO, False, imei),
 288                 )
 289             for imei in activesubs - neededsubs:
 290                 zsub.setsockopt(
 291                     zmq.UNSUBSCRIBE,
 292                     topic(GPS_POSITIONING.PROTO, True, imei),
 293                 )
 294                 zsub.setsockopt(
 295                     zmq.UNSUBSCRIBE,
 296                     topic(WIFI_POSITIONING.PROTO, False, imei),
 297                 )
 298             activesubs = neededsubs
 299             log.debug("Subscribed to: %s", activesubs)
 300             tosend = []
 301             topoll = []
 302             tostop = []
 303             towrite = set()
 304             events = poller.poll()
 305             for sk, fl in events:
 306                 if sk is zsub:
 307                     while True:
 308                         try:
 309                             zmsg = Bcast(zsub.recv(zmq.NOBLOCK))
 310                             msg = parse_message(zmsg.packet, zmsg.is_incoming)
 311                             log.debug("Got %s with %s", zmsg, msg)
 312                             tosend.append(
 313                                 {
 314                                     "imei": zmsg.imei,
 315                                     "timestamp": str(
 316                                         datetime.fromtimestamp(
 317                                             zmsg.when
 318                                         ).astimezone(tz=timezone.utc)
 319                                     ),
 320                                     "longitude": msg.longitude,
 321                                     "latitude": msg.latitude,
 322                                 }
 323                             )
 324                         except zmq.Again:
 325                             break
 326                 elif sk == tcpfd:
 327                     clntsock, clntaddr = tcpl.accept()
 328                     topoll.append((clntsock, clntaddr))
 329                 elif fl & zmq.POLLIN:
 330                     received = clients.recv(sk)
 331                     if received is None:
 332                         log.debug("Client gone from fd %d", sk)
 333                         tostop.append(sk)
 334                         towait.discard(fd)
 335                     else:
 336                         for msg in received:
 337                             log.debug("Received from %d: %s", sk, msg)
 338                             if msg.get("type", None) == "subscribe":
 339                                 imeis = msg.get("imei")
 340                                 numback = msg.get("backlog", 5)
 341                                 for imei in imeis:
 342                                     tosend.extend(backlog(imei, numback))
 343                         towrite.add(sk)
 344                 elif fl & zmq.POLLOUT:
 345                     log.debug("Write now open for fd %d", sk)
 346                     towrite.add(sk)
 347                     towait.discard(sk)
 348                 else:
 349                     log.debug("Stray event: %s on socket %s", fl, sk)
 350             # poll queue consumed, make changes now
 351             for fd in tostop:
 352                 poller.unregister(fd)
 353                 clients.stop(fd)
 354             for zmsg in tosend:
 355                 log.debug("Sending to the clients: %s", zmsg)
 356                 towrite |= clients.send(zmsg)
 357             for clntsock, clntaddr in topoll:
 358                 fd = clients.add(clntsock, clntaddr)
 359                 poller.register(fd, flags=zmq.POLLIN)
 360             # Deal with actually writing the data out
 361             trywrite = towrite - towait
 362             morewait = clients.write(trywrite)
 363             log.debug(
 364                 "towait %s, tried %s, still busy %s",
 365                 towait,
 366                 trywrite,
 367                 morewait,
 368             )
 369             for fd in morewait - trywrite:  # new fds waiting for write
 370                 poller.modify(fd, flags=zmq.POLLIN | zmq.POLLOUT)
 371             for fd in trywrite - morewait:  # no longer waiting for write
 372                 poller.modify(fd, flags=zmq.POLLIN)
 373             towait &= trywrite
 374             towait |= morewait
 375     except KeyboardInterrupt:
 376         pass
 377
 378
 379 if __name__.endswith("__main__"):
 380     runserver(common.init(log))