add: new bui-monitor tool

Handle a pool of burp client processes to have a more predictable amount
of burp client and allow some process parallelisation.

burpui/__main__.py

@@ -33,7 +33,7 @@ def parse_args(mode=True, name=None):
     parser.add_argument('-i', '--migrations', dest='migrations', help='migrations directory', metavar='<MIGRATIONSDIR>')
     parser.add_argument('remaining', nargs=REMAINDER)
     if mode:
-        parser.add_argument('-m', '--mode', dest='mode', help='application mode', metavar='<agent|server|celery|manage|legacy>')
+        parser.add_argument('-m', '--mode', dest='mode', help='application mode', metavar='<agent|server|celery|manage|monitor|legacy>')
 
     options, unknown = parser.parse_known_args()
     if mode and options.mode and options.mode not in ['celery', 'manage', 'server']:
@@ -65,6 +65,8 @@ def main():
         celery()
     elif options.mode == 'manage':
         manage()
+    elif options.mode == 'monitor':
+        monitor(options)
     elif options.mode == 'legacy':
         legacy(options, unknown)
     else:
@@ -131,10 +133,32 @@ def agent(options=None):
         conf = lookup_file(conf)
     check_config(conf)
 
-    agent = Agent(conf, options.log, options.logfile, options.debug)
+    agent = Agent(conf, options.log, options.logfile)
     trio.run(agent.run)
 
 
+def monitor(options=None):
+    import trio
+    from burpui.engines.monitor import MonitorPool
+    from burpui.utils import lookup_file
+    from burpui._compat import patch_json
+
+    patch_json()
+
+    if not options:
+        options, _ = parse_args(mode=False, name='bui-agent')
+
+    conf = ['buimonitor.cfg', 'buimonitor.sample.cfg']
+    if options.config:
+        conf = lookup_file(options.config, guess=False)
+    else:
+        conf = lookup_file(conf)
+    check_config(conf)
+
+    monitor = MonitorPool(conf, options.log, options.logfile)
+    trio.run(monitor.run)
+
+
 def celery():
     from burpui.utils import lookup_file
 

burpui/engines/agent.py

