flexget.plugins.modify_headers
Covered: 40 lines
Missed: 7 lines
Skipped 20 lines
Percent: 85 %
 1
import logging
 2
import urllib2
 3
from flexget.plugin import *
 5
log = logging.getLogger('headers')
 8
class HTTPHeadersProcessor(urllib2.BaseHandler):
11
    handler_order = urllib2.HTTPHandler.handler_order - 10
13
    def __init__(self, headers={}):
14
        self.headers = headers
16
    def http_request(self, request):
17
        for name, value in self.headers.iteritems():
18
            if not request.has_header(name):
19
                log.debug('Adding %s: %s' % (name, value))
20
                request.add_unredirected_header(name.capitalize(), value.strip())
21
        return request
23
    def http_response(self, request, response):
24
        return response
26
    https_request = http_request
27
    https_response = http_response
30
class PluginHeaders:
31
    """Allow setting up any headers in all requests (which use urllib2)
33
    Example:
35
    headers:
36
      cookie: uid=<YOUR UID>; pass=<YOUR PASS>
37
    """
39
    def validator(self):
40
        from flexget import validator
41
        config = validator.factory('dict')
42
        config.accept_any_key('text')
43
        config.accept_any_key('integer')
44
        return config
46
    def on_feed_start(self, feed):
47
        """Feed starting"""
48
        config = feed.config['headers']
49
        if urllib2._opener:
50
            log.debug('Adding HTTPCaptureHeaderHandler to default opener')
51
            urllib2._opener.add_handler(HTTPHeadersProcessor(config))
52
        else:
53
            log.debug('Creating new opener and installing it')
54
            opener = urllib2.build_opener(HTTPHeadersProcessor(config))
55
            urllib2.install_opener(opener)
57
    def on_feed_exit(self, feed):
58
        """Feed exiting, remove additions"""
59
        if urllib2._opener:
60
            log.debug('Removing urllib2 default opener')
62
            urllib2.install_opener(None)
64
    on_feed_abort = on_feed_exit
66
register_plugin(PluginHeaders, 'headers')