]> www.average.org Git - loctrkd.git/blobdiff - gps303/wsgateway.py
typeckecking: annotate wsgateway.py
[loctrkd.git] / gps303 / wsgateway.py
index cc1cb7ec08dc0db5baafbfd89545e60bdced85d2..c0a47768254bff754a2425ea1cd40ae597002a06 100644 (file)
 """ Websocket Gateway """
 
+from configparser import ConfigParser
+from datetime import datetime, timezone
+from json import dumps, loads
 from logging import getLogger
 from socket import socket, AF_INET6, SOCK_STREAM, SOL_SOCKET, SO_REUSEADDR
 from time import time
+from typing import Any, cast, Dict, List, Optional, Set, Tuple
 from wsproto import ConnectionType, WSConnection
-from wsproto.events import AcceptConnection, CloseConnection, Message, Ping, Request, TextMessage
+from wsproto.events import (
+    AcceptConnection,
+    CloseConnection,
+    Event,
+    Message,
+    Ping,
+    Request,
+    TextMessage,
+)
+from wsproto.utilities import RemoteProtocolError
 import zmq
 
 from . import common
-from .zmsg import LocEvt
+from .evstore import initdb, fetch
+from .gps303proto import (
+    GPS_POSITIONING,
+    STATUS,
+    WIFI_POSITIONING,
+    parse_message,
+)
+from .zmsg import Bcast, topic
 
 log = getLogger("gps303/wsgateway")
+htmlfile = None
+
+
+def backlog(imei: str, numback: int) -> List[Dict[str, Any]]:
+    result = []
+    for is_incoming, timestamp, packet in fetch(
+        imei,
+        [(True, GPS_POSITIONING.PROTO), (False, WIFI_POSITIONING.PROTO)],
+        numback,
+    ):
+        msg = parse_message(packet, is_incoming=is_incoming)
+        result.append(
+            {
+                "type": "location",
+                "imei": imei,
+                "timestamp": str(
+                    datetime.fromtimestamp(timestamp).astimezone(
+                        tz=timezone.utc
+                    )
+                ),
+                "longitude": msg.longitude,
+                "latitude": msg.latitude,
+                "accuracy": "gps"
+                if isinstance(msg, GPS_POSITIONING)
+                else "approximate",
+            }
+        )
+    return result
+
+
+def try_http(data: bytes, fd: int, e: Exception) -> bytes:
+    global htmlfile
+    try:
+        lines = data.decode().split("\r\n")
+        request = lines[0]
+        headers = lines[1:]
+        op, resource, proto = request.split(" ")
+        log.debug(
+            "HTTP %s for %s, proto %s from fd %d, headers: %s",
+            op,
+            resource,
+            proto,
+            fd,
+            headers,
+        )
+        if op == "GET":
+            if htmlfile is None:
+                return (
+                    f"{proto} 500 No data configured\r\n"
+                    f"Content-Type: text/plain\r\n\r\n"
+                    f"HTML data not configured on the server\r\n".encode()
+                )
+            else:
+                try:
+                    with open(htmlfile, "rb") as fl:
+                        htmldata = fl.read()
+                    length = len(htmldata)
+                    return (
+                        f"{proto} 200 Ok\r\n"
+                        f"Content-Type: text/html; charset=utf-8\r\n"
+                        f"Content-Length: {len(htmldata):d}\r\n\r\n"
+                    ).encode("utf-8") + htmldata
+                except OSError:
+                    return (
+                        f"{proto} 500 File not found\r\n"
+                        f"Content-Type: text/plain\r\n\r\n"
+                        f"HTML file could not be opened\r\n".encode()
+                    )
+        else:
+            return (
+                f"{proto} 400 Bad request\r\n"
+                "Content-Type: text/plain\r\n\r\n"
+                "Bad request\r\n".encode()
+            )
+    except ValueError:
+        log.warning("Unparseable data from fd %d: %s", fd, data)
+        raise e
 
 
 class Client:
     """Websocket connection to the client"""
 
-    def __init__(self, sock, addr):
+    def __init__(self, sock: socket, addr: Tuple[str, int]) -> None:
         self.sock = sock
         self.addr = addr
         self.ws = WSConnection(ConnectionType.SERVER)
         self.ws_data = b""
