Ziirish's Home :: Blog

Ziirish's Pub

 
 

While working on Burp-UI some weeks ago, I did some refactoring using the magic getattribute method.

I did not find a lot of examples on the Internet because it is usually not recommended to play with this function, but I think the use-case I'm about to show you is interesting.

Introduction

Burp-UI allows you to run several agents if you have several Burp servers and all these agents will be available through one unique UI.

The following schema shows you how it works:

Burp-UI agents

Burp-UI agents (direct link)

Every agent must implement the BUIbackend interface, but in fact an Agent is just a basic server acting between the actual backend and the UI engine.

It reads the commands sent by the UI engine and pass them to the backend. This way, you can have both Burp 2 and Burp 1 servers all centralized in the same dashboard.

On the other side, the UI engine also have a special backend to interact with the good agent. This backend also have to respect the interface.

If we do the math, everytime I wanted to add a new feature I had to implement 4 functions: one per backend + the agent. That's a lot of work. So I decided to cleanup a bit this mess and here is the result.

The special UI backend is just a proxy between the UI and the agent. There is no business code in it. The agent is also here to interact between the special remote backend and the actual local backend that to the business. The Burp 2 backend inherit the Burp 1 one because both of them share some code. So if we redo the math at this point, we see we can almost divide the efforts by 4 because in most cases, only the Burp 1 backend will carry the business code.

The Agent

Indeed, the Agent code looked like this:

class BUIAgent:
    def __init__(self, conf=None, debug=False):
        # some code [...]
        module = 'burpui.misc.backend.burp{0}'.format(self.vers)
        try:
            mod = __import__(module, fromlist=['Burp'])
            Client = mod.Burp
            self.backend = Client(conf=conf)
        except Exception, e:
            self.app.logger.error('Failed loading backend for Burp version %d: %s', self.vers, str(e))
            sys.exit(2)

        self.methods = {
                'status': self.backend.status,
                'get_backup_logs': self.backend.get_backup_logs,
                'get_counters': self.backend.get_counters,
                'is_backup_running': self.backend.is_backup_running,
                'is_one_backup_running': self.backend.is_one_backup_running,
                'get_all_clients': self.backend.get_all_clients,
                'get_client': self.backend.get_client,
                'get_tree': self.backend.get_tree,
                'restore_files': self.backend.restore_files,
                'read_conf': self.backend.read_conf,
                'store_conf': self.backend.store_conf,
                'get_parser_attr': self.backend.get_parser_attr
            }

        self.server = AgentServer((self.bind, self.port), AgentTCPHandler, self)

class AgentTCPHandler(SocketServer.BaseRequestHandler):
    "One instance per connection.  Override handle(self) to customize action."
    def handle(self):
        # self.request is the client connection
        try:
            r, _, _ = select.select([self.request], [], [], 5)
            if not r:
                raise Exception ('Socket timed-out')
            lengthbuf = self.request.recv(8)
            length, = struct.unpack('!Q', lengthbuf)
            data = self.recvall(length)
            j = json.loads(data)
            _, w, _ = select.select([], [self.request], [], 5)
            if not w:
                raise Exception ('Socket timed-out')
            if j['func'] not in self.server.agent.methods:
                self.request.sendall('KO')
                return
            self.request.sendall('OK')
            if j['args']:
                res = json.dumps(self.server.agent.methods[j['func']](**j['args']))
            else:
                res = json.dumps(self.server.agent.methods[j['func']]())

As you can see, the Agent just keeps a list of methods and associate each one of them with the backend's function. It means anytime I wanted to add a new function, I had to update this list...

And here comes the magic! The python abc module allows us to define abstract base classes. So I used this module to define the BUIbackend interface. Now if we want to inherit the BUIbackend abstract class, we must implement its abstract methods. Thanks to this, we know exactly what methods are implemented by any class inheriting the interface.

The new agent code now looks like this:

class BurpHandler(BUIbackend):
    # These functions MUST be implemented because we inherit an abstract class.
    # The hack here is to get the list of the functions and let the interpreter
    # think we don't have to implement them.
    # Thanks to this list, we know what function are implemented by our backend.
    foreign = BUIbackend.__abstractmethods__
    BUIbackend.__abstractmethods__ = frozenset()

    def __init__(self, vers=1, conf=None):
        self.vers = vers

        module = 'burpui.misc.backend.burp{0}'.format(self.vers)
        try:
            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
            mod = __import__(module, fromlist=['Burp'])
            Client = mod.Burp
            self.backend = Client(conf=conf)
        except Exception as e:
            sys.exit(2)

    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':
            return object.__getattribute__(self, name)
        # now we can retrieve the 'foreign' list and know if the object called
        # is in the backend
        if name in self.foreign:
            return getattr(self.backend, name)
        return object.__getattribute__(self, name)

And now the handle code calls the methods like this: res = json.dumps(getattr(self.cli, j['func'])(**j['args']))

Pretty nice, no?

The multi backend

The special backend acting as a proxy between the UI and the agent is called multi and it looked like this:

class Burp(BUIbackend):

    def __init__(self, app=None, conf=None):
        self.app = app
        self.servers = {}
        self.app.config['SERVERS'] = []
        self.running = {}
        if conf:
            config = ConfigParser.ConfigParser()
            with open(conf) as fp:
                config.readfp(fp)
                for sec in config.sections():
                    r = re.match('^Agent:(.+)$', sec)
                    if r:
                        try:
                            host = config.get(sec, 'host')
                            port = config.getint(sec, 'port')
                            password = config.get(sec, 'password')
                            ssl = config.getboolean(sec, 'ssl')
                        except Exception, e:
                            self.app.logger.error(str(e))
                        self.servers[r.group(1)] = NClient(app, host, port, password, ssl)

        self.app.logger.debug(self.servers)
        for key, serv in self.servers.iteritems():
            self.app.config['SERVERS'].append(key)

    """
    Utilities functions
    """

    def status(self, query='\n', agent=None):
        """
        status connects to the burp status port, ask the given 'question' and
        parses the output in an array
        """
        return self.servers[agent].status(query)

    def get_backup_logs(self, n, c, forward=False, agent=None):
        """
        parse_backup_log parses the log.gz of a given backup and returns a dict
        containing different stats used to render the charts in the reporting view
        """
        return self.servers[agent].get_backup_logs(n, c, forward)

    """
    And so on...
    """

As you can see, this backend creates a NClient object per agent found in the configuration and it implements every method provided by the interface. You may also notice the code of each methods looks very similar.

Now what about the NClient object? Well, here it is (or was, actually):

class NClient(BUIbackend):

    def __init__(self, app=None, host=None, port=None, password=None, ssl=None):
        self.host = host
        self.port = port
        self.password = password
        self.ssl = ssl
        self.connected = False
        self.app = app

    """
    Some network related code
    """

    def do_command(self, data=None):
        self.conn()
        res = '[]'
        if not data:
            return res
        try:
            data['password'] = self.password
            raw = json.dumps(data)
            length = len(raw)
            self.sock.sendall(struct.pack('!Q', length))
            self.sock.sendall(raw)
            r, _, _ = select.select([self.sock], [], [], 5)
            if not r:
                raise Exception ('Socket timed-out')
            tmp = self.sock.recv(2)
            if 'OK' != tmp:
                return res
            r, _, _ = select.select([self.sock], [], [], 5)
            if not r:
                raise Exception ('Socket timed-out')
            lengthbuf = self.sock.recv(8)
            length, = struct.unpack('!Q', lengthbuf)
            r, _, _ = select.select([self.sock], [], [], 5)
            if not r:
                raise Exception ('Socket timed-out')
            res = self.sock.recv(length)
        except Exception, e:
            self.app.logger.error(str(e))
        finally:
            self.close()
            return res

    """
    Utilities functions
    """

    def status(self, query='\n', agent=None):
        """
        status connects to the burp status port, ask the given 'question' and
        parses the output in an array
        """
        data = {'func': 'status', 'args': {'query': query}}
        return json.loads(self.do_command(data))

    def get_backup_logs(self, n, c, forward=False, agent=None):
        """
        parse_backup_log parses the log.gz of a given backup and returns a dict
        containing different stats used to render the charts in the reporting view
        """
        data = {'func': 'get_backup_logs', 'args': {'n': n, 'c': c, 'forward': forward}}
        return json.loads(self.do_command(data))

Again, the code of each backend method is kinda boring because we just need to encode its arguments in a json and pass them through a socket.

The new code is sexier imho:

INTERFACE_METHODS = BUIbackend.__abstractmethods__


class ProxyCall(object):
    """Class to dispatch call of unknown methods in order to dynamically
    call agents one without maintaining the explicit list of methods.
    """
    def __init__(self, proxy, method, network=False):
        self.proxy = proxy
        self.method = method
        self.network = network

    def __call__(self, *args, **kwargs):
        """This is were the proxy call (and the magic) occurs"""
        # retrieve the original function prototype
        proto = getattr(BUIbackend, self.method)
        args_name = list(proto.__code__.co_varnames)
        # skip self
        args_name.pop(0)
        # we transform unnamed arguments to named ones
        # example:
        #     def my_function(toto, tata=None, titi=None):
        #
        #     x = my_function('blah', titi='blih')
        #
        # => {'toto': 'blah', 'titi': 'blih'}
        encoded_args = {}
        for idx, opt in enumerate(args):
            encoded_args[args_name[idx]] = opt
        encoded_args.update(kwargs)

        # Special case for network calls
        if self.network:
            data = {'func': self.method, 'args': encoded_args}
            if self.method == 'restore_files':
                return self.proxy.do_command(data)
            return json.loads(self.proxy.do_command(data))
        # normal case for "standard" interface
        if 'agent' not in encoded_args:
            raise AttributeError(str(encoded_args))
        agentName = encoded_args['agent']
        # we don't need this argument anymore
        del encoded_args['agent']
        try:
            agent = self.proxy.servers[agentName]
        except KeyError:
            # This exception should be forwarded to the final user
            if not agentName:
                msg = "You must provide an agent name"
            else:
                msg = "Agent '{}' not found".format(agentName)
            raise BUIserverException(msg)
        return getattr(agent, self.method)(**encoded_args)

