Merge branch 'parallel-for-agent' into 'master'

Parallel for agent

See merge request ziirish/burp-ui!103

burpui/decorators.py

@@ -11,8 +11,6 @@ import time
 import datetime
 
 from functools import wraps
-from wsgiref.handlers import format_date_time
-from flask_restplus.utils import unpack
 
 
 def browser_cache(expires=None):
@@ -31,6 +29,9 @@ def browser_cache(expires=None):
             return render_template('index.html')
 
     """
+    from wsgiref.handlers import format_date_time
+    from flask_restplus.utils import unpack
+
     def cache_decorator(view):
         @wraps(view)
         def cache_func(*args, **kwargs):
@@ -64,5 +65,25 @@ def implement(func):
     to indicate we don't want the default "magic" implementation and use the
     custom implementation instead.
     """
-    func.__ismethodimplemented__ = True
+    try:
+        func.__ismethodimplemented__ = True
+    except AttributeError:
+        # properties seem immutable
+        pass
+    return func
+
+
+def usetriorun(func):
+    """A decorator indicating the method uses trio.run
+
+    Such functions should always be written like this:
+
+    ::
+        def function(self, *args, **kwargs):
+            return trio.run(self._async_function, partial(*args, **kwargs))
+    """
+    try:
+        func.__isusingtriorun__ = True
+    except AttributeError:
+        pass
     return func

burpui/engines/agent.py

@@ -15,8 +15,6 @@ import json
 import logging
 import trio
 
-from logging.handlers import RotatingFileHandler
-
 from ..exceptions import BUIserverException
 from ..misc.backend.interface import BUIbackend
 from .._compat import pickle, to_bytes, to_unicode
@@ -53,20 +51,24 @@ class BurpHandler(BUIbackend):
     BUIbackend.__abstractmethods__ = frozenset()
 
     def __init__(self, backend='burp2', logger=None, conf=None):
-        self.backend = backend
+        self.backend_name = backend
+        self.is_async = backend == 'parallel'
         self.logger = logger
 
         top = __name__
-        if '.' in self.backend:
-            module = self.backend
+        if '.' in self.backend_name:
+            module = self.backend_name
         else:
             if '.' in top:
                 top = top.split('.')[0]
-            module = '{0}.misc.backend.{1}'.format(top, self.backend)
+            module = '{0}.misc.backend.{1}'.format(top, self.backend_name)
         try:
             sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
             mod = __import__(module, fromlist=['Burp'])
-            Client = mod.Burp
+            if self.is_async:
+                Client = mod.AsyncBurp
+            else:
+                Client = mod.Burp
             self.backend = Client(conf=conf)
             stats = self.backend.statistics()
             if 'alive' not in stats or not stats['alive']:
@@ -78,7 +80,7 @@ class BurpHandler(BUIbackend):
     def __getattribute__(self, name):
         # always return this value because we need it and if we don't do that
         # we'll end up with an infinite loop
-        if name == 'foreign' or name == 'backend':
+        if name in ['foreign', 'backend', 'logger', 'is_async']:
             return object.__getattribute__(self, name)
         # now we can retrieve the 'foreign' list and know if the object called
         # is in the backend
@@ -97,36 +99,10 @@ class BUIAgent(BUIbackend):
     def __init__(self, conf=None, level=0, logfile=None):
         self.padding = 1
         level = level or 0
-        if level > logging.NOTSET:
-            levels = [
-                logging.CRITICAL,
-                logging.ERROR,
-                logging.WARNING,
-                logging.INFO,
-                logging.DEBUG,
-            ]
-            if level >= len(levels):
-                level = len(levels) - 1
-            lvl = levels[level]
-            self.logger.setLevel(lvl)
-            if lvl > logging.DEBUG:
-                LOG_FORMAT = '[%(asctime)s] %(levelname)s in %(module)s.%(funcName)s: %(message)s'
-            else:
-                LOG_FORMAT = (
-                    '-' * 80 + '\n' +
-                    '%(levelname)s in %(module)s.%(funcName)s [%(pathname)s:%(lineno)d]:\n' +
-                    '%(message)s\n' +
-                    '-' * 80
-                )
-            if logfile:
-                handler = RotatingFileHandler(logfile, maxBytes=1024 * 1024 * 100, backupCount=20)
-            else:
-                handler = logging.StreamHandler()
-            handler.setLevel(lvl)
-            handler.setFormatter(logging.Formatter(LOG_FORMAT))
-            self.logger.addHandler(handler)
-            self.logger.info('conf: {}'.format(conf))
-            self.logger.info('level: {}'.format(logging.getLevelName(lvl)))
+        self.logger.init_logger(config=dict(level=level, logfile=logfile))
+        lvl = self.logger.getEffectiveLevel()
+        self.logger.info('conf: {}'.format(conf))
+        self.logger.info('level: {}'.format(logging.getLevelName(lvl)))
         if not conf:
             raise IOError('No configuration file found')
 
