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"
   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 class Bcast(_Zmsg):
 104     """Zmq message to broadcast what was received from the terminal"""
 105
 106     KWARGS = (
 107         ("is_incoming", True),
 108         ("proto", "UNKNOWN"),
 109         ("imei", None),
 110         ("when", None),
 111         ("peeraddr", None),
 112         ("packet", b""),
 113     )
 114
 115     @property
 116     def packed(self) -> bytes:
 117         return (
 118             pack(
 119                 "!B16s16sd",
 120                 int(self.is_incoming),
 121                 self.proto[:16].ljust(16, "\0").encode(),
 122                 b"0000000000000000"
 123                 if self.imei is None
 124                 else self.imei.encode(),
 125                 0 if self.when is None else self.when,
 126             )
 127             + pack_peer(self.peeraddr)
 128             + self.packet
 129         )
 130
 131     def decode(self, buffer: bytes) -> None:
 132         is_incoming, proto, imei, when = unpack("!B16s16sd", buffer[:41])
 133         self.is_incoming = bool(is_incoming)
 134         self.proto = proto.decode()
 135         self.imei = (
 136             None if imei == b"0000000000000000" else imei.decode().strip("\0")
 137         )
 138         self.when = when
 139         self.peeraddr = unpack_peer(buffer[41:59])
 140         self.packet = buffer[59:]
 141
 142
 143 class Resp(_Zmsg):
 144     """Zmq message received from a third party to send to the terminal"""
 145
 146     KWARGS = (("imei", None), ("when", None), ("packet", b""))
 147
 148     @property
 149     def packed(self) -> bytes:
 150         return (
 151             pack(
 152                 "!16sd",
 153                 "0000000000000000"
 154                 if self.imei is None
 155                 else self.imei.encode(),
 156                 0 if self.when is None else self.when,
 157             )
 158             + self.packet
 159         )
 160
 161     def decode(self, buffer: bytes) -> None:
 162         imei, when = unpack("!16sd", buffer[:24])
 163         self.imei = (
 164             None if imei == b"0000000000000000" else imei.decode().strip("\0")
 165         )
 166
 167         self.when = when
 168         self.packet = buffer[24:]