gps303/collector.py

   1 """ TCP server that communicates with terminals """
   2
   3 from configparser import ConfigParser
   4 from importlib import import_module
   5 from logging import getLogger
   6 from os import umask
   7 from socket import (
   8     socket,
   9     AF_INET6,
  10     SOCK_STREAM,
  11     SOL_SOCKET,
  12     SO_KEEPALIVE,
  13     SO_REUSEADDR,
  14 )
  15 from struct import pack
  16 from time import time
  17 from typing import Any, cast, Dict, List, Optional, Tuple, Union
  18 import zmq
  19
  20 from . import common
  21 from .zmsg import Bcast, Resp
  22
  23 log = getLogger("gps303/collector")
  24
  25 MAXBUFFER: int = 4096
  26
  27
  28 class ProtoModule:
  29     class Stream:
  30         @staticmethod
  31         def enframe(buffer: bytes) -> bytes:
  32             ...
  33
  34         def recv(self, segment: bytes) -> List[Union[bytes, str]]:
  35             ...
  36
  37         def close(self) -> bytes:
  38             ...
  39
  40     @staticmethod
  41     def probe_buffer(buffer: bytes) -> bool:
  42         ...
  43
  44     @staticmethod
  45     def parse_message(packet: bytes, is_incoming: bool = True) -> Any:
  46         ...
  47
  48     @staticmethod
  49     def inline_response(packet: bytes) -> Optional[bytes]:
  50         ...
  51
  52     @staticmethod
  53     def is_goodbye_packet(packet: bytes) -> bool:
  54         ...
  55
  56     @staticmethod
  57     def imei_from_packet(packet: bytes) -> Optional[str]:
  58         ...
  59
  60     @staticmethod
  61     def proto_of_message(packet: bytes) -> int:
  62         ...
  63
  64     @staticmethod
  65     def proto_by_name(name: str) -> int:
  66         ...
  67
  68
  69 pmods: List[ProtoModule] = []
  70
  71
  72 class Client:
  73     """Connected socket to the terminal plus buffer and metadata"""
  74
  75     def __init__(self, sock: socket, addr: Tuple[str, int]) -> None:
  76         self.sock = sock
  77         self.addr = addr
  78         self.pmod: Optional[ProtoModule] = None
  79         self.stream: Optional[ProtoModule.Stream] = None
  80         self.imei: Optional[str] = None
  81
  82     def close(self) -> None:
  83         log.debug("Closing fd %d (IMEI %s)", self.sock.fileno(), self.imei)
  84         self.sock.close()
  85         if self.stream:
  86             rest = self.stream.close()
  87         else:
  88             rest = b""
  89         if rest:
  90             log.warning("%d bytes in buffer on close: %s", len(rest), rest)
  91
  92     def recv(self) -> Optional[List[Tuple[float, Tuple[str, int], bytes]]]:
  93         """Read from the socket and parse complete messages"""
  94         try:
  95             segment = self.sock.recv(MAXBUFFER)
  96         except OSError as e:
  97             log.warning(
  98                 "Reading from fd %d (IMEI %s): %s",
  99                 self.sock.fileno(),
 100                 self.imei,
 101                 e,
 102             )
 103             return None
 104         if not segment:  # Terminal has closed connection
 105             log.info(
 106                 "EOF reading from fd %d (IMEI %s)",
 107                 self.sock.fileno(),
 108                 self.imei,
 109             )
 110             return None
 111         if self.stream is None:
 112             for pmod in pmods:
 113                 if pmod.probe_buffer(segment):
 114                     self.pmod = pmod
 115                     self.stream = pmod.Stream()
 116                     break
 117         if self.stream is None:
 118             log.info(
 119                 "unrecognizable %d bytes of data %s from fd %d",
 120                 len(segment),
 121                 segment[:32].hex(),
 122                 self.sock.fileno(),
 123             )
 124             return []
 125         when = time()
 126         msgs = []
 127         for elem in self.stream.recv(segment):
 128             if isinstance(elem, bytes):
 129                 msgs.append((when, self.addr, elem))
 130             else:
 131                 log.warning(
 132                     "%s from fd %d (IMEI %s)",
 133                     elem,
 134                     self.sock.fileno(),
 135                     self.imei,
 136                 )
 137         return msgs
 138
 139     def send(self, buffer: bytes) -> None:
 140         assert self.stream is not None
 141         try:
 142             self.sock.send(self.stream.enframe(buffer))
 143         except OSError as e:
 144             log.error(
 145                 "Sending to fd %d (IMEI %s): %s",
 146                 self.sock.fileno(),
 147                 self.imei,
 148                 e,
 149             )
 150
 151
 152 class Clients:
 153     def __init__(self) -> None:
 154         self.by_fd: Dict[int, Client] = {}
 155         self.by_imei: Dict[str, Client] = {}
 156
 157     def add(self, clntsock: socket, clntaddr: Tuple[str, int]) -> int:
 158         fd = clntsock.fileno()
 159         log.info("Start serving fd %d from %s", fd, clntaddr)
 160         self.by_fd[fd] = Client(clntsock, clntaddr)
 161         return fd
 162
 163     def stop(self, fd: int) -> None:
 164         clnt = self.by_fd[fd]
 165         log.info("Stop serving fd %d (IMEI %s)", clnt.sock.fileno(), clnt.imei)
 166         clnt.close()
 167         if clnt.imei:
 168             del self.by_imei[clnt.imei]
 169         del self.by_fd[fd]
 170
 171     def recv(
 172         self, fd: int
 173     ) -> Optional[
 174         List[Tuple[ProtoModule, Optional[str], float, Tuple[str, int], bytes]]
 175     ]:
 176         clnt = self.by_fd[fd]
 177         msgs = clnt.recv()
 178         if msgs is None:
 179             return None
 180         result = []
 181         for when, peeraddr, packet in msgs:
 182             assert clnt.pmod is not None
 183             if clnt.imei is None:
 184                 imei = clnt.pmod.imei_from_packet(packet)
 185                 if imei is not None:
 186                     log.info("LOGIN from fd %d (IMEI %s)", fd, imei)
 187                     clnt.imei = imei
 188                     oldclnt = self.by_imei.get(clnt.imei)
 189                     if oldclnt is not None:
 190                         log.info(
 191                             "Orphaning fd %d with the same IMEI",
 192                             oldclnt.sock.fileno(),
 193                         )
 194                         oldclnt.imei = None
 195                     self.by_imei[clnt.imei] = clnt
 196                 else:
 197                     log.warning(
 198                         "Login message from %s: %s, but client imei unfilled",
 199                         peeraddr,
 200                         packet,
 201                     )
 202             result.append((clnt.pmod, clnt.imei, when, peeraddr, packet))
 203             log.debug(
 204                 "Received from %s (IMEI %s): %s",
 205                 peeraddr,
 206                 clnt.imei,
 207                 packet.hex(),
 208             )
 209         return result
 210
 211     def response(self, resp: Resp) -> Optional[ProtoModule]:
 212         if resp.imei in self.by_imei:
 213             clnt = self.by_imei[resp.imei]
 214             clnt.send(resp.packet)
 215             return clnt.pmod
 216         else:
 217             log.info("Not connected (IMEI %s)", resp.imei)
 218             return None
 219
 220
 221 def runserver(conf: ConfigParser, handle_hibernate: bool = True) -> None:
 222     global pmods
 223     pmods = [
 224         cast(ProtoModule, import_module("." + modnm, __package__))
 225         for modnm in conf.get("collector", "protocols").split(",")
 226     ]
 227     # Is this https://github.com/zeromq/pyzmq/issues/1627 still not fixed?!
 228     zctx = zmq.Context()  # type: ignore
 229     zpub = zctx.socket(zmq.PUB)  # type: ignore
 230     zpull = zctx.socket(zmq.PULL)  # type: ignore
 231     oldmask = umask(0o117)
 232     zpub.bind(conf.get("collector", "publishurl"))
 233     zpull.bind(conf.get("collector", "listenurl"))
 234     umask(oldmask)
 235     tcpl = socket(AF_INET6, SOCK_STREAM)
 236     tcpl.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
 237     tcpl.bind(("", conf.getint("collector", "port")))
 238     tcpl.listen(5)
 239     tcpfd = tcpl.fileno()
 240     poller = zmq.Poller()  # type: ignore
 241     poller.register(zpull, flags=zmq.POLLIN)
 242     poller.register(tcpfd, flags=zmq.POLLIN)
 243     clients = Clients()
 244     try:
 245         while True:
 246             tosend = []
 247             topoll = []
 248             tostop = []
 249             events = poller.poll(1000)
 250             for sk, fl in events:
 251                 if sk is zpull:
 252                     while True:
 253                         try:
 254                             msg = zpull.recv(zmq.NOBLOCK)
 255                             zmsg = Resp(msg)
 256                             tosend.append(zmsg)
 257                         except zmq.Again:
 258                             break
 259                 elif sk == tcpfd:
 260                     clntsock, clntaddr = tcpl.accept()
 261                     clntsock.setsockopt(SOL_SOCKET, SO_KEEPALIVE, 1)
 262                     topoll.append((clntsock, clntaddr))
 263                 elif fl & zmq.POLLIN:
 264                     received = clients.recv(sk)
 265                     if received is None:
 266                         log.debug("Terminal gone from fd %d", sk)
 267                         tostop.append(sk)
 268                     else:
 269                         for pmod, imei, when, peeraddr, packet in received:
 270                             proto = pmod.proto_of_message(packet)
 271                             zpub.send(
 272                                 Bcast(
 273                                     proto=proto,
 274                                     imei=imei,
 275                                     when=when,
 276                                     peeraddr=peeraddr,
 277                                     packet=packet,
 278                                 ).packed
 279                             )
 280                             if (
 281                                 pmod.is_goodbye_packet(packet)
 282                                 and handle_hibernate
 283                             ):
 284                                 log.debug(
 285                                     "Goodbye from fd %d (IMEI %s)",
 286                                     sk,
 287                                     imei,
 288                                 )
 289                                 tostop.append(sk)
 290                             respmsg = pmod.inline_response(packet)
 291                             if respmsg is not None:
 292                                 tosend.append(
 293                                     Resp(imei=imei, when=when, packet=respmsg)
 294                                 )
 295                 else:
 296                     log.debug("Stray event: %s on socket %s", fl, sk)
 297             # poll queue consumed, make changes now
 298             for zmsg in tosend:
 299                 log.debug("Sending to the client: %s", zmsg)
 300                 rpmod = clients.response(zmsg)
 301                 if rpmod is not None:
 302                     zpub.send(
 303                         Bcast(
 304                             is_incoming=False,
 305                             proto=rpmod.proto_of_message(zmsg.packet),
 306                             when=zmsg.when,
 307                             imei=zmsg.imei,
 308                             packet=zmsg.packet,
 309                         ).packed
 310                     )
 311             for fd in tostop:
 312                 poller.unregister(fd)  # type: ignore
 313                 clients.stop(fd)
 314             for clntsock, clntaddr in topoll:
 315                 fd = clients.add(clntsock, clntaddr)
 316                 poller.register(fd, flags=zmq.POLLIN)
 317     except KeyboardInterrupt:
 318         zpub.close()
 319         zpull.close()
 320         zctx.destroy()  # type: ignore
 321         tcpl.close()
 322
 323
 324 if __name__.endswith("__main__"):
 325     runserver(common.init(log))