@@ -194,7 +170,10 @@ class BUIAgent(BUIbackend):
                 elif j['func'] == 'agent_version':
                     res = json.dumps(__version__)
                 elif j['func'] == 'restore_files':
-                    res, err = getattr(self.client, j['func'])(**j['args'])
+                    if self.client.is_async:
+                        res, err = await getattr(self.client, j['func'])(**j['args'])
+                    else:
+                        res, err = getattr(self.client, j['func'])(**j['args'])
                     if err:
                         await server_stream.send_all(b'ER')
                         await server_stream.send_all(struct.pack('!Q', len(err)))
@@ -258,6 +237,7 @@ class BUIAgent(BUIbackend):
                         os.unlink(path)
                         res = json.dumps(True)
                 else:
+                    callback = getattr(self.client, j['func'])
                     if j['args']:
                         if 'pickled' in j and j['pickled']:
                             # de-serialize arguments if needed
@@ -280,9 +260,16 @@ class BUIAgent(BUIbackend):
                             data = b64decode(pickles)
                             data = data.replace(b'burpui.datastructures', to_bytes(f'{mod}.datastructures'))
                             j['args'] = pickle.loads(data)
-                        res = json.dumps(getattr(self.client, j['func'])(**j['args']))
+
+                        if self.client.is_async:
+                            res = json.dumps(await callback(**j['args']))
+                        else:
+                            res = json.dumps(callback(**j['args']))
                     else:
-                        res = json.dumps(getattr(self.client, j['func'])())
+                        if self.client.is_async:
+                            res = json.dumps(await callback())
+                        else:
+                            res = json.dumps(callback())
                 self.logger.info(f'result: {res}')
                 await server_stream.send_all(b'OK')
             except (BUIserverException, Exception) as exc:

burpui/misc/backend/burp1.py

@@ -541,9 +541,7 @@ class Burp(BUIbackend):
 
     def get_clients_report(self, clients, agent=None):
         """See :func:`burpui.misc.backend.interface.BUIbackend.get_clients_report`"""
-        ret = {}
-        cls = []
-        bkp = []
+        data = []
         for cli in clients:
             if not cli:
                 continue
@@ -551,19 +549,27 @@ class Burp(BUIbackend):
             if not client or not client[-1]:
                 continue
             stats = self.get_backup_logs(client[-1]['number'], cli['name'])
+            data.append((cli, client, stats))
+        return self._do_get_clients_report(data)
+
+    def _do_get_clients_report(self, data):
+        ret = {}
+        cls = []
+        bkp = []
+        for client, backups, stats in data:
             os = stats['os'] if 'os' in stats else "unknown"
             totsize = stats['totsize'] if 'totsize' in stats else 0
             total = stats['total']['total'] if \
                 'total' in stats and 'total' in stats['total'] else 0
             cls.append({
-                'name': cli['name'],
+                'name': client['name'],
                 'stats': {
                     'os': os,
                     'totsize': totsize,
                     'total': total
                 }
             })
-            bkp.append({'name': cli['name'], 'number': len(client)})
+            bkp.append({'name': client['name'], 'number': len(backups)})
         ret = {'clients': cls, 'backups': bkp}
         return ret
 

burpui/misc/backend/burp2.py

@@ -376,7 +376,11 @@ class Burp(Burp1):
         except KeyError:
             self.logger.warning('Client not found')
             return ret
+        return self._do_get_counters(client)
 