+        self.ready = False
+        self.imeis: Set[str] = set()
 
-    def close(self):
+    def close(self) -> None:
         log.debug("Closing fd %d", self.sock.fileno())
         self.sock.close()
 
-    def recv(self):
+    def recv(self) -> Optional[List[Dict[str, Any]]]:
         try:
             data = self.sock.recv(4096)
-        except OSError:
+        except OSError as e:
             log.warning(
                 "Reading from fd %d: %s",
                 self.sock.fileno(),
@@ -44,102 +143,208 @@ class Client:
             )
             self.ws.receive_data(None)
             return None
-        self.ws.receive_data(data)
-        msgs = []
-        for event in self.ws.events():
-            if isinstance(event, Request):
-                log.debug("WebSocket upgrade on fd %d", self.sock.fileno())
-                #self.ws_data += self.ws.send(event.response())  # Why not?!
-                self.ws_data += self.ws.send(AcceptConnection())
-            elif isinstance(event, (CloseConnection, Ping)):
-                log.debug("%s on fd %d", event, self.sock.fileno())
-                self.ws_data += self.ws.send(event.response())
-            elif isinstance(event, TextMessage):
-                # TODO: save imei "subscription"
-                log.debug("%s on fd %d", event, self.sock.fileno())
-                msgs.append(event.data)
-            else:
-                log.warning("%s on fd %d", event, self.sock.fileno())
-        if self.ws_data:  # Temp hack
-            self.write()
-        return msgs
-
-    def send(self, imei, message):
-        # TODO: filter only wanted imei got from the client
-        self.ws_data += self.ws.send(Message(data=message))
-
-    def write(self):
         try:
-            sent = self.sock.send(self.ws_data)
-            self.ws_data = self.ws_data[sent:]
-        except OSError as e:
-            log.error(
-                "Sending to fd %d (IMEI %s): %s",
+            self.ws.receive_data(data)
+        except RemoteProtocolError as e:
+            log.debug(
+                "Websocket error on fd %d, try plain http (%s)",
                 self.sock.fileno(),
-                self.imei,
                 e,
             )
-            self.ws_data = b""
+            self.ws_data = try_http(data, self.sock.fileno(), e)
+            self.write()  # TODO this is a hack
+            log.debug("Sending HTTP response to %d", self.sock.fileno())
+            msgs = None
+        else:
+            msgs = []
+            for event in self.ws.events():
+                if isinstance(event, Request):
+                    log.debug("WebSocket upgrade on fd %d", self.sock.fileno())
+                    # self.ws_data += self.ws.send(event.response())  # Why not?!
+                    self.ws_data += self.ws.send(AcceptConnection())
+                    self.ready = True
+                elif isinstance(event, (CloseConnection, Ping)):
+                    log.debug("%s on fd %d", event, self.sock.fileno())
+                    self.ws_data += self.ws.send(event.response())
+                elif isinstance(event, TextMessage):
+                    log.debug("%s on fd %d", event, self.sock.fileno())
+                    msg = loads(event.data)
+                    msgs.append(msg)
+                    if msg.get("type", None) == "subscribe":
+                        self.imeis = set(msg.get("imei", []))
+                        log.debug(
+                            "subs list on fd %s is %s",
+                            self.sock.fileno(),
+                            self.imeis,
+                        )
+                else:
+                    log.warning("%s on fd %d", event, self.sock.fileno())
+        return msgs
+
+    def wants(self, imei: str) -> bool:
+        log.debug(
+            "wants %s? set is %s on fd %d",
+            imei,
+            self.imeis,
+            self.sock.fileno(),
+        )
+        return imei in self.imeis
+
+    def send(self, message: Dict[str, Any]) -> None:
+        if self.ready and message["imei"] in self.imeis:
+            self.ws_data += self.ws.send(Message(data=dumps(message)))
+
+    def write(self) -> bool:
+        if self.ws_data:
+            try:
+                sent = self.sock.send(self.ws_data)
+                self.ws_data = self.ws_data[sent:]
+            except OSError as e:
+                log.error(
+                    "Sending to fd %d: %s",
+                    self.sock.fileno(),
+                    e,
+                )
+                self.ws_data = b""
+        return bool(self.ws_data)
 
 
 class Clients:
-    def __init__(self):
-        self.by_fd = {}
+    def __init__(self) -> None:
+        self.by_fd: Dict[int, Client] = {}
 
-    def add(self, clntsock, clntaddr):
+    def add(self, clntsock: socket, clntaddr: Tuple[str, int]) -> int:
         fd = clntsock.fileno()
         log.info("Start serving fd %d from %s", fd, clntaddr)
         self.by_fd[fd] = Client(clntsock, clntaddr)
         return fd
 
-    def stop(self, fd):
+    def stop(self, fd: int) -> None:
         clnt = self.by_fd[fd]
         log.info("Stop serving fd %d", clnt.sock.fileno())
         clnt.close()
         del self.by_fd[fd]
 
-    def recv(self, fd):
+    def recv(self, fd: int) -> Optional[List[Dict[str, Any]]]:
         clnt = self.by_fd[fd]
-        msgs = clnt.recv()
-        if msgs is None:
-            return None
-        result = []
-        for msg in msgs:
-            log.debug("Received: %s", msg)
-        return result
+        return clnt.recv()
+
+    def send(self, msg: Dict[str, Any]) -> Set[int]:
+        towrite = set()
+        for fd, clnt in self.by_fd.items():
+            if clnt.wants(msg["imei"]):
+                clnt.send(msg)
+                towrite.add(fd)
+        return towrite
+
+    def write(self, towrite: Set[int]) -> Set[int]:
+        waiting = set()
+        for fd, clnt in [(fd, self.by_fd.get(fd)) for fd in towrite]:
+            if clnt and clnt.write():
+                waiting.add(fd)
+        return waiting
 
-    def send(self, msgs):
+    def subs(self) -> Set[str]:
+        result = set()
         for clnt in self.by_fd.values():
-            clnt.send(msgs)
-            clnt.write()
-
-def runserver(conf):
-    zctx = zmq.Context()
-    zsub = zctx.socket(zmq.SUB)
-    zsub.connect(conf.get("lookaside", "publishurl"))
-    zsub.setsockopt(zmq.SUBSCRIBE, b"")
+            result |= clnt.imeis
+        return result
+
+
+def runserver(conf: ConfigParser) -> None:
+    global htmlfile
+
+    initdb(conf.get("storage", "dbfn"))
+    htmlfile = conf.get("wsgateway", "htmlfile")
+    # Is this https://github.com/zeromq/pyzmq/issues/1627 still not fixed?!
+    zctx = zmq.Context()  # type: ignore
+    zsub = zctx.socket(zmq.SUB)  # type: ignore
+    zsub.connect(conf.get("collector", "publishurl"))
     tcpl = socket(AF_INET6, SOCK_STREAM)
     tcpl.setblocking(False)
     tcpl.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
     tcpl.bind(("", conf.getint("wsgateway", "port")))
     tcpl.listen(5)
     tcpfd = tcpl.fileno()
-    poller = zmq.Poller()
+    poller = zmq.Poller()  # type: ignore
     poller.register(zsub, flags=zmq.POLLIN)
     poller.register(tcpfd, flags=zmq.POLLIN)
     clients = Clients()
+    activesubs: Set[str] = set()
     try:
+        towait: Set[int] = set()
         while True:
+            neededsubs = clients.subs()
+            for imei in neededsubs - activesubs:
+                zsub.setsockopt(
+                    zmq.SUBSCRIBE,
+                    topic(GPS_POSITIONING.PROTO, True, imei),
+                )
+                zsub.setsockopt(
+                    zmq.SUBSCRIBE,
+                    topic(WIFI_POSITIONING.PROTO, False, imei),
+                )
+                zsub.setsockopt(
+                    zmq.SUBSCRIBE,
+                    topic(STATUS.PROTO, True, imei),
+                )
+            for imei in activesubs - neededsubs:
+                zsub.setsockopt(
+                    zmq.UNSUBSCRIBE,
+                    topic(GPS_POSITIONING.PROTO, True, imei),
+                )
+                zsub.setsockopt(
+                    zmq.UNSUBSCRIBE,
+                    topic(WIFI_POSITIONING.PROTO, False, imei),
+                )
+                zsub.setsockopt(
+                    zmq.UNSUBSCRIBE,
+                    topic(STATUS.PROTO, True, imei),
+                )
+            activesubs = neededsubs
+            log.debug("Subscribed to: %s", activesubs)
             tosend = []
             topoll = []
             tostop = []
-            events = poller.poll(1000)
+            towrite = set()
+            events = poller.poll()
             for sk, fl in events:
                 if sk is zsub:
                     while True:
                         try:
-                            msg = zsub.recv(zmq.NOBLOCK)
-                            tosend.append(LocEvt(msg))
+                            zmsg = Bcast(zsub.recv(zmq.NOBLOCK))
+                            msg = parse_message(zmsg.packet, zmsg.is_incoming)
+                            log.debug("Got %s with %s", zmsg, msg)
+                            if isinstance(msg, STATUS):
+                                tosend.append(
+                                    {
+                                        "type": "status",
+                                        "imei": zmsg.imei,
+                                        "timestamp": str(
+                                            datetime.fromtimestamp(
+                                                zmsg.when
+                                            ).astimezone(tz=timezone.utc)
+                                        ),
+                                        "battery": msg.batt,
+                                    }
+                                )
+                            else:
+                                tosend.append(
+                                    {
+                                        "type": "location",
+                                        "imei": zmsg.imei,
+                                        "timestamp": str(
+                                            datetime.fromtimestamp(
+                                                zmsg.when
+                                            ).astimezone(tz=timezone.utc)
+                                        ),
+                                        "longitude": msg.longitude,
+                                        "latitude": msg.latitude,
+                                        "accuracy": "gps"
+                                        if zmsg.is_incoming
+                                        else "approximate",
+                                    }
+                                )
                         except zmq.Again:
                             break
                 elif sk == tcpfd:
@@ -148,26 +353,50 @@ def runserver(conf):
                 elif fl & zmq.POLLIN:
                     received = clients.recv(sk)
                     if received is None:
-                        log.debug(
-                            "Client gone from fd %d", sk
-                        )
+                        log.debug("Client gone from fd %d", sk)
                         tostop.append(sk)
+                        towait.discard(sk)
                     else:
-                        for msg in received:
-                            log.debug("Received from %d: %s", sk, msg)
+                        for wsmsg in received:
+                            log.debug("Received from %d: %s", sk, wsmsg)
+                            if wsmsg.get("type", None) == "subscribe":
+                                # Have to live w/o typeckeding from json
+                                imeis = cast(List[str], wsmsg.get("imei"))
+                                numback: int = wsmsg.get("backlog", 5)
+                                for imei in imeis:
+                                    tosend.extend(backlog(imei, numback))
+                        towrite.add(sk)
+                elif fl & zmq.POLLOUT:
+                    log.debug("Write now open for fd %d", sk)
+                    towrite.add(sk)
+                    towait.discard(sk)
                 else:
                     log.debug("Stray event: %s on socket %s", fl, sk)
             # poll queue consumed, make changes now
             for fd in tostop:
-                poller.unregister(fd)
+                poller.unregister(fd)  # type: ignore
                 clients.stop(fd)
-            for zmsg in tosend:
-                log.debug("Sending to the client: %s", zmsg)
-                clients.send(zmsg)
+            for wsmsg in tosend:
+                log.debug("Sending to the clients: %s", wsmsg)
+                towrite |= clients.send(wsmsg)
             for clntsock, clntaddr in topoll:
                 fd = clients.add(clntsock, clntaddr)
                 poller.register(fd, flags=zmq.POLLIN)
-            # TODO: Handle write overruns (register for POLLOUT)
+            # Deal with actually writing the data out
+            trywrite = towrite - towait
+            morewait = clients.write(trywrite)
+            log.debug(
+                "towait %s, tried %s, still busy %s",
+                towait,
+                trywrite,
+                morewait,
+            )
+            for fd in morewait - trywrite:  # new fds waiting for write
+                poller.modify(fd, flags=zmq.POLLIN | zmq.POLLOUT)  # type: ignore
+            for fd in trywrite - morewait:  # no longer waiting for write
+                poller.modify(fd, flags=zmq.POLLIN)  # type: ignore
+            towait &= trywrite
+            towait |= morewait
     except KeyboardInterrupt:
         pass