@@ -1,6 +1,6 @@
 # -*- coding: utf8 -*-
 """
-.. module:: burpui.agent
+.. module:: burpui.engines.agent
     :platform: Unix
     :synopsis: Burp-UI agent module.
 
@@ -39,7 +39,7 @@ BUI_DEFAULTS = {
         'sslcert': '',
         'sslkey': '',
         'backend': 'burp2',
-        'password': 'password',
+        'password': 'azerty',
     },
 }
 
@@ -91,8 +91,7 @@ class BurpHandler(BUIbackend):
 class BUIAgent(BUIbackend):
     BUIbackend.__abstractmethods__ = frozenset()
 
-    def __init__(self, conf=None, level=0, logfile=None, debug=False):
+    def __init__(self, conf=None, level=0, logfile=None):
-        self.debug = debug
         self.padding = 1
         level = level or 0
         if level > logging.NOTSET:
@@ -172,7 +171,7 @@ class BUIAgent(BUIbackend):
             if not lengthbuf:
                 return
             length, = struct.unpack('!Q', lengthbuf)
-            data = await server_stream.receive_some(length)
+            data = await self.receive_all(server_stream, length)
             self.logger.info(f'recv: {data!r}')
             txt = to_unicode(data)
             if txt == 'RE':
@@ -180,7 +179,7 @@ class BUIAgent(BUIbackend):
             j = json.loads(txt)
             if j['password'] != self.password:
                 self.logger.warning('-----> Wrong Password <-----')
-                await server_stream.send_all(b'ok')
+                await server_stream.send_all(b'KO')
                 return
             try:
                 if j['func'] == 'proxy_parser':
@@ -287,9 +286,7 @@ class BUIAgent(BUIbackend):
                 res = str(exc)
                 self.logger.error(res, exc_info=exc)
                 self.logger.warning(f'Forwarding Exception: {res}')
-                await server_stream.send_all(struct.pack('!Q', len(res)))
+
-                await server_stream.send_all(to_bytes(res))
-                return
             await server_stream.send_all(struct.pack('!Q', len(res)))
             await server_stream.send_all(to_bytes(res))
         except AttributeError as exc:
@@ -298,13 +295,12 @@ class BUIAgent(BUIbackend):
         except Exception as exc:
             self.logger.error(f'!!! {exc} !!!', exc_info=exc)
 
-    async def receive_all(self, stream, length=1024):
+    async def receive_all(self, stream: trio.StapledStream, length=1024, bsize=None):
         buf = b''
-        bsize = 1024
+        bsize = bsize if bsize is not None else 1024
+        bsize = min(bsize, length)
         received = 0
         tries = 0
-        if length < bsize:
-            bsize = length
         while received < length:
             newbuf = await stream.receive_some(bsize)
             if not newbuf:

burpui/engines/monitor.py

@@ -0,0 +1,208 @@
+# -*- coding: utf8 -*-
+"""
+.. module:: burpui.engines.monitor
+    :platform: Unix
+    :synopsis: Burp-UI monitor pool module.
+
+.. moduleauthor:: Ziirish <hi+burpui@ziirish.me>
+
+"""
+import ssl
+import trio
+import json
+import struct
+import logging
+
+from itertools import count
+from logging.handlers import RotatingFileHandler
+
+from ..exceptions import BUIserverException
+from ..misc.backend.burp.utils import Monitor
+from ..config import config
+from .._compat import to_bytes, to_unicode
+from ..desc import __version__
+
+
+CONNECTION_COUNTER = count()
+
+
+BUI_DEFAULTS = {
+    'Global': {
+        'port': 11111,
+        'bind': '::1',
+        'ssl': False,
+        'sslcert': '',
+        'sslkey': '',
+        'password': 'password123456',
+        'pool': 5,
+    },
+}
+
+
+class MonitorPool:
+    logger = logging.getLogger('burp-ui')  # type: logging.Logger
+
+    def __init__(self, conf=None, level=0, logfile=None, debug=False):
+        self.debug = debug
+        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(f'conf: {conf}')
+            self.logger.info('level: {}'.format(logging.getLevelName(lvl)))
+        if not conf:
+            raise IOError('No configuration file found')
+
+        # Raise exception if errors are encountered during parsing
+        self.conf = config
+        self.conf.parse(conf, True, BUI_DEFAULTS)
+        self.conf.default_section('Global')
+        self.port = self.conf.safe_get('port', 'integer')
+        self.bind = self.conf.safe_get('bind')
+        self.ssl = self.conf.safe_get('ssl', 'boolean')
+        self.sslcert = self.conf.safe_get('sslcert')
+        self.sslkey = self.conf.safe_get('sslkey')
+        self.password = self.conf.safe_get('password')
+        self.pool = self.conf.safe_get('pool', 'integer')
+
+        self.burpbin = self.conf.safe_get('burpbin', section='Burp')
+        self.bconfcli = self.conf.safe_get('bconfcli', section='Burp')
+        self.timeout = self.conf.safe_get('timeout', 'integer', section='Burp')
+
+        self.conf.setdefault('BUI_MONITOR', True)
+
+        self.monitor_pool = trio.Queue(self.pool)
+
+    def _ssl_context(self):
+        if not self.ssl:
+            return None
+        ctx = ssl.SSLContext()
+        ctx.load_cert_chain(self.sslcert, self.sslkey)
+        return ctx
+
+    async def receive_all(self, stream: trio.StapledStream, length=1024, bsize=None):
+        buf = b''
+        bsize = bsize if bsize is not None else 1024
+        bsize = min(bsize, length)
+        received = 0
+        tries = 0
+        while received < length:
+            newbuf = await stream.receive_some(bsize)
+            if not newbuf:
+                # 3 successive read failure => raise exception
+                if tries > 3:
+                    raise Exception('Unable to read full response')
+                tries += 1
+                await trio.sleep(0.1)
+                continue
+            # reset counter
+            tries = 0
+            buf += newbuf
+            received += len(newbuf)
+        return buf
+
+    async def handle(self, server_stream: trio.StapledStream):
+        ident = next(CONNECTION_COUNTER)
+        self.logger.info(f'{ident} - handle_request: started')
+        t0 = trio.current_time()
+        lengthbuf = await server_stream.receive_some(8)
+        if not lengthbuf:
+            return
+        length, = struct.unpack('!Q', lengthbuf)
+        data = await self.receive_all(server_stream, length)
+        self.logger.info(f'{ident} - recv: {data!r}')
+        txt = to_unicode(data)
+        if txt == 'RE':
+            return
+        req = json.loads(txt)
+        if req['password'] != self.password:
+            self.logger.warning(f'{ident} -----> Wrong Password <-----')
+            await server_stream.send_all(b'KO')
+            return
+        try:
+            if req.get('func') == 'monitor_version':
+                response = json.dumps(__version__)
+            else:
+                query = req['query']
+                self.logger.info(f'{ident} - Waiting for a monitor...')
+                t1 = trio.current_time()
+                mon = await self.monitor_pool.get()  # type: Monitor
+                t2 = trio.current_time()
+                t = t2 - t1
+                self.logger.info(f'{ident} - Waited {t:.3f}s')
+
+                response = mon.status(query, timeout=self.timeout, cache=req.get('cache', True))
+                response = json.dumps(response)
+                self.logger.info(f'{ident} - Releasing monitor')
+                await self.monitor_pool.put(mon)
+            self.logger.debug(f'{ident} - Sending: {response}')
+            await server_stream.send_all(b'OK')
+        except BUIserverException as exc:
+            await server_stream.send_all(b'ER')
+            response = str(exc)
+            self.logger.error(response, exc_info=exc)
+            self.logger.warning(f'Forwarding Exception: {response}')
+
+        await server_stream.send_all(struct.pack('!Q', len(response)))
+        await server_stream.send_all(to_bytes(response))
+
+        t3 = trio.current_time()
+        t = t3 - t0
+        self.logger.info(f'{ident} - Completed in {t:.3f}s')
+
+    async def launch_monitor(self, id):
+        self.logger.info(f'Starting client n°{id}')
+        mon = Monitor(self.burpbin, self.bconfcli, timeout=self.timeout, ident=id)
+        # warm up monitor
+        mon.status()
+        await self.monitor_pool.put(mon)
+
+    async def cleanup_monitor(self):
+        while not self.monitor_pool.empty():
+            self.logger.info('killing proc')
+            mon = await self.monitor_pool.get()  # noqa
+            del mon
+
+    async def run(self):
+        self.logger.info('Starting clients...')
+        async with trio.open_nursery() as nursery:
+            for i in range(self.pool):
+                nursery.start_soon(self.launch_monitor, i + 1)
+        self.logger.info(f'Ready to serve requests on {self.bind}:{self.port}')
+        try:
+            ctx = self._ssl_context()
+            if ctx:
+                await trio.serve_ssl_over_tcp(self.handle, self.port, ctx, host=self.bind)
+            else:
+                await trio.serve_tcp(self.handle, self.port, host=self.bind)
+        except KeyboardInterrupt:
+            pass
+
+        self.logger.info('Cleaning up')
+        async with trio.open_nursery() as nursery:
+            nursery.start_soon(self.cleanup_monitor)

burpui/misc/backend/burp/utils.py

@@ -48,7 +48,7 @@ class Monitor(object):
     _last_status_cleanup = datetime.datetime.now()
     _time_to_cache = datetime.timedelta(seconds=3)
 
-    def __init__(self, burpbin, burpconf, app=None, timeout=5):
+    def __init__(self, burpbin, burpconf, app=None, timeout=5, ident=None):
         """
         :param app: ``Burp-UI`` server instance in order to access logger
                        and/or some global settings
