X-Git-Url: http://average.org/gitweb/?a=blobdiff_plain;f=gps303%2Fzmsg.py;h=2d497c0cd5a98997bb3aa89c101a39fb29213c54;hb=18eda7307e92eeee6a9e2fdd2e810f98d25df654;hp=8179616485f0d9f5b89148e4dcf12d21df4294ae;hpb=5e1e7a4d37a1e149d5e899dada7b55a863cd8e64;p=loctrkd.git

diff --git a/gps303/zmsg.py b/gps303/zmsg.py
index 8179616..2d497c0 100644
--- a/gps303/zmsg.py
+++ b/gps303/zmsg.py
@@ -2,36 +2,47 @@
 
 import ipaddress as ip
 from struct import pack, unpack
+from typing import Any, cast, Optional, Tuple, Type, Union
 
 __all__ = "Bcast", "Resp", "topic"
 
 
-def pack_peer(peeraddr):
-    try:
-        if peeraddr is None:
-            saddr = "::"
-            port = 0
-        else:
-            saddr, port, _x, _y = peeraddr
-        addr = ip.ip_address(saddr)
-    except ValueError:
+def pack_peer(
+    peeraddr: Union[None, Tuple[str, int], Tuple[str, int, Any, Any]]
+) -> bytes:
+    if peeraddr is None:
+        addr: Union[ip.IPv4Address, ip.IPv6Address] = ip.IPv6Address(0)
+        port = 0
+    elif len(peeraddr) == 2:
+        peeraddr = cast(Tuple[str, int], peeraddr)
         saddr, port = peeraddr
-        a4 = ip.ip_address(saddr)
-        addr = ip.IPv6Address(b"\0\0\0\0\0\0\0\0\0\0\xff\xff" + a4.packed)
+        addr = ip.ip_address(saddr)
+    elif len(peeraddr) == 4:
+        peeraddr = cast(Tuple[str, int, Any, Any], peeraddr)
+        saddr, port, _x, _y = peeraddr
+        addr = ip.ip_address(saddr)
+    if isinstance(addr, ip.IPv4Address):
+        addr = ip.IPv6Address(b"\0\0\0\0\0\0\0\0\0\0\xff\xff" + addr.packed)
     return addr.packed + pack("!H", port)
 
 
-def unpack_peer(buffer):
+def unpack_peer(
+    buffer: bytes,
+) -> Tuple[str, int]:
     a6 = ip.IPv6Address(buffer[:16])
     port = unpack("!H", buffer[16:])[0]
-    addr = a6.ipv4_mapped
-    if addr is None:
-        addr = a6
-    return (addr, port)
+    a4 = a6.ipv4_mapped
+    if a4 is not None:
+        return (str(a4), port)
+    elif a6 == ip.IPv6Address("::"):
+        return ("", 0)
+    return (str(a6), port)
 
 
 class _Zmsg:
-    def __init__(self, *args, **kwargs):
+    KWARGS: Tuple[Tuple[str, Any], ...]
+
+    def __init__(self, *args: Any, **kwargs: Any) -> None:
         if len(args) == 1:
             self.decode(args[0])
         elif bool(kwargs):
@@ -46,7 +57,7 @@ class _Zmsg:
                 + str(kwargs)
             )
 
-    def __repr__(self):
+    def __repr__(self) -> str:
         return "{}({})".format(
             self.__class__.__name__,
             ", ".join(
@@ -62,26 +73,30 @@ class _Zmsg:
             ),
         )
 
-    def __eq__(self, other):
-        return all(
-            [getattr(self, k) == getattr(other, k) for k, _ in self.KWARGS]
-        )
+    def __eq__(self, other: object) -> bool:
+        if isinstance(other, self.__class__):
+            return all(
+                [getattr(self, k) == getattr(other, k) for k, _ in self.KWARGS]
+            )
+        return NotImplemented
 
-    def decode(self, buffer):
+    def decode(self, buffer: bytes) -> None:
         raise NotImplementedError(
             self.__class__.__name__ + "must implement `decode()` method"
         )
 
     @property
-    def packed(self):
+    def packed(self) -> bytes:
         raise NotImplementedError(
             self.__class__.__name__ + "must implement `packed()` property"
         )
 
 
-def topic(proto, is_incoming=True, imei=None):
-    return (
-        pack("BB", is_incoming, proto) + b"" if imei is None else imei.encode()
+def topic(
+    proto: int, is_incoming: bool = True, imei: Optional[str] = None
+) -> bytes:
+    return pack("BB", is_incoming, proto) + (
+        b"" if imei is None else pack("16s", imei.encode())
     )
 
 
@@ -98,10 +113,16 @@ class Bcast(_Zmsg):
     )
 
     @property
-    def packed(self):
+    def packed(self) -> bytes:
         return (
-            pack("BB", int(self.is_incoming), self.proto)
-            + ("0000000000000000" if self.imei is None else self.imei).encode()
+            pack(
+                "BB16s",
+                int(self.is_incoming),
+                self.proto,
+                b"0000000000000000"
+                if self.imei is None
+                else self.imei.encode(),
+            )
             + (
                 b"\0\0\0\0\0\0\0\0"
                 if self.when is None
@@ -111,10 +132,10 @@ class Bcast(_Zmsg):
             + self.packet
         )
 
-    def decode(self, buffer):
+    def decode(self, buffer: bytes) -> None:
         self.is_incoming = bool(buffer[0])
         self.proto = buffer[1]
-        self.imei = buffer[2:18].decode()
+        self.imei: Optional[str] = buffer[2:18].decode()
         if self.imei == "0000000000000000":
             self.imei = None
         self.when = unpack("!d", buffer[18:26])[0]
@@ -128,16 +149,23 @@ class Resp(_Zmsg):
     KWARGS = (("imei", None), ("when", None), ("packet", b""))
 
     @property
-    def packed(self):
+    def packed(self) -> bytes:
         return (
-            "0000000000000000" if self.imei is None else self.imei.encode()
-        ) + (
+            pack(
+                "16s",
+                "0000000000000000"
+                if self.imei is None
+                else self.imei.encode(),
+            )
+            + (
                 b"\0\0\0\0\0\0\0\0"
                 if self.when is None
                 else pack("!d", self.when)
-            ) + self.packet
+            )
+            + self.packet
+        )
 
-    def decode(self, buffer):
+    def decode(self, buffer: bytes) -> None:
         self.imei = buffer[:16].decode()
         self.when = unpack("!d", buffer[16:24])[0]
         self.packet = buffer[24:]