path: root/wsgitools/filters.py

__all__ = []

import sys
import time
import gzip
try:
    import cStringIO as StringIO
except ImportError:
    import StringIO

__all__.append("CloseableIterator")
class CloseableIterator:
    """Concatenating iterator with close attribute."""
    def __init__(self, close_function, *iterators):
        """If close_function is not None, it will be the close attribute of
        the created iterator object. Further parameters specify iterators
        that are to be concatenated."""
        if close_function is not None:
            self.close = close_function
        self.iterators = map(iter, iterators)
    def __iter__(self):
        """iterator interface"""
        return self
    def next(self):
        """iterator interface"""
        if not self.iterators:
            raise StopIteration
        try:
            return self.iterators[0].next()
        except StopIteration:
            self.iterators.pop(0)
            return self.next()

__all__.append("CloseableList")
class CloseableList(list):
    """A list with a close attribute."""
    def __init__(self, close_function, *args):
        """If close_function is not None, it will be the close attribute of
        the created list object. Other parameters are passed to the list
        constructor."""
        if close_function is not None:
            self.close = close_function
        list.__init__(self, *args)
    def __iter__(self):
        return CloseableIterator(getattr(self, "close", None),
                                 list.__iter__(self))

__all__.append("BaseWSGIFilter")
class BaseWSGIFilter:
    """Generic WSGI filter class to be used with WSGIFilterMiddleware.

    For each request a filter object gets created.
    The environment is then passed through filter_environ.
    Possible exceptions are filtered by filter_exc_info.
    After that for each (header, value) tuple filter_header is used.
    The resulting list is filtered through filter_headers.
    Any data is filtered through filter_data.
    In order to possibly append data the append_data method is invoked.
    When the request has finished handle_close is invoked.

    All methods do not modify the passed data by default. Passing the
    BaseWSGIFilter class to a WSGIFilterMiddleware will result in not modifying
    the request at all.
    """
    def __init__(self):
        """This constructor does nothing and can safely be overwritten. It is
        only listed here to document that it must be callable without additional
        parameters."""
        pass
    def filter_environ(self, environ):
        """Receives a dict with the environment passed to the wsgi application
        and a dict must be returned. The default is to return the same dict."""
        return environ
    def filter_exc_info(self, exc_info):
        """Receives either None or a tuple passed as third argument to
        start_response from the wrapped wsgi application. Either None or such a
        tuple must be returned."""
        return exc_info
    def filter_status(self, status):
        """Receives a status string passed as first argument to start_response
        from the wrapped wsgi application. A valid HTTP status string must be
        returned."""
        return status
    def filter_header(self, headername, headervalue):
        """This function is invoked for each (headername, headervalue) tuple in
        the second argument to the start_response from the wrapped wsgi
        application. Such a value or None for discarding the header must be
        returned."""
        return (headername, headervalue)
    def filter_headers(self, headers):
        """A list of headers passed as the second argument to the start_response
        from the wrapped wsgi application is passed to this function and such a
        list must also be returned."""
        return headers
    def filter_data(self, data):
        """For each string that is either written by the write callable or
        returned from the wrapped wsgi application this method is invoked. It
        must return a string."""
        return data
    def append_data(self):
        """This function can be used to append data to the response. A list of
        strings or some kind of iterable yielding strings has to be returned.
        The default is to return an empty list.
        """
        return []
    def handle_close(self):
        """This method is invoked after the request has finished."""
        pass

__all__.append("WSGIFilterMiddleware")
class WSGIFilterMiddleware:
    """This wsgi middleware can be used with specialized BaseWSGIFilters to
    modify wsgi requests and/or reponses."""
    def __init__(self, app, filterclass):
        """app is a wsgi application.
        filterclass is a subclass of BaseWSGIFilter or some class that
                    implements the interface."""
        self.app = app
        self.filterclass = filterclass
    def __call__(self, environ, start_response):
        """wsgi interface"""
        assert isinstance(environ, dict)
        filter = self.filterclass()
        environ = filter.filter_environ(environ)

        def modified_start_response(status, headers, exc_info=None):
            assert isinstance(status, str)
            assert isinstance(headers, list)
            exc_info = filter.filter_exc_info(exc_info)
            status = filter.filter_status(status)
            headers = (filter.filter_header(h, v) for h, v in headers)
            headers = [h for h in headers if h is not None]
            headers = filter.filter_headers(headers)
            write = start_response(status, headers, exc_info)
            def modified_write(data):
                write(filter.filter_data(data))
            return modified_write

        ret = self.app(environ, modified_start_response)
        assert hasattr(ret, "__iter__")

        def modified_close():
            filter.handle_close()
            getattr(ret, "close", lambda:0)()

        if isinstance(ret, list):
            return CloseableList(modified_close,
                                 [filter.filter_data(data) for data in ret]
                                 + list(filter.append_data()))
        ret = iter(ret)
        def late_append_data():
            """Invoke filter.append_data() after filter_data() has seen all
            data."""
            for data in filter.append_data():
                yield data
        return CloseableIterator(modified_close,
                                 (filter.filter_data(data) for data in ret),
                                 late_append_data())