@@ -71,6 +71,7 @@ class Monitor(object):
         self.client_version = None
         self.server_version = None
         self.batch_list_supported = False
+        self.ident = ident or id(self)
 
         self._burp_client_ok = False
         version = ''
@@ -112,6 +113,7 @@ class Monitor(object):
 
     def _exit(self):
         """try not to leave child process server side"""
+        self.logger.debug(f'Exiting {self.ident}')
         self._terminate_burp()
         self._kill_burp()
 
@@ -248,7 +250,7 @@ class Monitor(object):
         try:
             timeout = timeout or self.timeout
             query = sanitize_string(query.rstrip())
-            self.logger.info(f"query: '{query}'")
+            self.logger.info(f"{self.ident} - query: '{query}'")
             query = '{0}\n'.format(query)
 
             self._cleanup_cache()

docs/buiagent.rst

@@ -126,8 +126,8 @@ Each option is commented, but here is a more detailed documentation:
   `Burp-UI versions <advanced_usage.html#versions>`__ for more details)
 - *password*: The shared secret between the `Burp-UI`_ server and `bui-agent`_.
 
-As with `Burp-UI`_, you need a specific section depending on the *version*
+As with `Burp-UI`_, you need a specific ``[Burp]`` section.
-value. Please refer to the `Burp-UI versions <advanced_usage.html#versions>`__
+Please refer to the `Burp-UI versions <advanced_usage.html#options>`__
 section for more details.
 
 Example

