title: An extra argument on a flexible decorator
author: depierre
published: 2014-06-25
categories: Python
keywords: python, decorator, function, extra, argument


Recently, I worked on a fix for [an issue that was reported on OWTF's github
page](https://github.com/owtf/owtf/issues/242), which shows that OWTF had some
permission errors that were not well processed.

After two first attempts to fix that bug (that were not satisfying in my
opinion), I proposed another solution that is cleaner and more stable.

My fix requires [to
decorate](http://simeonfranklin.com/blog/2012/jul/1/python-decorators-in-12-steps/)
some basic i/o python functions in order to intercept the `OSError` and
`IOError` and properly exit OWTF.  
Since the project is multi threaded, not catching such errors would freeze the
program (that is what was happening in the issue I mentioned above).

I thought I could show how the fix works and how I added an extra argument to
these i/o functions.

<!---summary-->

# Quick view on the issue

A couple of days ago, an issue has been reported concerning some permission
problems.  
The scenario that triggers the bug is the following:

1. Run OWTF as root
2. Run OWTF a second time as a simple user

And OWTF would freeze (*please excuse the CamelCase, I swore to myself that I
would fix that as soon as I can*):

    :::py3tb
    Traceback (most recent call last):
        File "./owtf.py", line 410, in
            Core = core.Init(RootDir, OwtfPid) # Initialise Framework
        File "framework/core.py", line 408, in Init
            return Core(RootDir, OwtfPid)
        File "framework/core.py", line 61, in init
            self.CreateTempStorageDirs(OwtfPid)
        File "framework/core.py", line 87, in CreateTempStorageDirs
            os.makedirs(temp_storage)
        File "/usr/lib/python2.7/os.py", line 157, in makedirs
            mkdir(name, mode)
    OSError: [Errno 13] Permission denied: '/tmp/owtf/8117'

The expected behavior should be that OWTF would stop each of its child
processes, print the error and properly exit, instead of just printing the
error and freeze.

# Proposed fixes

In order to fix this issue, two previous fixes were proposed. But they were not
the best answers that we should implement.  
Let me explain why.

## Force the user to be root

Well, I think you already know why this is not the right answer just by reading
the title.

[The first
fix](https://github.com/owtf/owtf/commit/74095814bcc23dcd0667f2c9d54ac28fd17a3587)
consisted of checking if the user was root before initializing OWTF. If he is
not root exit, otherwise continue.

I was surprised at first become since I started using OWTF, I never needed to
run it as root.

![How about no](/static/images/blog/meme_how_about_no.jpg)

**This is bad** because the user should not be demanded to run the tool as root
if the tool runs fine without...

## Check all directories on start

[The second
fix](https://github.com/owtf/owtf/commit/3157669da1b4dfe331f2b146f4879244912a7dec)
looked cleaner.

Before initializing OWTF, a function would check all the files it would use and
see if it had the correct permissions on them. If not then it would print an
error message and exit before going any further.

I had a problem with this solution:

+ What would happen if, one day, we would like to move all *foo* files to a
  new *bar* directory?  
+ Or what would happen if a new mandatory feature would use its own
  directory for storing its outputs?

Well, running OWTF as root and then as a normal user will at first freeze like
before the fix.  
Then we would have to update the function that checks the permissions and
finally we would be able to correctly run OWTF.

And that is the problem! We would be required to update some codes! **That
means more maintenance and it kind of sucks...**

# The decorator solution

After [reversing the first
fix](https://github.com/owtf/owtf/commit/11716fcfa0f98cdb2529ad6e6c8f93f8e98af569)
and checking the second one, I thought it would be better for OWTF to decorate
often-used i/o python functions in order to catch `OSError` and `IOError` and
properly terminate OWTF.

Why do I think this is better? Well, as soon as the functions are decorated, we
would never have to maintain the code.

With the previous example, if we decide to move all *foo* files to the *bar*
directory, it will necessarily use one of the decorated functions and any error
would be caught.

## How to decorate?

OWTF's `Core` class, which represents the core of the project (what a
coincidence), provides a nice function named `FrameworkAbort`. This will ensure
that everything stops properly before returning and that's exactly what we are
looking for :)

So the idea is that whenever an `OSError`/`IOError` is raised the
`FrameworkAbort` function should be called.

Using a nice python decorator, this gives us the following snippet:

    :::python
    # Simplifyied code.
    def catch_error(func):
        """Decorator on I/O functions.

        If an error is detected, force OWTF to quit properly.

        """
        def io_error(*args, **kwargs):
            """Call the original function while checking for errors."""
            try:
                return func(*args, **kwargs)
            except (OSError, IOError) as e:
                FrameworkAbort(
                    "Error when calling '%s'! %s." % (func.__name__, str(e)))
                raise e
        return io_perm

    # Decorate the `open` function.
    open = catch_error(open)

And would be used like that the normal `open` function (primary objective of a
decorator you would say and you would be right).

I used the [`*args, **kwargs` python
feature](http://freepythontips.wordpress.com/2013/08/04/args-and-kwargs-in-python-explained/)
to have a nice flexible decorator.  
That way it can decorate anything like `open`, `os.makedirs`, `shutil.rmtree`,
etc., even though each of the function requires different arguments.

Sadly this fix is buggy :'(  
I did not see it at first but it definitely can break some stuffs!

## But this is buggy

Let me show you how this version is broken in some ways.

If you try something like this with the decorated `open` function:

    :::python
    try:
        a = open('non-existing_file', 'r')
    except IOError:  # Never catched because OWTF exited before.
        pass

The line 3 will never be reached because the `FrameworkAbort` will exit the
program. It means that the decorator implicitly force every i/o decorated
functions to succeed otherwise it will exit.

So if you wanted to do something like that:

    :::python
    try:
        with open('maybe-existing_file', 'r') as f:
            data = f.read()
    except IOError:
        data = ''

Well... You would be screwed :/

So I needed a way to specify that OWTF should not exit if the i/o operation
failed but was not mandatory. The only solution I found interesting was adding
an extra argument to the decorated functions.

Because we cannot explicitly specify an extra argument in the decorator inner
function (because of `*args, **kwargs`), we will have to work around.

## Add an extra argument to any decorated function

If you take some times and have a closer look at the current decorator
implementation, you surely already found the answer.

In fact, since we are using `**kwargs`, that will contain the dictionary of
every key-word arguments, we could simply check if our own one is in there
with:

    :::python
    owtf_clean = kwargs.get('owtf_clean', True)

Using `get` will do two things for us:

+ First if `'owtf_clean'` is not in the dictionary, it means that the user did
  not specify it and we therefore use the default `True` value.
+ Second, if it exists, it means the user did override it and we therefore
  retrieve its actual value.
+ So the perfect behavior for a key-word argument with a default value.

But that is still not enough. If we use `get`, this extra argument will be
passed to the decorated function and it will yell an error like:

    :::pycon
    # Call to the original `open` function.
    >>> open('non-existing_file', 'r', owtf_clean=True)
    Traceback (most recent call last):
      File "<stdin>", line 1, in <module>
    TypeError: 'owtf_clean' is an invalid keyword argument for this function

So instead of using `get`, we should use `pop` which will do the same job as
`get` but also it will remove that key from the dictionary.  
All in all, this gives us the following decorator:

    :::python
    def catch_error(func):
        """Decorator on I/O functions.

        If an error is detected, force OWTF to quit properly.

        """
        def io_error(*args, **kwargs):
            """Call the original function while checking for errors.

            If `owtf_clean` parameter is not explicitely passed or if it is
            set to `True`, it force OWTF to properly exit.

            """
            owtf_clean = kwargs.pop('owtf_clean', True)
            try:
                return func(*args, **kwargs)
            except (OSError, IOError) as e:
                if owtf_clean:
                    FrameworkAbort(
                        "Error when calling '%s'! %s." % (func.__name__, str(e)))
                raise e
        return io_error

Going back to the previous buggy example:

    :::python
    try:
        with open('maybe-existing_file', 'r', owtf_clean=False) as f:
            data = f.read()
    except IOError:
        data = ''

This works now like a charm :) You can check [the committed
version](https://github.com/owtf/owtf/blob/1bed8dc5d6b2343d39bb7bcfd95b981c1c74f07b/framework/core.py#L417)
on OWTF's repo.

# Conclusion

Simple, if you need to pass an extra argument to a function you should use a
decorator.  
But if you need your decorator to be as flexible as possible you should mix
`*args, **kwargs` and the `pop` dict method to retrieve your extra argument and
still have your decorated functions working :)

Be careful though, the name you give to the key-word argument should not be
already reserved by any of the functions you decorate. Otherwise it will never
be passed to the original function since you popped it before!