0001import re
0002from paste.response import header_value
0003
0004class Filter(object):
0005    """
0006    Class that implements WSGI output-filtering middleware
0007    """
0008
0009    # If this is true, then conditional requests will be diabled
0010    # (e.g., If-Modified-Since)
0011    force_no_conditional = True
0012
0013    conditional_headers = [
0014        'HTTP_IF_MODIFIED_SINCE',
0015        'HTTP_IF_NONE_MATCH',
0016        ]
0017
0018    # If true, then any status code will be filtered; otherwise only
0019    # 200 OK responses are filtered
0020    filter_all_status = False
0021
0022    # If you provide this (a string or list of string mimetypes) then
0023    # only content with this mimetype will be filtered
0024    filter_content_types = ('text/html', )
0025
0026    # If this is set, then HTTPEncode will be used to decode the value
0027    # given provided mimetype and this output
0028    format_output = None
0029
0030    # You can also use a specific format object, which forces the
0031    # parsing with that format
0032    format = None
0033
0034    # If you aren't using a format but you want unicode instead of
0035    # 8-bit strings, then set this to true
0036    decode_unicode = False
0037
0038    # When we get unicode back from the filter, we'll use this
0039    # encoding and update the Content-Type:
0040    output_encoding = 'utf8'
0041
0042    def __init__(self, app):
0043        self.app = app
0044        if (self.format is not None
0045            and self.filter_content_types is Filter.filter_content_types):
0046            self.filter_content_types = format.content_types
0047
0048    def __call__(self, environ, start_response):
0049        if self.force_no_conditional:
0050            for key in self.conditional_headers:
0051                if key in environ:
0052                    del environ[key]
0053        # @@: I should actually figure out a way to deal with some
0054        # encodings, particular since stuff we don't care about like
0055        # text/javascript could be gzipped usefully.
0056        if 'HTTP_ACCEPT_ENCODING' in environ:
0057            del environ['HTTP_ACCEPT_ENCODING']
0058        shortcutted = []
0059        captured = []
0060        written_output = []
0061        def replacement_start_response(status, headers, exc_info=None):
0062            if not self.should_filter(status, headers, exc_info):
0063                shortcutted.append(None)
0064                return start_response(status, headers, exc_info)
0065            if exc_info is not None and shortcutted:
0066                raise exc_info[0], exc_info[1], exc_info[2]
0067            # Otherwise we don't care about exc_info...
0068            captured[:] = [status, headers]
0069            return written_output.append
0070        app_iter = self.app(environ, replacement_start_response)
0071        if shortcutted:
0072            # We chose not to filter
0073            return app_iter
0074        if not captured or written_output:
0075            # This app hasn't called start_response We can't do
0076            # anything magic with it; or it used the start_response
0077            # writer, and we still can't do anything with it
0078            try:
0079                for chunk in app_iter:
0080                    written_output.append(chunk)
0081            finally:
0082                if hasattr(app_iter, 'close'):
0083                    app_iter.close()
0084            app_iter = written_output
0085        try:
0086            return self.filter_output(
0087                environ, start_response,
0088                captured[0], captured[1], app_iter)
0089        finally:
0090            if hasattr(app_iter, 'close'):
0091                app_iter.close()
0092
0093    def paste_deploy_middleware(cls, app, global_conf, **app_conf):
0094        # You may wish to override this to make it convert the
0095        # arguments or use global_conf.  To declare your entry
0096        # point use:
0097        # setup(
0098        #   entry_points="""
0099        #   [paste.filter_app_factory]
0100        #   myfilter = myfilter:MyFilter.paste_deploy_middleware
0101        #   """)
0102        return cls(app, **app_conf)
0103
0104    paste_deploy_middleware = classmethod(paste_deploy_middleware)
0105
0106    def should_filter(self, status, headers, exc_info):
0107        if not self.filter_all_status:
0108            if not status.startswith('200'):
0109                return False
0110        content_type = header_value(headers, 'content-type')
0111        if content_type and ';' in content_type:
0112            content_type = content_type.split(';', 1)[0]
0113        if content_type in self.filter_content_types:
0114            return True
0115        return False
0116
0117    _charset_re = re.compile(
0118        r'charset="?([a-z0-9-_.]+)"?', re.I)
0119
0120    # @@: I should do something with these:
0121    #_meta_equiv_type_re = re.compile(
0122    #    r'<meta[^>]+http-equiv="?content-type"[^>]*>', re.I)
0123    #_meta_equiv_value_re = re.compile(
0124    #    r'value="?[^">]*"?', re.I)
0125
0126    def filter_output(self, environ, start_response,
0127                      status, headers, app_iter):
0128        content_type = header_value(headers, 'content-type')
0129        if ';' in content_type:
0130            content_type = content_type.split(';', 1)[0]
0131        if self.format_output:
0132            import httpencode
0133            format = httpencode.find_format_match(self.format_output, content_type)
0134        else:
0135            format = self.format
0136        if format:
0137            data = format.parse_response(headers, app_iter)
0138        else:
0139            data = ''.join(app_iter)
0140            if self.decode_unicode:
0141                # @@: Need to calculate encoding properly
0142                full_ct = header_value(headers, 'content-type') or ''
0143                match = self._charset_re.search(full_ct)
0144                if match:
0145                    encoding = match.group(0)
0146                else:
0147                    # @@: Obviously not a great guess
0148                    encoding = 'utf8'
0149                data = data.decode(encoding, 'replace')
0150        new_output = self.filter(
0151            environ, headers, data)
0152        if format:
0153            app = format.responder(new_output, headers)
0154            return app(environ, start_response)
0155        else:
0156            enc_data = []
0157            encoding = self.output_encoding
0158            if not isinstance(new_output, basestring):
0159                for chunk in new_output:
0160                    if isinstance(chunk, unicode):
0161                        chunk = chunk.encode(encoding)
0162                    enc_data.append(chunk)
0163            elif isinstance(new_output, unicode):
0164                enc_data.append(new_output.encode(encoding))
0165            else:
0166                enc_data.append(new_output)
0167            start_response(status, headers)
0168            return enc_data
0169
0170    def filter(self, environ, headers, data):
0171        raise NotImplementedError