docs/buimonitor.rst

@@ -0,0 +1,149 @@
+bui-monitor
+===========
+
+The `bui-monitor`_ is a `Burp`_ client monitor processes pool.
+
+This pool only supports the `burp2`_ backend.
+
+The goal of this pool is to have a consistent amount of burp client processes
+related to your `Burp-UI`_ stack.
+
+Before this pool, you could have 1 process per `Burp-UI`_ instance (so if you
+use gunicorn with several workers, that would multiply the amount of processes),
+you also had 1 process per `celery`_ worker instance (which is one per CPU core
+available on your machine by default).
+In the end, it could be difficult to anticipate the resources to provision
+beforehand.
+Also, this wasn't very scalable.
+
+If you choose to use the `bui-monitor`_ pool with the appropriate backend (the
+`async`_ one), you can now take advantage of some requests parallelisation.
+
+Cherry on the cake, the `async`_ backend is available within both the *local*
+`Burp-UI`_ process but also within the `bui-agent`_!
+
+
+Architecture
+------------
+
+The architecture is described bellow:
+
+::
+
+	+---------------------+
+	|                     |
+	|     celery          |
+	|                     |
+	+---------------------+
+	| +-----------------+ |                                 +----------------------+
+	| |                 | |                                 |                      |
+	| |  worker 1       +----------------+------------------>     bui-monitor      |
+	| |                 | |              |                  |                      |
+	| +-----------------+ |              |                  +----------------------+
+	| +-----------------+ |              |                  | +------------------+ |
+	| |                 | |              |                  | |                  | |
+	| |  worker n       +----------------+                  | | burp -a m   (1)  | |
+	| |                 | |              |                  | |                  | |
+	| +-----------------+ |              |                  | +------------------+ |
+	+---------------------+              |                  | +------------------+ |
+										 |                  | |                  | |
+	+---------------------+              |                  | | burp -a m   (2)  | |
+	|                     |              |                  | |                  | |
+	|     burp-ui         |              |                  | +------------------+ |
+	|                     |              |                  | +------------------+ |
+	+---------------------+              |                  | |                  | |
+	| +-----------------+ |              |                  | | burp -a m   (n)  | |
+	| |                 | |              |                  | |                  | |
+	| |  worker 1       +----------------+                  | +------------------+ |
+	| |                 | |              |                  +----------------------+
+	| +-----------------+ |              |
+	| +-----------------+ |              |
+	| |                 | |              |
+	| |  worker n       +----------------+
+	| |                 | |
+	| +-----------------+ |
+	+---------------------+
+
+
+Requirements
+------------
+
+The monitor pool is powered by asyncio through trio.
+It is part of the `Burp-UI`_ package.
+You can launch it with the ``bui-monitor`` command.
+
+Configuration
+-------------
+
+There is a specific `buimanager.cfg`_ configuration file with a ``[Global]``
+section as below:
+
+::
+
+	# Burp-UI monitor configuration file
+	[Global]
+	# On which port is the application listening
+	port = 11111
+	# On which address is the application listening
+	# '::1' is the default for local IPv6
+	# set it to '127.0.0.1' if you want to listen on local IPv4 address
+	bind = ::1
+	# Pool size: number of 'burp -a m' process to load
+	pool = 5
+	# enable SSL
+	ssl = true
+	# ssl cert
+	sslcert = /var/lib/burp/ssl/server/ssl_cert-server.pem
+	# ssl key
+	sslkey = /var/lib/burp/ssl/server/ssl_cert-server.key
+	# monitor password
+	password = password123456
+
+	## burp backend specific options
+	#[Burp]
+	## burp binary
+	#burpbin = /usr/sbin/burp
+	## burp client configuration file used for the restoration
+	#bconfcli = /etc/burp/burp.conf
+	## how many time to wait for the monitor to answer (in seconds)
+	#timeout = 15
+
+
+Each option is commented, but here is a more detailed documentation:
+
+- *port*: On which port is `bui-monitor`_ listening.
+- *bind*: On which address is `bui-monitor`_ listening.
+- *pool*: Number of burp client processes to launch.
+- *ssl*: Whether to communicate with the `Burp-UI`_ server over SSL or not.
+- *sslcert*: What SSL certificate to use when SSL is enabled.
+- *sslkey*: What SSL key to use when SSL is enabled.
+- *password*: The shared secret between the `Burp-UI`_ server and `bui-monitor`_.
+
+As with `Burp-UI`_, you need the ``[Burp]`` section to specify `Burp`_ client options. There are fewer options because we only launch client processes.
+
+Service
+=======
+
+I have no plan to implement daemon features, but there are a lot of tools
+available to help you achieve such a behavior.
+
+To run bui-monitor as a service, a systemd file is provided. You can use it like
+this:
+
+::
+
+    cp /usr/local/share/burpui/contrib/systemd/bui-monitor.service /etc/systemd/system/
+    systemctl daemon-reload
+    systemctl enable bui-monitor.service
+    systemctl start bui-monitor.service
+
+
+
+.. _Burp: http://burp.grke.org/
+.. _Burp-UI: https://git.ziirish.me/ziirish/burp-ui
+.. _buimonitor.cfg: https://git.ziirish.me/ziirish/burp-ui/blob/master/share/burpui/etc/buimonitor.sample.cfg
+.. _bui-agent: buiagent.html
+.. _bui-monitor: buimonitor.html
+.. _burp2: advanced_usage.html#burp2
+.. _async: advanced_usage.html#async
+.. _celery: http://www.celeryproject.org/

