level3 url token authentication

# MIT Licence, Copyright 2015 by Fabian Topfstedt <[email protected]>
import datetime
import hashlib
import hmac
import urllib
import urlparse


def remove_from_qs(qs, param_name):
    """
    Gets a query string, removes all occurances of the given param_name and
    returns the resulting query string.
    """
    data = urlparse.parse_qs(qs, keep_blank_values=True)
    data.pop(param_name, None)
    return urllib.urlencode(data, doseq=True)


def add_to_qs(qs, param_name, param_value):
    """
    Gets a query string, adds a parameter and its value and returns the
    resulting query string.
    """
    data = urlparse.parse_qs(qs, keep_blank_values=True)
    values = data.get(param_name, [])
    values.append(param_value)
    data[param_name] = values
    return urllib.urlencode(data, doseq=True)


class UrlSigner(object):
    """
    A class to do Level3 URL Token Authentication.
    """
    def __init__(self, secret_nr, secret, ignored_params=["clientid"],
                 not_valid_before_param_name="stime",
                 not_valid_after_param_name="etime",
                 hash_param_name="encoded"):
        """
        secret_nr: A number 0-9 identifying the row in the secret table
        secret: A character string secret between 20 and 64 bytes long
        ignored_params: A list of parameter names to be ignored when hashing
        not_valid_before_param_name: A string
        not_valid_after_param_name: A string
        hash_param_name: A string
        """
        self.secret_nr = str(secret_nr)
        self.secret = str(secret)
        self.ignored_params = ignored_params
        self.not_valid_before_param_name = not_valid_before_param_name
        self.not_valid_after_param_name = not_valid_after_param_name
        self.hash_param_name = hash_param_name

    def _format_dt(self, dt):
        """
        Formats a datetime object to yyyymmddHHMMSS.
        """
        return dt.strftime("%Y%m%d%H%M%S")

    def _generate_token(self, path, query):
        """
        Generates the token that will be passed as the encoded paramter.
        """
        for ignored_param in self.ignored_params:
            query = remove_from_qs(query, ignored_param)
        uri = "%s?%s" % (path, query)
        hexdigest = hmac.new(self.secret, uri, hashlib.sha1).hexdigest()
        token = "%1.1s%20.20s" % (self.secret_nr, hexdigest)
        return token

    def sign_url(self, url, not_valid_before=None, not_valid_after=None,
                 expires_in_s=None):
        """
        Takes a URL and returns its signed representation.

        Not_valid_before and not_valid_after are optional UTC datetime objects.
        You can also set the validity time in a relative manner using the
        expires_in_s parameter, where requests are not valid before now and
        not valid after the given amount of seconds from now.
        """
        parsed = urlparse.urlparse(url)
        scheme, netloc, path, params, query, fragment = list(parsed)

        if not_valid_before:
            query = add_to_qs(query, self.not_valid_before_param_name,
                              self._format_dt(not_valid_before))

        if not_valid_after:
            query = add_to_qs(query, self.not_valid_after_param_name,
                              self._format_dt(not_valid_after))

        if expires_in_s:
            now = datetime.datetime.utcnow()
            then = now + datetime.timedelta(seconds=expires_in_s)
            query = add_to_qs(query, self.not_valid_before_param_name,
                              self._format_dt(now))
            query = add_to_qs(query, self.not_valid_after_param_name,
                              self._format_dt(then))

        token = self._generate_token(path, query)
        query = add_to_qs(query, self.hash_param_name, token)

        signed_url = urlparse.urlunparse([scheme, netloc, path, params, query,
                                          fragment])
        return signed_url