1"""
2NTLM authenticating pool, contributed by erikcederstran
3
4Issue #10, see: http://code.google.com/p/urllib3/issues/detail?id=10
5"""
6from __future__ import absolute_import
7
8try:
9    from http.client import HTTPSConnection
10except ImportError:
11    from httplib import HTTPSConnection
12from logging import getLogger
13from ntlm import ntlm
14
15from urllib3 import HTTPSConnectionPool
16
17
18log = getLogger(__name__)
19
20
21class NTLMConnectionPool(HTTPSConnectionPool):
22    """
23    Implements an NTLM authentication version of an urllib3 connection pool
24    """
25
26    scheme = 'https'
27
28    def __init__(self, user, pw, authurl, *args, **kwargs):
29        """
30        authurl is a random URL on the server that is protected by NTLM.
31        user is the Windows user, probably in the DOMAIN\\username format.
32        pw is the password for the user.
33        """
34        super(NTLMConnectionPool, self).__init__(*args, **kwargs)
35        self.authurl = authurl
36        self.rawuser = user
37        user_parts = user.split('\\', 1)
38        self.domain = user_parts[0].upper()
39        self.user = user_parts[1]
40        self.pw = pw
41
42    def _new_conn(self):
43        # Performs the NTLM handshake that secures the connection. The socket
44        # must be kept open while requests are performed.
45        self.num_connections += 1
46        log.debug('Starting NTLM HTTPS connection no. %d: https://%s%s' %
47                  (self.num_connections, self.host, self.authurl))
48
49        headers = {}
50        headers['Connection'] = 'Keep-Alive'
51        req_header = 'Authorization'
52        resp_header = 'www-authenticate'
53
54        conn = HTTPSConnection(host=self.host, port=self.port)
55
56        # Send negotiation message
57        headers[req_header] = (
58            'NTLM %s' % ntlm.create_NTLM_NEGOTIATE_MESSAGE(self.rawuser))
59        log.debug('Request headers: %s' % headers)
60        conn.request('GET', self.authurl, None, headers)
61        res = conn.getresponse()
62        reshdr = dict(res.getheaders())
63        log.debug('Response status: %s %s' % (res.status, res.reason))
64        log.debug('Response headers: %s' % reshdr)
65        log.debug('Response data: %s [...]' % res.read(100))
66
67        # Remove the reference to the socket, so that it can not be closed by
68        # the response object (we want to keep the socket open)
69        res.fp = None
70
71        # Server should respond with a challenge message
72        auth_header_values = reshdr[resp_header].split(', ')
73        auth_header_value = None
74        for s in auth_header_values:
75            if s[:5] == 'NTLM ':
76                auth_header_value = s[5:]
77        if auth_header_value is None:
78            raise Exception('Unexpected %s response header: %s' %
79                            (resp_header, reshdr[resp_header]))
80
81        # Send authentication message
82        ServerChallenge, NegotiateFlags = \
83            ntlm.parse_NTLM_CHALLENGE_MESSAGE(auth_header_value)
84        auth_msg = ntlm.create_NTLM_AUTHENTICATE_MESSAGE(ServerChallenge,
85                                                         self.user,
86                                                         self.domain,
87                                                         self.pw,
88                                                         NegotiateFlags)
89        headers[req_header] = 'NTLM %s' % auth_msg
90        log.debug('Request headers: %s' % headers)
91        conn.request('GET', self.authurl, None, headers)
92        res = conn.getresponse()
93        log.debug('Response status: %s %s' % (res.status, res.reason))
94        log.debug('Response headers: %s' % dict(res.getheaders()))
95        log.debug('Response data: %s [...]' % res.read()[:100])
96        if res.status != 200:
97            if res.status == 401:
98                raise Exception('Server rejected request: wrong '
99                                'username or password')
100            raise Exception('Wrong server response: %s %s' %
101                            (res.status, res.reason))
102
103        res.fp = None
104        log.debug('Connection established')
105        return conn
106
107    def urlopen(self, method, url, body=None, headers=None, retries=3,
108                redirect=True, assert_same_host=True):
109        if headers is None:
110            headers = {}
111        headers['Connection'] = 'Keep-Alive'
112        return super(NTLMConnectionPool, self).urlopen(method, url, body,
113                                                       headers, retries,
114                                                       redirect,
115                                                       assert_same_host)
116