Добавил методы сравнения в структуры, и класс ShapeItem для инфы из NAS

agent/structs.py

@@ -1,6 +1,6 @@
 # -*- coding: utf8 -*-
 from abc import ABCMeta, abstractmethod
-from struct import pack, unpack
+from struct import pack, unpack, calcsize
 from .utils import int2ip, ip2int
 
 
@@ -13,14 +13,26 @@ class BaseStruct(object, metaclass=ABCMeta):
     def deserialize(self, data, *args):
         """создаём объект из бинарной строки"""
 
+    def __ne__(self, other):
+        return not self == other
 
-class IpStruct(object):
+
+class IpStruct(BaseStruct):
 
     def __init__(self, ip):
         if type(ip) is int:
             self.__ip = ip
         else:
-            self.__ip = ip2int(ip)
+            self.__ip = ip2int(str(ip))
+
+    def serialize(self):
+        dt = pack("!I", int(self.__ip))
+        return dt
+
+    def deserialize(self, data, *args):
+        dt = unpack("!I", data)
+        self.__ip = int(dt[0])
+        return self
 
     def get_str(self):
         return int2ip(self.__ip)
@@ -28,14 +40,21 @@ class IpStruct(object):
     def get_int(self):
         return self.__ip
 
+    def __eq__(self, other):
+        assert isinstance(other, IpStruct)
+        return self.__ip == other.__ip
+
+    def __str__(self):
+        return int2ip(self.__ip)
+
 
 # Как обслуживается абонент
 class TariffStruct(BaseStruct):
 
     def __init__(self, tariff_id=0, speedIn=None, speedOut=None):
         self.tid = tariff_id
-        self.speedIn = float(speedIn) or 0.0625
-        self.speedOut = float(speedOut) or 0.0625
+        self.speedIn = float(speedIn or 0.0625)
+        self.speedOut = float(speedOut or 0.0625)
 
     def serialize(self):
         dt = pack("!Iff", int(self.tid), float(self.speedIn), float(self.speedOut))
@@ -48,6 +67,15 @@ class TariffStruct(BaseStruct):
         self.speedOut = float(dt[2])
         return self
 
+    def __eq__(self, other):
+        assert isinstance(other, TariffStruct)
+        # не сравниваем id, т.к. тарифы с одинаковыми скоростями для NAS одинаковы
+        # Да и иногда не удобно доставать из nas id тарифы из базы
+        return self.speedIn == other.speedIn and self.speedOut == other.speedOut
+
+    def __str__(self):
+        return "Id=%d, speedIn=%.2f, speedOut=%.2f" % (self.tid, self.speedIn, self.speedOut)
+
 
 # Абонент из базы
 class AbonStruct(BaseStruct):
@@ -64,13 +92,43 @@ class AbonStruct(BaseStruct):
         dt = pack("!LII", self.uid, self.ip.get_int(), self.tariff.tid)
         return dt
 
-    def deserialize(self, data, all_tarifs=None):
+    def deserialize(self, data, tariff=None):
         dt = unpack("!LII", data)
         self.uid = dt[0]
         self.ip = IpStruct(dt[1])
-        tarifs = [trf for trf in all_tarifs if trf.tid == dt[2]]
-        if len(tarifs) < 1:
-            raise IndexError
-        assert isinstance(tarifs[0], TariffStruct)
-        self.tariff = tarifs[0]
+        if tariff is not None:
+            assert isinstance(tariff, TariffStruct)
+            self.tariff = tariff
+        return self
+
+    def __eq__(self, other):
+        assert isinstance(other, AbonStruct)
+        r = self.uid == other.uid and self.ip == other.ip
+        r = r and self.tariff == other.tariff
+        return r
+
+    def __str__(self):
+        return "uid=%d, ip=%s, tariff=%s" % (self.uid, self.ip, self.tariff)
+
+
+# Правило шейпинга в фаере, или ещё можно сказать услуга абонента на NAS
+class ShapeItem(BaseStruct):
+    def __init__(self, abon, sid):
+        self.abon = abon
+        self.sid = sid
+
+    def serialize(self):
+        abon_pack = self.abon.serialize()
+        dt = pack('!L', self.sid)
+        return dt+abon_pack
+
+    def deserialize(self, data, *args):
+        sz = calcsize('!L')
+        dt = unpack('!L', data[:sz])
+        self.sid = dt
+        self.abon.deserialize(data[sz:])
         return self
+
+    def __eq__(self, other):
+        assert isinstance(other, ShapeItem)
+        return self.sid == other.sid and self.abon == other.abon