1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
|
# -*- coding: utf-8 -*-
# _http.py
# Copyright (C) 2017 LEAP
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program. If not, see <http://www.gnu.org/licenses/>.
"""
A twisted-based HTTP client that:
- is pinned to a specific TLS certificate,
- does token authentication using the Authorization header,
- can do bandwidth throttling.
"""
import base64
import os
import sys
from twisted.internet import reactor
from twisted.protocols.policies import ThrottlingFactory
from twisted.protocols.policies import ThrottlingProtocol
from twisted.web.iweb import IAgent
from twisted.web.client import Agent as _Agent
from twisted.web.client import CookieAgent
from twisted.web.client import HTTPConnectionPool
from twisted.web.http_headers import Headers
from cookielib import CookieJar
from treq.client import HTTPClient as _HTTPClient
from zope.interface import implementer
from leap.common.http import getPolicyForHTTPS
__all__ = ['HTTPClient']
class HTTPClient(_HTTPClient):
def __init__(self, uuid, token, cert_file):
agent = Agent(uuid, token, cert_file)
jar = CookieJar()
self._agent = CookieAgent(agent, jar)
super(self.__class__, self).__init__(self._agent)
def set_token(self, token):
self._agent.set_token(token)
class HTTPThrottlingProtocol(ThrottlingProtocol):
def request(self, *args, **kwargs):
return self.wrappedProtocol.request(*args, **kwargs)
def throttleWrites(self):
if hasattr(self, 'producer') and self.producer:
self.producer.pauseProducing()
def unthrottleWrites(self):
if hasattr(self, 'producer') and self.producer:
self.producer.resumeProducing()
class HTTPThrottlingFactory(ThrottlingFactory):
protocol = HTTPThrottlingProtocol
class ThrottlingHTTPConnectionPool(HTTPConnectionPool):
maxPersistentPerHost = 1 # throttling happens "host-wise"
maxConnectionCount = sys.maxsize # max number of concurrent connections
readLimit = 1 * 10 ** 6 # max bytes we should read per second
writeLimit = 1 * 10 ** 6 # max bytes we should write per second
def _newConnection(self, key, endpoint):
def quiescentCallback(protocol):
self._putConnection(key, protocol)
factory = self._factory(quiescentCallback, repr(endpoint))
throttlingFactory = HTTPThrottlingFactory(
factory,
maxConnectionCount=self.maxConnectionCount,
readLimit=self.readLimit,
writeLimit=self.writeLimit)
return endpoint.connect(throttlingFactory)
@implementer(IAgent)
class Agent(_Agent):
def __init__(self, uuid, token, cert_file, throttling=False):
self._uuid = uuid
self._token = None
self._creds = None
self.set_token(token)
factory = getPolicyForHTTPS(cert_file)
pool = self._get_pool()
_Agent.__init__(self, reactor, contextFactory=factory, pool=pool)
def _get_pool(self):
throttling = bool(os.environ.get('SOLEDAD_THROTTLING'))
persistent = bool(os.environ.get('SOLEDAD_HTTP_PERSIST'))
if throttling:
klass = ThrottlingHTTPConnectionPool
else:
klass = HTTPConnectionPool
return klass(reactor, persistent=persistent)
def set_token(self, token):
self._token = token
self._creds = self._encoded_creds()
def _encoded_creds(self):
creds = '%s:%s' % (self._uuid, self._token)
encoded = base64.b64encode(creds)
return 'Token %s' % encoded
def request(self, method, uri, headers=None, bodyProducer=None):
# authenticate the request
headers = headers or Headers()
headers.addRawHeader('Authorization', self._creds)
# perform the authenticated request
return _Agent.request(
self, method, uri, headers=headers, bodyProducer=bodyProducer)
|