class NClient(BUIbackend):
    # These functions MUST be implemented because we inherit an abstract class.
    # The hack here is to get the list of the functions and let the interpreter
    # think we don't have to implement them.
    # Thanks to this list, we know what function are implemented by our backend.
    foreign = INTERFACE_METHODS
    BUIbackend.__abstractmethods__ = frozenset()

    def __init__(self, app=None, host=None, port=None, password=None, ssl=None, timeout=5):
        self.host = host
        self.port = port
        self.password = password
        self.ssl = ssl
        self.connected = False
        self.app = app
        self.timeout = timeout or 5

    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':
            return object.__getattribute__(self, name)
        # now we can retrieve the 'foreign' list and know if the object called
        # needs a dynamic implementation
        if name in self.foreign:
            return ProxyCall(self, name, network=True)
        return object.__getattribute__(self, name)

Now this is getting pretty convenient because every backend method is dynamically implemented thanks to the __getattribute__ method. But what if I want to overwrite a method for some reason?

Well, using a custom decorator, we can and here is what it looks like:

def implement(func):
    """A decorator indicating the method is implemented.
    For the agent and the 'multi' backend, we inherit the backend interface but
    we don't really implement it because we just act as a proxy.
    But maintaining the exhaustive list of methods in several places to always
    implement the same "proxy" thing was painful so I ended up cheating to
    dynamically implement those methods thanks to the __getattribute__ magic
    function.
    But sometimes we want to implement specific things, hence this decorator
    to indicate we don't want the default "magic" implementation and use the
    custom implementation instead.
    """
    func.__ismethodimplemented__ = True
    return func

class Burp(BUIbackend):
    """The :class:`burpui.misc.backend.multi.Burp` class provides a consistent
    backend to interact with ``agents``.
    It is actually the *real* multi backend implementing the
    :class:`burpui.misc.backend.interface.BUIbackend` class.
    For each agent found in the configuration, it will load a
    :class:`burpui.misc.backend.multi.NClient` class.
    :param server: ``Burp-UI`` server instance in order to access logger
                   and/or some global settings
    :type server: :class:`burpui.server.BUIServer`
    :param conf: Configuration file to use
    :type conf: str
    """
    # These functions MUST be implemented because we inherit an abstract class.
    # The hack here is to get the list of the functions and let the interpreter
    # think we don't have to implement them.
    # Thanks to this list, we know what function are implemented by our backend.
    foreign = INTERFACE_METHODS
    BUIbackend.__abstractmethods__ = frozenset()

    def __init__(self, server=None, conf=None):
        """
        :param server: Application context
        :type server: :class:`burpui.server.BUIServer`
        """
        self.app = server
        self.acl_handler = server.acl_handler
        self.servers = {}
        self.app.config['SERVERS'] = []
        self.running = {}
        if conf:
            config = ConfigParser.ConfigParser()
            with open(conf) as fp:
                config.readfp(fp)
                for sec in config.sections():
                    r = re.match('^Agent:(.+)$', sec)
                    if r:
                        ssl = False
                        host = self._safe_config_get(config.get, 'host', sec)
                        port = self._safe_config_get(config.getint, 'port', sec, cast=int)
                        password = self._safe_config_get(config.get, 'password', sec)
                        ssl = self._safe_config_get(config.getboolean, 'ssl', sec, cast=bool)
                        timeout = self._safe_config_get(config.getint, 'timeout', sec, cast=int)

                        self.servers[r.group(1)] = NClient(self.app, host, port, password, ssl, timeout)

        self.app.config['SERVERS'] = [ x for x in viewkeys(self.servers) ]

    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':
            return object.__getattribute__(self, name)
        # now we can retrieve the 'foreign' list and know if the object called
        # needs to be "proxyfied"
        if name in self.foreign:
            proxy = True
            func = None
            try:
                func = object.__getattribute__(self, name)
                proxy = not getattr(func, '__ismethodimplemented__', False)
            except:
                pass
            self.logger.debug('func: {} - {}'.format(name, proxy))
            if proxy:
                return ProxyCall(self, name)
            elif func:
                return func
        return object.__getattribute__(self, name)

    @implement
    def is_one_backup_running(self, agent=None):
        """See :func:`burpui.misc.backend.interface.BUIbackend.is_one_backup_running`"""
        r = []
        if agent:
            r = self.servers[agent].is_one_backup_running(agent)
            self.running[agent] = r
        else:
            r = self._backup_running_parallel()

            self.running = r
        self.refresh = time.time()
        return r

Every method decorated with the implement decorator will be returned as is by the magic __getattribute__ function. The others will be dynamically generated!

Conclusion

I hope you now see what magic means when talking about python magic methods ;)