+    def _do_get_counters(self, data):
+        ret = {}
+        client = data
         # check the client is currently backing-up
         if 'run_status' not in client or client['run_status'] != 'running':
             return ret
@@ -486,13 +490,14 @@ class Burp(Burp1):
             query = self.status('c:{0}\n'.format(name))
         except BUIserverException:
             return False
-        if not query:
-            return False
-        try:
-            return query['clients'][0]['run_status'] in ['running']
-        except KeyError:
-            self.logger.warning('Client not found')
-            return False
+        return self._do_is_backup_running(query)
+
+    def _do_is_backup_running(self, data):
+        if data:
+            try:
+                return data['clients'][0]['run_status'] in ['running']
+            except KeyError:
+                pass
         return False
 
     def is_one_backup_running(self, agent=None):
@@ -504,7 +509,11 @@ class Burp(Burp1):
             clients = self.get_all_clients()
         except BUIserverException:
             return ret
-        for client in clients:
+        return self._do_is_one_backup_running(clients)
+
+    def _do_is_one_backup_running(self, data):
+        ret = []
+        for client in data:
             if client['state'] in ['running']:
                 ret.append(client['name'])
         return ret
@@ -631,6 +640,11 @@ class Burp(Burp1):
         except (KeyError, IndexError):
             self.logger.warning('Client not found')
             return ret
+        return self._do_get_client_status(client)
+
+    def _do_get_client_status(self, data):
+        ret = {}
+        client = data
         ret['state'] = self._status_human_readable(client['run_status'])
         infos = client['backups']
         if ret['state'] in ['running']:
@@ -641,7 +655,7 @@ class Burp(Burp1):
                     if 'action' in child and child['action'] == 'backup':
                         ret['phase'] = child['phase']
                         break
-            counters = self.get_counters(name)
+            counters = self._do_get_counters(client)
             if 'percent' in counters:
                 ret['percent'] = counters['percent']
             else:
@@ -736,6 +750,10 @@ class Burp(Burp1):
         query = self.status('c:{0}:b:{1}\n'.format(name, backup))
         if not query:
             return False
+        return self._do_is_backup_deletable(query)
+
+    def _do_is_backup_deletable(self, data):
+        query = data
         try:
             flags = query['clients'][0]['backups'][0]['flags']
             return 'deletable' in flags

burpui/misc/backend/interface.py

@@ -9,10 +9,11 @@
 """
 import os
 import re
-import logging
 
 from abc import ABCMeta, abstractmethod
 
+from ...tools.logging import logger
+
 G_BURPPORT = 4972
 G_BURPHOST = '::1'
 G_BURPBIN = '/usr/sbin/burp'
@@ -44,7 +45,7 @@ class BUIbackend(object, metaclass=ABCMeta):
     # Defaults config parameters
     defaults = {}
 
-    logger = logging.getLogger('burp-ui')
+    logger = logger
 
     def __init__(self, server=None, conf=None):  # pragma: no cover
         """
@@ -1143,3 +1144,6 @@ class BUIbackend(object, metaclass=ABCMeta):
         returns the version of the given agent.
         """
         raise NotImplementedError("Sorry, the current Backend does not implement this method!")  # pragma: no cover
+
+
+BUIBACKEND_INTERFACE_METHODS = BUIbackend.__abstractmethods__.copy()

burpui/misc/backend/utils/burp2.py

@@ -11,13 +11,13 @@ import re
 import subprocess
 import json
 import datetime
-import logging
 
 from select import select
 
 from ...._compat import to_bytes, to_unicode
 from ....exceptions import BUIserverException
 from ....security import sanitize_string
+from ....tools.logging import logger
 from .constant import BURP_LIST_BATCH, BURP_MINIMAL_VERSION, BURP_STATUS_FORMAT_V2, \
     BURP_STATUS_DELIMITER
 
@@ -39,7 +39,7 @@ class Monitor(object):
     :param timeout: Time to wait for an answer
     :type timeout: int
     """
-    logger = logging.getLogger('burp-ui')
+    logger = logger
 
     # cache status results
     _status_cache = {}

pkgs/burp-ui-agent/burpui_agent-decoy/decorators.py

+../../../burpui/decorators.py
\ No newline at end of file