flexget.plugins.modify.headers
Covered: 38 lines
Missed: 15 lines
Skipped 22 lines
Percent: 71 %
 1
import logging
 2
import urllib2
 3
from flexget.plugin import register_plugin, priority
 5
log = logging.getLogger('headers')
 8
class HTTPHeadersProcessor(urllib2.BaseHandler):
11
    handler_order = urllib2.HTTPHandler.handler_order - 10
13
    def __init__(self, headers=None):
14
        if headers:
15
            self.headers = headers
16
        else:
17
            self.headers = {}
19
    def http_request(self, request):
20
        for name, value in self.headers.iteritems():
21
            if not request.has_header(name):
22
                log.debug('Adding %s: %s' % (name, value))
23
                request.add_unredirected_header(name.capitalize(), value.strip())
24
        return request
26
    def http_response(self, request, response):
27
        return response
29
    https_request = http_request
30
    https_response = http_response
33
class PluginHeaders(object):
34
    """Allow setting up any headers in all requests (which use urllib2)
36
    Example:
38
    headers:
39
      cookie: uid=<YOUR UID>; pass=<YOUR PASS>
40
    """
42
    def validator(self):
43
        from flexget import validator
44
        config = validator.factory('dict')
45
        config.accept_valid_keys('text', key_type='text')
46
        return config
48
    @priority(130)
49
    def on_feed_start(self, feed, config):
50
        """Feed starting"""
52
        if feed.requests.headers:
53
            feed.requests.headers.update(config)
54
        else:
55
            feed.requests.headers = config
57
        if urllib2._opener:
58
            log.debug('Adding HTTPHeadersProcessor to default opener')
59
            urllib2._opener.add_handler(HTTPHeadersProcessor(config))
60
        else:
61
            log.debug('Creating new opener and installing it')
62
            opener = urllib2.build_opener(HTTPHeadersProcessor(config))
63
            urllib2.install_opener(opener)
65
    def on_feed_exit(self, feed, config):
66
        """Feed exiting, remove additions"""
67
        if urllib2._opener:
68
            log.debug('Removing urllib2 default opener')
70
            urllib2.install_opener(None)
72
    on_feed_abort = on_feed_exit
74
register_plugin(PluginHeaders, 'headers', api_ver=2)