Перевёл метод чтения списков на генераторы (ленивые вычисления). И пофиксил пару багов

agent/mod_mikrotik.py

@@ -4,8 +4,9 @@ import binascii
 from hashlib import md5
 from .core import BaseTransmitter, NasFailedResult, NasNetworkError
 from mydefs import ping
-from .structs import TariffStruct, IpStruct
+from .structs import TariffStruct, AbonStruct, IpStruct
 from . import settings
+from djing.settings import DEBUG
 
 
 class ApiRos:
@@ -15,18 +16,17 @@ class ApiRos:
         self.currenttag = 0
 
     def login(self, username, pwd):
-        for repl, attrs in self.talk(["/login"]):
+        for repl, attrs in self.talk_iter(["/login"]):
             chal = binascii.unhexlify(attrs['=ret'])
         md = md5()
         md.update(b'\x00')
         md.update(bytes(pwd, 'utf-8'))
         md.update(chal)
-        self.talk(["/login", "=name=" + username,
-                   "=response=00" + binascii.hexlify(md.digest()).decode('utf-8')])
+        for r in self.talk_iter(["/login", "=name=" + username,
+                   "=response=00" + binascii.hexlify(md.digest()).decode('utf-8')]): pass
 
-    def talk(self, words):
+    def talk_iter(self, words):
         if self.writeSentence(words) == 0: return
-        r = []
         while 1:
             i = self.readSentence()
             if len(i) == 0: continue
@@ -38,8 +38,8 @@ class ApiRos:
                     attrs[w] = ''
                 else:
                     attrs[w[:j]] = w[j+1:]
-            r.append((reply, attrs))
-            if reply == '!done': return r
+            yield (reply, attrs)
+            if reply == '!done': return
 
     def writeSentence(self, words):
         ret = 0
@@ -57,14 +57,16 @@ class ApiRos:
             r.append(w)
 
     def writeWord(self, w):
-        print("<<< " + w)
+        if DEBUG:
+            print("<<< " + w)
         b = bytes(w, "utf-8")
         self.writeLen(len(b))
         self.writeBytes(b)
 
     def readWord(self):
         ret = self.readBytes(self.readLen()).decode('utf-8')
-        print(">>> " + ret)
+        if DEBUG:
+            print(">>> " + ret)
         return ret
 
     def writeLen(self, l):
@@ -130,7 +132,6 @@ class ApiRos:
         return ret
 
 
-# TODO: Реализовать передачу в шейпер срок действия тарифа
 class MikrotikTransmitter(BaseTransmitter):
     def __init__(self, login=None, password=None, ip=None, port=None):
         ip = ip or settings.NAS_IP
@@ -144,25 +145,36 @@ class MikrotikTransmitter(BaseTransmitter):
         except ConnectionRefusedError:
             raise NasNetworkError('Подключение к %s отклонено (Connection Refused)' % ip)
 
+    def _exec_cmd_iter(self, cmd):
+        assert isinstance(cmd, list)
+        result_iter = self.ar.talk_iter(cmd)
+        for rt in result_iter:
+            if rt[0] == '!trap':
+                raise NasFailedResult(rt[1]['=message'])
+            yield rt
+
     def _exec_cmd(self, cmd):
         assert isinstance(cmd, list)
-        result = self.ar.talk(cmd)
-        for rt in result:
+        result_iter = self.ar.talk_iter(cmd)
+        res = []
+        for rt in result_iter:
             if rt[0] == '!trap':
                 raise NasFailedResult(rt[1]['=message'])
-        return result
+            res.append(rt[1])
+        return res
 
     # ищем правило по имени, и возвращаем всю инфу о найденном правиле
     def _find_queue(self, name):
         ret = self._exec_cmd(['/queue/simple/print', '?name=%s' % name])
-        return ret[0][1]
+        return ret[0]
 
     def add_user_range(self, user_list):
-        return list(map(self.add_user, user_list))
+        for usr in user_list:
+            self.add_user(usr)
 
-    def remove_user_range(self, user_list):
-        names = ['uid%d' % usr.uid for usr in user_list]
-        return self._exec_cmd(['/queue/simple/remove', '=.id=%s' % ','.join(names)])
+    def remove_user_range(self, user_ids):
+        names = ['%d' % usr for usr in user_ids]
+        return self._exec_cmd(['/queue/simple/remove', *names])
 
     # добавляем правило шейпинга для указанного ip и со скоростью max-limit=Upload/Download
     # Мы уверены что user это инстанс класса agent.structs.AbonStruct
@@ -178,41 +190,63 @@ class MikrotikTransmitter(BaseTransmitter):
     # удаляем правило шейпера по имени правила
     def remove_user(self, user):
         uid = user if type(user) is int else user.uid
-        self._exec_cmd(['/queue/simple/remove', '=.id=uid%d' % uid])
+        return self._exec_cmd(['/queue/simple/remove', '=name=uid%d' % uid])
 
     # обновляем основную инфу абонента
     def update_user(self, user):
         assert isinstance(user.tariff, TariffStruct)
         assert isinstance(user.ip, IpStruct)
-        self._exec_cmd(['/queue/simple/set', '=.id=uid%d' % user.uid,
+        return self._exec_cmd(['/queue/simple/set', '=name=uid%d' % user.uid,
             '=max-limit=%fM/%fM' % (user.tariff.speedOut, user.tariff.speedIn),
             '=target-addresses=%s/32' % user.ip.get_str()
         ])
 
+    # читаем абонентов, возващаем абнента и номер в микротике
+    def read_users_iter(self):
+        ret_it = self._exec_cmd_iter(['/queue/simple/print', '=detail'])
+        for re in ret_it:
+            if re[0] == '!done': return
+            speeds = re[1]['=limit-at'].split('/')
+            speeds = [sp.replace('M','') for sp in speeds]
+            abon = AbonStruct(
+                uid=int(re[1]['=name'][3:]),
+                ip=IpStruct(re[1]['=target-addresses'][:-3]),
+                tariff=TariffStruct(speedIn=speeds[0], speedOut=speeds[1])
+            )
+            yield abon
+
+    # то же что и выше, только получаем номера в микротике
+    def read_users_mikroids_iter(self):
+        ret_it = self._exec_cmd_iter(['/queue/simple/print', '=detail'])
+        for re in ret_it:
+            if re[0] == '!done': return
+            yield int(re[1]['=.id'].replace('*', ''), base=16)
+
     # приостановливаем обслуживание абонента
+    # в @user передаём номер в микротике
     def pause_user(self, user):
-        self._exec_cmd(['/queue/simple/disable', '=.id=uid%d' % user.uid])
+        self._exec_cmd(['/queue/simple/disable', user])
 
     # продолжаем обслуживание абонента
+    # в @user передаём номер в микротике
     def start_user(self, user):
-        self._exec_cmd(['/queue/simple/enable', '=.id=uid%d' % user.uid])
+        self._exec_cmd(['/queue/simple/enable', user])
 
     # Тарифы хранить нам не надо, так что методы тарифов ниже не реализуем
     def add_tariff_range(self, tariff_list):
         pass
 
-    # todo: реальзовать
+    # соответственно и удалять тарифы не надо
     def remove_tariff_range(self, tariff_list):
         pass
 
-    # todo: реальзовать
+    # и добавлять тоже
     def add_tariff(self, tariff):
         pass
 
-    # todo: реальзовать
+    # и обновлять
     def update_tariff(self, tariff):
         pass
 
-    # todo: реальзовать
     def remove_tariff(self, tid):
         pass