gps303/collector.py

   1 """ TCP server that communicates with terminals """
   2
   3 from getopt import getopt
   4 from logging import getLogger
   5 from logging.handlers import SysLogHandler
   6 from socket import socket, AF_INET, SOCK_STREAM, SOL_SOCKET, SO_REUSEADDR
   7 from time import time
   8 from struct import pack
   9 import zmq
  10
  11 from . import common
  12 from .gps303proto import HIBERNATION, LOGIN, parse_message, proto_of_message
  13
  14 log = getLogger("gps303/collector")
  15
  16
  17 class Bcast:
  18     """Zmq message to broadcast what was received from the terminal"""
  19
  20     def __init__(self, imei, msg):
  21         self.as_bytes = (
  22             pack("B", proto_of_message(msg))
  23             + ("0000000000000000" if imei is None else imei).encode()
  24             + msg
  25         )
  26
  27
  28 class Resp:
  29     """Zmq message received from a third party to send to the terminal"""
  30
  31     def __init__(self, *args, **kwargs):
  32         if not kwargs and len(args) == 1 and isinstance(args[0], bytes):
  33             self.imei = msg[:16].decode()
  34             self.payload = msg[16:]
  35         elif len(args) == 0:
  36             self.imei = kwargs["imei"]
  37             self.payload = kwargs["payload"]
  38
  39
  40 class Client:
  41     """Connected socket to the terminal plus buffer and metadata"""
  42
  43     def __init__(self, sock, addr):
  44         self.sock = sock
  45         self.addr = addr
  46         self.buffer = b""
  47         self.imei = None
  48
  49     def close(self):
  50         log.debug("Closing fd %d (IMEI %s)", self.sock.fileno(), self.imei)
  51         self.sock.close()
  52         self.buffer = b""
  53         self.imei = None
  54
  55     def recv(self):
  56         """Read from the socket and parse complete messages"""
  57         try:
  58             segment = self.sock.recv(4096)
  59         except OSError:
  60             log.warning(
  61                 "Reading from fd %d (IMEI %s): %s",
  62                 self.sock.fileno(),
  63                 self.imei,
  64                 e,
  65             )
  66             return None
  67         if not segment:  # Terminal has closed connection
  68             log.info(
  69                 "EOF reading from fd %d (IMEI %s)",
  70                 self.sock.fileno(),
  71                 self.imei,
  72             )
  73             return None
  74         when = time()
  75         self.buffer += segment
  76         msgs = []
  77         while True:
  78             framestart = self.buffer.find(b"xx")
  79             if framestart == -1:  # No frames, return whatever we have
  80                 break
  81             if framestart > 0:  # Should not happen, report
  82                 log.warning(
  83                     'Undecodable data "%s" from fd %d (IMEI %s)',
  84                     self.buffer[:framestart].hex(),
  85                     self.sock.fileno(),
  86                     self.imei,
  87                 )
  88                 self.buffer = self.buffer[framestart:]
  89             # At this point, buffer starts with a packet
  90             frameend = self.buffer.find(b"\r\n", 4)
  91             if frameend == -1:  # Incomplete frame, return what we have
  92                 break
  93             packet = self.buffer[2:frameend]
  94             self.buffer = self.buffer[frameend + 2 :]
  95             if proto_of_message(packet) == LOGIN.PROTO:
  96                 self.imei = parse_message(packet).imei
  97                 log.info(
  98                     "LOGIN from fd %d (IMEI %s)", self.sock.fileno(), self.imei
  99                 )
 100             msgs.append(packet)
 101         return msgs
 102
 103     def send(self, buffer):
 104         try:
 105             self.sock.send(b"xx" + buffer + b"\r\n")
 106         except OSError as e:
 107             log.error(
 108                 "Sending to fd %d (IMEI %s): %s",
 109                 self.sock.fileno,
 110                 self.imei,
 111                 e,
 112             )
 113
 114
 115 class Clients:
 116     def __init__(self):
 117         self.by_fd = {}
 118         self.by_imei = {}
 119
 120     def add(self, clntsock, clntaddr):
 121         fd = clntsock.fileno()
 122         self.by_fd[fd] = Client(clntsock, clntaddr)
 123         return fd
 124
 125     def stop(self, fd):
 126         clnt = self.by_fd[fd]
 127         log.info("Stop serving fd %d (IMEI %s)", clnt.sock.fileno(), clnt.imei)
 128         clnt.close()
 129         if clnt.imei:
 130             del self.by_imei[clnt.imei]
 131         del self.by_fd[fd]
 132
 133     def recv(self, fd):
 134         clnt = self.by_fd[fd]
 135         msgs = clnt.recv()
 136         if msgs is None:
 137             return None
 138         result = []
 139         for msg in msgs:
 140             if proto_of_message(msg) == LOGIN.PROTO:  # Could do blindly...
 141                 self.by_imei[clnt.imei] = clnt
 142             result.append((clnt.imei, msg))
 143         return result
 144
 145     def response(self, resp):
 146         if resp.imei in self.by_imei:
 147             self.by_imei[resp.imei].send(resp.payload)
 148
 149
 150 def runserver(conf):
 151     zctx = zmq.Context()
 152     zpub = zctx.socket(zmq.PUB)
 153     zpub.bind(conf.get("collector", "publishurl"))
 154     zsub = zctx.socket(zmq.SUB)
 155     zsub.connect(conf.get("collector", "listenurl"))
 156     tcpl = socket(AF_INET, SOCK_STREAM)
 157     tcpl.setsockopt(SOL_SOCKET, SO_REUSEADDR, 1)
 158     tcpl.bind(("", conf.getint("collector", "port")))
 159     tcpl.listen(5)
 160     tcpfd = tcpl.fileno()
 161     poller = zmq.Poller()
 162     poller.register(zsub, flags=zmq.POLLIN)
 163     poller.register(tcpfd, flags=zmq.POLLIN)
 164     clients = Clients()
 165     try:
 166         while True:
 167             tosend = []
 168             topoll = []
 169             tostop = []
 170             events = poller.poll(10)
 171             for sk, fl in events:
 172                 if sk is zsub:
 173                     while True:
 174                         try:
 175                             msg = zsub.recv(zmq.NOBLOCK)
 176                             tosend.append(Resp(msg))
 177                         except zmq.Again:
 178                             break
 179                 elif sk == tcpfd:
 180                     clntsock, clntaddr = tcpl.accept()
 181                     topoll.append((clntsock, clntaddr))
 182                 else:
 183                     for imei, msg in clients.recv(sk):
 184                         zpub.send(Bcast(imei, msg).as_bytes)
 185                         if (
 186                             msg is None
 187                             or proto_of_message(msg) == HIBERNATION.PROTO
 188                         ):
 189                             log.debug(
 190                                 "HIBERNATION from fd %d (IMEI %s)", sk, imei
 191                             )
 192                             tostop.append(sk)
 193                         elif proto_of_message(msg) == LOGIN.PROTO:
 194                             clients.response(Resp(imei=imei, payload=LOGIN.response()))
 195             # poll queue consumed, make changes now
 196             for fd in tostop:
 197                 poller.unregister(fd)
 198                 clients.stop(fd)
 199             for zmsg in tosend:
 200                 clients.response(zmsg)
 201             for clntsock, clntaddr in topoll:
 202                 fd = clients.add(clntsock, clntaddr)
 203                 poller.register(fd)
 204     except KeyboardInterrupt:
 205         pass
 206
 207
 208 if __name__.endswith("__main__"):
 209     runserver(common.init(log))