loctrkd/zmsg.py

   1 """ Zeromq messages """
   2
   3 import ipaddress as ip
   4 from struct import pack, unpack
   5 from typing import Any, cast, Optional, Tuple, Type, Union
   6
   7 __all__ = "Bcast", "Resp", "topic", "rtopic"
   8
   9
  10 def pack_peer(  # 18 bytes
  11     peeraddr: Union[None, Tuple[str, int], Tuple[str, int, Any, Any]]
  12 ) -> bytes:
  13     if peeraddr is None:
  14         addr: Union[ip.IPv4Address, ip.IPv6Address] = ip.IPv6Address(0)
  15         port = 0
  16     elif len(peeraddr) == 2:
  17         peeraddr = cast(Tuple[str, int], peeraddr)
  18         saddr, port = peeraddr
  19         addr = ip.ip_address(saddr)
  20     elif len(peeraddr) == 4:
  21         peeraddr = cast(Tuple[str, int, Any, Any], peeraddr)
  22         saddr, port, _x, _y = peeraddr
  23         addr = ip.ip_address(saddr)
  24     if isinstance(addr, ip.IPv4Address):
  25         addr = ip.IPv6Address(b"\0\0\0\0\0\0\0\0\0\0\xff\xff" + addr.packed)
  26     return addr.packed + pack("!H", port)
  27
  28
  29 def unpack_peer(
  30     buffer: bytes,
  31 ) -> Tuple[str, int]:
  32     a6 = ip.IPv6Address(buffer[:16])
  33     port = unpack("!H", buffer[16:])[0]
  34     a4 = a6.ipv4_mapped
  35     if a4 is not None:
  36         return (str(a4), port)
  37     elif a6 == ip.IPv6Address("::"):
  38         return ("", 0)
  39     return (str(a6), port)
  40
  41
  42 class _Zmsg:
  43     KWARGS: Tuple[Tuple[str, Any], ...]
  44
  45     def __init__(self, *args: Any, **kwargs: Any) -> None:
  46         if len(args) == 1:
  47             self.decode(args[0])
  48         elif bool(kwargs):
  49             for k, v in self.KWARGS:
  50                 setattr(self, k, kwargs.get(k, v))
  51         else:
  52             raise RuntimeError(
  53                 self.__class__.__name__
  54                 + ": both args "
  55                 + str(args)
  56                 + " and kwargs "
  57                 + str(kwargs)
  58             )
  59
  60     def __repr__(self) -> str:
  61         return "{}({})".format(
  62             self.__class__.__name__,
  63             ", ".join(
  64                 [
  65                     "{}={}".format(
  66                         k,
  67                         'bytes.fromhex("{}")'.format(getattr(self, k).hex())
  68                         if isinstance(getattr(self, k), bytes)
  69                         else getattr(self, k),
  70                     )
  71                     for k, _ in self.KWARGS
  72                 ]
  73             ),
  74         )
  75
  76     def __eq__(self, other: object) -> bool:
  77         if isinstance(other, self.__class__):
  78             return all(
  79                 [getattr(self, k) == getattr(other, k) for k, _ in self.KWARGS]
  80             )
  81         return NotImplemented
  82
  83     def decode(self, buffer: bytes) -> None:
  84         raise NotImplementedError(
  85             self.__class__.__name__ + "must implement `decode()` method"
  86         )
  87
  88     @property
  89     def packed(self) -> bytes:
  90         raise NotImplementedError(
  91             self.__class__.__name__ + "must implement `packed()` property"
  92         )
  93
  94
  95 def topic(
  96     proto: str, is_incoming: bool = True, imei: Optional[str] = None
  97 ) -> bytes:
  98     return pack("B16s", is_incoming, proto.encode()) + (
  99         b"" if imei is None else pack("16s", imei.encode())
 100     )
 101
 102
 103 def rtopic(imei: str) -> bytes:
 104     return pack("16s", imei.encode())
 105
 106
 107 class Bcast(_Zmsg):
 108     """Zmq message to broadcast what was received from the terminal"""
 109
 110     KWARGS = (
 111         ("is_incoming", True),
 112         ("proto", "UNKNOWN"),
 113         ("imei", None),
 114         ("when", None),
 115         ("peeraddr", None),
 116         ("packet", b""),
 117     )
 118
 119     @property
 120     def packed(self) -> bytes:
 121         return (
 122             pack(
 123                 "!B16s16sd",
 124                 int(self.is_incoming),
 125                 self.proto[:16].ljust(16, "\0").encode(),
 126                 b"0000000000000000"
 127                 if self.imei is None
 128                 else self.imei.encode(),
 129                 0 if self.when is None else self.when,
 130             )
 131             + pack_peer(self.peeraddr)
 132             + self.packet
 133         )
 134
 135     def decode(self, buffer: bytes) -> None:
 136         is_incoming, proto, imei, when = unpack("!B16s16sd", buffer[:41])
 137         self.is_incoming = bool(is_incoming)
 138         self.proto = proto.decode().rstrip("\0")
 139         self.imei = (
 140             None if imei == b"0000000000000000" else imei.decode().strip("\0")
 141         )
 142         self.when = when
 143         self.peeraddr = unpack_peer(buffer[41:59])
 144         self.packet = buffer[59:]
 145
 146
 147 class Resp(_Zmsg):
 148     """Zmq message received from a third party to send to the terminal"""
 149
 150     KWARGS = (("imei", None), ("when", None), ("packet", b""))
 151
 152     @property
 153     def packed(self) -> bytes:
 154         return (
 155             pack(
 156                 "!16sd",
 157                 "0000000000000000"
 158                 if self.imei is None
 159                 else self.imei.encode(),
 160                 0 if self.when is None else self.when,
 161             )
 162             + self.packet
 163         )
 164
 165     def decode(self, buffer: bytes) -> None:
 166         imei, when = unpack("!16sd", buffer[:24])
 167         self.imei = (
 168             None if imei == b"0000000000000000" else imei.decode().strip("\0")
 169         )
 170
 171         self.when = when
 172         self.packet = buffer[24:]
 173
 174
 175 class Rept(_Zmsg):
 176     """Broadcast zmq message with "rectified" proto-agnostic json data"""
 177
 178     KWARGS = (("imei", None), ("payload", ""), ("pmod", None))
 179
 180     @property
 181     def packed(self) -> bytes:
 182         return (
 183             pack(
 184                 "16s16s",
 185                 b"0000000000000000"
 186                 if self.imei is None
 187                 else self.imei.encode(),
 188                 b"                "
 189                 if self.pmod is None
 190                 else self.pmod.encode(),
 191             )
 192             + self.payload.encode()
 193         )
 194
 195     def decode(self, buffer: bytes) -> None:
 196         imei = buffer[:16]
 197         self.imei = (
 198             None if imei == b"0000000000000000" else imei.decode().strip("\0")
 199         )
 200         pmod = buffer[16:32]
 201         self.pmod = (
 202             None if pmod == b"                " else pmod.decode().strip("\0")
 203         )
 204         self.payload = buffer[32:].decode()