__all__.append("RequestLogWSGIFilter")
class RequestLogWSGIFilter(BaseWSGIFilter):
    """This filter logs all requests in the apache log file format."""
    @classmethod
    def creator(cls, log):
        """Returns a function creating RequestLogWSGIFilters on given log file.
        log has to be a file-like object."""
        return lambda:cls(log)
    def __init__(self, log=sys.stdout):
        assert hasattr(log, "write")
        self.log = log
        self.time = time.strftime("%d/%b/%Y:%T %z")
        self.length = 0
    def filter_environ(self, environ):
        """BaseWSGIFilter interface"""
        assert isinstance(environ, dict)
        self.remote = environ.get("REMOTE_ADDR", "?")
        self.user = environ.get("REMOTE_USER", "-")
        self.reqmethod = environ["REQUEST_METHOD"]
        self.path = environ["SCRIPT_NAME"] + environ["PATH_INFO"]
        self.proto = environ.get("SERVER_PROTOCOL", None)
        self.referrer = environ.get("HTTP_REFERER", "-")
        self.useragent = environ.get("HTTP_USER_AGENT", "-")
        return environ
    def filter_status(self, status):
        """BaseWSGIFilter interface"""
        assert isinstance(status, str)
        self.status = status.split()[0]
        return status
    def filter_data(self, data):
        """BaseWSGIFilter interface"""
        self.length += len(data)
        return data
    def handle_close(self):
        """BaseWSGIFilter interface"""
        line = '%s %s - [%s]' % (self.remote, self.user, self.time)
        line = '%s "%s %s' % (line, self.reqmethod, self.path)
        if self.proto is not None:
            line = "%s %s" % (line, self.proto)
        line = '%s" %s %d' % (line, self.status, self.length)
        if self.referrer is not None:
            line = '%s "%s"' % (line, self.referrer)
        else:
            line += " -"
        if self.useragent is not None:
            line = '%s "%s"' % (line, self.useragent)
        else:
            line += " -"
        print >> self.log,  line

__all__.append("TimerWSGIFilter")
class TimerWSGIFilter(BaseWSGIFilter):
    """Replaces a specific string in the data returned from the filtered wsgi
    application with the time the request took. The string has to be exactly
    eight bytes long, defaults to "?GenTime" and must be an element of the
    iterable returned by the filtered application. If the application returns
    something like ["spam?GenTime", "?GenTime spam", "?GenTime"] only the last
    occurance get's replaced."""
    @classmethod
    def creator(cls, pattern):
        """Returns a function creating TimerWSGIFilters with a given pattern
        beeing a string of exactly eight bytes."""
        return lambda:cls(pattern)
    def __init__(self, pattern="?GenTime"):
        self.pattern = pattern
        self.start = time.time()
    def filter_data(self, data):
        """BaseWSGIFilter interface"""
        if data == self.pattern:
            return "%8.3g" % (time.time() - self.start)
        return data

__all__.append("EncodeWSGIFilter")
class EncodeWSGIFilter(BaseWSGIFilter):
    """Encodes all body data (no headers) with given charset."""
    @classmethod
    def creator(cls, charset):
        """Returns a function creating EncodeWSGIFilters with a given charset.
        """
        return lambda:cls(charset)
    def __init__(self, charset="utf-8"):
        self.charset = charset
    def filter_data(self, data):
        """BaseWSGIFilter interface"""
        return data.encode(self.charset)
    def filter_header(self, header, value):
        """BaseWSGIFilter interface"""
        if header.lower() != "content-type":
            return (header, value)
        return (header, "%s; charset=%s" % (value, self.charset))

__all__.append("GzipWSGIFilter")
class GzipWSGIFilter(BaseWSGIFilter):
    """Compresses content using gzip."""
    def __init__(self):
        self.compress = False
        self.sio = None
        self.gzip = None
    def filter_environ(self, environ):
        """BaseWSGIFilter interface"""
        assert isinstance(environ, dict)
        if "HTTP_ACCEPT_ENCODING" in environ:
            acceptenc = environ["HTTP_ACCEPT_ENCODING"].split(',')
            acceptenc = map(str.strip, acceptenc)
            if "gzip" in acceptenc:
                self.compress = True
                self.sio = StringIO.StringIO()
                self.gzip = gzip.GzipFile(fileobj=self.sio, mode="w")
        return environ
    def filter_headers(self, headers):
        """BaseWSGIFilter interface"""
        assert isinstance(headers, list)
        if self.compress:
            headers.append(("Content-encoding", "gzip"))
        return headers
    def filter_data(self, data):
        """BaseWSGIFilter interface"""
        if not self.compress:
            return data
        self.gzip.write(data)
        self.gzip.flush()
        data = self.sio.getvalue()
        self.sio.truncate(0)
        return data
    def append_data(self):
        """BaseWSGIFilter interface"""
        if not self.compress:
            return []
        self.gzip.close()
        data = self.sio.getvalue()
        return [data]