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