def add_auth(self,
r: requests.PreparedRequest) -> requests.PreparedRequest:
rr = r.copy()
url = urlparse(r.url)
if 'Host' in r.headers:
netloc = r.headers['Host'].decode('utf-8')
else:
netloc = url.netloc
rr.url = urlunparse((url.scheme, netloc, url.path, url.params,
url.query, url.fragment))
if r.method == 'POST':
if r.body:
if isinstance(r.body, bytes):
body = dict(parse_qsl(r.body.decode("utf-8").strip()))
elif isinstance(r.body, str):
body = dict(parse_qsl(r.body.strip()))
r.body = urlencode(self.update_params(rr, body))
new_headers = r.headers
new_headers[
'Content-Type'] = 'application/x-www-form-urlencoded; charset=utf-8'
r.headers = new_headers
elif r.method == 'GET':
url = urlparse(r.url)
if url.query:
new_query = urlencode(
self.update_params(rr, dict(parse_qsl(url.query))))
r.url = urlunparse((url.scheme, url.netloc, url.path,
url.params, new_query, url.fragment))
return r
def __call__(self, request: PreparedRequest):
data = json.loads(request.body)
data['nonce'] = make_nonce()
signdata = json.dumps(data)
sign = hmac.new(self.api_secret, bytearray(signdata,
encoding=ENCODING),
hashlib.sha512).hexdigest()
request.headers.update({'Key': self.api_key, 'Sign': sign})
request.body = signdata
return request
def compress_body(request: requests.PreparedRequest, always: bool):
deflater = zlib.compressobj()
body_bytes = (request.body if isinstance(request.body, bytes) else
request.body.encode())
deflated_data = deflater.compress(body_bytes)
deflated_data += deflater.flush()
is_economical = len(deflated_data) < len(body_bytes)
if is_economical or always:
request.body = deflated_data
request.headers['Content-Encoding'] = 'deflate'
request.headers['Content-Length'] = str(len(deflated_data))
def __call__(self, r: PreparedRequest):
data = {
"bk_app_code": self.bk_app_code,
"bk_app_secret": self.bk_app_secret,
"bk_username": self.bk_username,
"operator": self.operator,
}
if self.bk_supplier_account:
data["bk_supplier_account"] = self.bk_supplier_account
r.body = update_request_body(r.body, data)
return r
def __call__(self, r: PreparedRequest):
# 针对get请求，添加auth参数到url中; 针对post请求，添加auth参数到body体中
auth_params = {
"app_code": self.app_code,
"app_secret": self.app_secret,
"bk_username": self.bk_username
}
if r.method in ["GET"]:
r.url = update_url_parameters(r.url, auth_params)
elif r.method in ["POST"]:
r.body = update_request_body(r.body, auth_params)
return r
def test_write_failure_event_to_db():
ServiceMonitor.create_table(wait=True)
monitor = WriteToDynamoMonitor()
request = PreparedRequest()
request.body = "The body of request"
request.headers = {"bill": "Some bill"}
response = Response()
response.status_code = 500
response.headers['ben'] = "Some Ben"
response.data = "The response text"
result = monitor.failure("some-service", request, response, 1108)
assert result.id is not None
fetched_result = ServiceMonitor.query(result.service_name, ServiceMonitor.id == result.id).next()
assert fetched_result.id == result.id
assert fetched_result.event_type == "FAILED"
assert fetched_result.timestamp == result.timestamp
assert fetched_result.ttl == result.ttl
def compress_request(
request: requests.PreparedRequest,
always: bool,
):
deflater = zlib.compressobj()
if isinstance(request.body, str):
body_bytes = request.body.encode()
elif hasattr(request.body, 'read'):
body_bytes = request.body.read()
else:
body_bytes = request.body
deflated_data = deflater.compress(body_bytes)
deflated_data += deflater.flush()
is_economical = len(deflated_data) < len(body_bytes)
if is_economical or always:
request.body = deflated_data
request.headers['Content-Encoding'] = 'deflate'
request.headers['Content-Length'] = str(len(deflated_data))
def sign(self, req: requests.PreparedRequest) -> requests.PreparedRequest:
method = req.method.upper()
assert method == 'POST', 'POST method only, method={}'.format(method)
body = req.body
assert body
data = dict(urllib.parse.parse_qsl(body))
data['nonce'] = self.get_nonce()
body = urllib.parse.urlencode(data)
req.body = body
secret = self.api_secret.encode()
signature = hmac.new(secret, body.encode(), hashlib.sha512).hexdigest()
headers = {
'key': self.api_key,
'sign': signature,
}
req.headers.update(headers)
return req
def _extract_response(page, encoding='utf8'):
history = []
set_cookies = []
res = None
try:
for i, url in enumerate(page['history']):
resource = page['resources'].pop(0)
while resource['request']['url'] != url:
resource = page['resources'].pop(0)
if resource['error']:
return resource['error'], None
request = resource['request']
req = PreparedRequest()
req.method = request['method'].encode(encoding)
req.url = request['url'].encode(encoding)
# Set Request Headers
req.headers = CaseInsensitiveDict()
for header in request['headers']:
req.headers[header['name'].encode(encoding)] = header['value'].encode(encoding)
# Set Request Cookies
req._cookies = RequestsCookieJar()
if set_cookies:
if 'Cookie' not in req.headers:
req.headers['Cookie'] = ""
else:
set_cookies.insert(0, '')
req.headers['Cookie'] += "; ".join(set_cookies)
if 'Cookie' in req.headers:
cookies = SimpleCookie()
cookies.load(req.headers['Cookie'])
for key, cookie in cookies.items():
req._cookies.set(key, cookie.value)
req.body = request.get('postData', None)
if req.body:
req.body = req.body.encode(encoding)
response = resource['endReply'] or resource['startReply']
res = Response()
res.encoding = encoding
res.url = response['url'].encode(encoding)
res.status_code = response['status']
for header in response['headers']:
res.headers[header['name'].encode(encoding)] = header['value'].encode(encoding)
if header['name'] == 'Set-Cookie':
set_cookies.append(res.headers[header['name']].rsplit(';', 1)[0])
res.history = list(history)
res.request = req
history.append(res)
res._content = re.sub(
(
'<html><head></head><body>'
'<pre style="word-wrap: break-word; white-space: pre-wrap;">(.*?)</pre>'
'</body></html>'
),
r'\1',
page['content'],
flags=re.DOTALL
).encode(encoding)
return None, res
except IndexError:
return {'errorCode': -1,
'errorString': 'An existing connection was forcibly closed by the remote host'}, None
def auth_request(self, r: P, nonce: str):
r.body = self.parse_data(r.body)
super().auth_request(r, nonce)
r.prepare_body(data=r.body, files=None)
def auth_request(self, r: P, nonce: str):
r.body = self.parse_data(r.body)
super().auth_request(r, nonce)
r.prepare_body(data=r.body, files=None)
