An extra argument on a flexible decorator

Published on Wednesday, 25 June 2014 in Python ; tagged with python, decorator, function, extra, argument ; text version

Recently, I worked on a fix for an issue that was reported on OWTF's github page, 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 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.

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):

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 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

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 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:

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 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:

# 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 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:

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:

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:

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

Using get will do two things for us:

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:

# 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:

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:

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 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!


contactdepier.re License WTFPL2