1420216e5SHanoh Haim"""
2420216e5SHanoh HaimNTLM authenticating pool, contributed by erikcederstran
3420216e5SHanoh Haim
4420216e5SHanoh HaimIssue #10, see: http://code.google.com/p/urllib3/issues/detail?id=10
5420216e5SHanoh Haim"""
6420216e5SHanoh Haimfrom __future__ import absolute_import
7420216e5SHanoh Haim
8420216e5SHanoh Haimfrom logging import getLogger
9420216e5SHanoh Haimfrom ntlm import ntlm
10420216e5SHanoh Haim
11420216e5SHanoh Haimfrom .. import HTTPSConnectionPool
12420216e5SHanoh Haimfrom ..packages.six.moves.http_client import HTTPSConnection
13420216e5SHanoh Haim
14420216e5SHanoh Haim
15420216e5SHanoh Haimlog = getLogger(__name__)
16420216e5SHanoh Haim
17420216e5SHanoh Haim
18420216e5SHanoh Haimclass NTLMConnectionPool(HTTPSConnectionPool):
19420216e5SHanoh Haim    """
20420216e5SHanoh Haim    Implements an NTLM authentication version of an urllib3 connection pool
21420216e5SHanoh Haim    """
22420216e5SHanoh Haim
23420216e5SHanoh Haim    scheme = 'https'
24420216e5SHanoh Haim
25420216e5SHanoh Haim    def __init__(self, user, pw, authurl, *args, **kwargs):
26420216e5SHanoh Haim        """
27420216e5SHanoh Haim        authurl is a random URL on the server that is protected by NTLM.
28420216e5SHanoh Haim        user is the Windows user, probably in the DOMAIN\\username format.
29420216e5SHanoh Haim        pw is the password for the user.
30420216e5SHanoh Haim        """
31420216e5SHanoh Haim        super(NTLMConnectionPool, self).__init__(*args, **kwargs)
32420216e5SHanoh Haim        self.authurl = authurl
33420216e5SHanoh Haim        self.rawuser = user
34420216e5SHanoh Haim        user_parts = user.split('\\', 1)
35420216e5SHanoh Haim        self.domain = user_parts[0].upper()
36420216e5SHanoh Haim        self.user = user_parts[1]
37420216e5SHanoh Haim        self.pw = pw
38420216e5SHanoh Haim
39420216e5SHanoh Haim    def _new_conn(self):
40420216e5SHanoh Haim        # Performs the NTLM handshake that secures the connection. The socket
41420216e5SHanoh Haim        # must be kept open while requests are performed.
42420216e5SHanoh Haim        self.num_connections += 1
43420216e5SHanoh Haim        log.debug('Starting NTLM HTTPS connection no. %d: https://%s%s',
44420216e5SHanoh Haim                  self.num_connections, self.host, self.authurl)
45420216e5SHanoh Haim
46420216e5SHanoh Haim        headers = {}
47420216e5SHanoh Haim        headers['Connection'] = 'Keep-Alive'
48420216e5SHanoh Haim        req_header = 'Authorization'
49420216e5SHanoh Haim        resp_header = 'www-authenticate'
50420216e5SHanoh Haim
51420216e5SHanoh Haim        conn = HTTPSConnection(host=self.host, port=self.port)
52420216e5SHanoh Haim
53420216e5SHanoh Haim        # Send negotiation message
54420216e5SHanoh Haim        headers[req_header] = (
55420216e5SHanoh Haim            'NTLM %s' % ntlm.create_NTLM_NEGOTIATE_MESSAGE(self.rawuser))
56420216e5SHanoh Haim        log.debug('Request headers: %s', headers)
57420216e5SHanoh Haim        conn.request('GET', self.authurl, None, headers)
58420216e5SHanoh Haim        res = conn.getresponse()
59420216e5SHanoh Haim        reshdr = dict(res.getheaders())
60420216e5SHanoh Haim        log.debug('Response status: %s %s', res.status, res.reason)
61420216e5SHanoh Haim        log.debug('Response headers: %s', reshdr)
62420216e5SHanoh Haim        log.debug('Response data: %s [...]', res.read(100))
63420216e5SHanoh Haim
64420216e5SHanoh Haim        # Remove the reference to the socket, so that it can not be closed by
65420216e5SHanoh Haim        # the response object (we want to keep the socket open)
66420216e5SHanoh Haim        res.fp = None
67420216e5SHanoh Haim
68420216e5SHanoh Haim        # Server should respond with a challenge message
69420216e5SHanoh Haim        auth_header_values = reshdr[resp_header].split(', ')
70420216e5SHanoh Haim        auth_header_value = None
71420216e5SHanoh Haim        for s in auth_header_values:
72420216e5SHanoh Haim            if s[:5] == 'NTLM ':
73420216e5SHanoh Haim                auth_header_value = s[5:]
74420216e5SHanoh Haim        if auth_header_value is None:
75420216e5SHanoh Haim            raise Exception('Unexpected %s response header: %s' %
76420216e5SHanoh Haim                            (resp_header, reshdr[resp_header]))
77420216e5SHanoh Haim
78420216e5SHanoh Haim        # Send authentication message
79420216e5SHanoh Haim        ServerChallenge, NegotiateFlags = \
80420216e5SHanoh Haim            ntlm.parse_NTLM_CHALLENGE_MESSAGE(auth_header_value)
81420216e5SHanoh Haim        auth_msg = ntlm.create_NTLM_AUTHENTICATE_MESSAGE(ServerChallenge,
82420216e5SHanoh Haim                                                         self.user,
83420216e5SHanoh Haim                                                         self.domain,
84420216e5SHanoh Haim                                                         self.pw,
85420216e5SHanoh Haim                                                         NegotiateFlags)
86420216e5SHanoh Haim        headers[req_header] = 'NTLM %s' % auth_msg
87420216e5SHanoh Haim        log.debug('Request headers: %s', headers)
88420216e5SHanoh Haim        conn.request('GET', self.authurl, None, headers)
89420216e5SHanoh Haim        res = conn.getresponse()
90420216e5SHanoh Haim        log.debug('Response status: %s %s', res.status, res.reason)
91420216e5SHanoh Haim        log.debug('Response headers: %s', dict(res.getheaders()))
92420216e5SHanoh Haim        log.debug('Response data: %s [...]', res.read()[:100])
93420216e5SHanoh Haim        if res.status != 200:
94420216e5SHanoh Haim            if res.status == 401:
95420216e5SHanoh Haim                raise Exception('Server rejected request: wrong '
96420216e5SHanoh Haim                                'username or password')
97420216e5SHanoh Haim            raise Exception('Wrong server response: %s %s' %
98420216e5SHanoh Haim                            (res.status, res.reason))
99420216e5SHanoh Haim
100420216e5SHanoh Haim        res.fp = None
101420216e5SHanoh Haim        log.debug('Connection established')
102420216e5SHanoh Haim        return conn
103420216e5SHanoh Haim
104420216e5SHanoh Haim    def urlopen(self, method, url, body=None, headers=None, retries=3,
105420216e5SHanoh Haim                redirect=True, assert_same_host=True):
106420216e5SHanoh Haim        if headers is None:
107420216e5SHanoh Haim            headers = {}
108420216e5SHanoh Haim        headers['Connection'] = 'Keep-Alive'
109420216e5SHanoh Haim        return super(NTLMConnectionPool, self).urlopen(method, url, body,
110420216e5SHanoh Haim                                                       headers, retries,
111420216e5SHanoh Haim                                                       redirect,
112420216e5SHanoh Haim                                                       assert_same_host)
113