docs/index.rst

@@ -32,6 +32,7 @@ Documentation
    gunicorn
    docker
    buiagent
+   buimonitor
    contributing
    changelog
    faq

setup.py

@@ -288,6 +288,7 @@ setup(
             'bui-celery=burpui.__main__:celery',
             'bui-manage=burpui.__main__:manage',
             'bui-agent-legacy=burpui.__main__:agent',
+            'bui-monitor=burpui.__main__:monitor',
             'burp-ui-legacy=burpui.__main__:legacy',
         ],
     },

share/burpui/etc/buimonitor.sample.cfg

@@ -0,0 +1,27 @@
+# Burp-UI monitor configuration file
+[Global]
+# On which port is the application listening
+port = 11111
+# On which address is the application listening
+# '::1' is the default for local IPv6
+# set it to '127.0.0.1' if you want to listen on local IPv4 address
+bind = ::1
+# Pool size: number of 'burp -a m' process to load
+pool = 5
+# enable SSL
+ssl = true
+# ssl cert
+sslcert = /var/lib/burp/ssl/server/ssl_cert-server.pem
+# ssl key
+sslkey = /var/lib/burp/ssl/server/ssl_cert-server.key
+# monitor password
+password = password123456
+
+## burp backend specific options
+#[Burp]
+## burp binary
+#burpbin = /usr/sbin/burp
+## burp client configuration file used for the restoration
+#bconfcli = /etc/burp/burp.conf
+## how many time to wait for the monitor to answer (in seconds)
+#timeout = 15

tools/bui-monitor

@@ -0,0 +1,3 @@
+#!/bin/bash
+
+python ./burpui -m monitor "$@"