hmac auth changes

This commit is contained in:
simon987 2019-03-01 21:17:09 -05:00
parent 3400ee44ea
commit 0154c8dfaf

View File

@ -4,7 +4,7 @@ import hmac
import json import json
import os import os
import time import time
from email.utils import formatdate
import requests import requests
API_TIMEOUT = 5 API_TIMEOUT = 5
@ -109,7 +109,7 @@ class Worker:
return None return None
def format_headers(ua: str = None, wid: int = None, signature: str = None): def format_headers(ts: str = None, ua: str = None, wid: int = None, signature: str = None):
headers = dict() headers = dict()
if ua is None: if ua is None:
@ -119,6 +119,7 @@ def format_headers(ua: str = None, wid: int = None, signature: str = None):
headers["X-Worker-Id"] = str(wid) headers["X-Worker-Id"] = str(wid)
headers["X-Signature"] = str(signature) headers["X-Signature"] = str(signature)
headers["Timestamp"] = str(ts)
return headers return headers
@ -197,8 +198,9 @@ class TaskTrackerApi:
def _http_get(self, endpoint: str, worker: Worker = None): def _http_get(self, endpoint: str, worker: Worker = None):
if worker is not None: if worker is not None:
signature = hmac.new(key=worker.secret, msg=endpoint.encode("utf8"), digestmod=hashlib.sha256).hexdigest() ts = formatdate(timeval=None, localtime=False, usegmt=True)
headers = format_headers(signature=signature, wid=worker.id) signature = hmac.new(key=worker.secret, msg=(endpoint + ts).encode("utf8"), digestmod=hashlib.sha256).hexdigest()
headers = format_headers(signature=signature, wid=worker.id, ts=ts)
else: else:
headers = format_headers() headers = format_headers()
retries = 0 retries = 0
@ -221,18 +223,19 @@ class TaskTrackerApi:
def _http_post(self, endpoint: str, body, worker: Worker = None): def _http_post(self, endpoint: str, body, worker: Worker = None):
body_bytes = json.dumps(body).encode("utf8") body = json.dumps(body)
if worker is not None: if worker is not None:
signature = hmac.new(key=worker.secret, msg=body_bytes, digestmod=hashlib.sha256).hexdigest() ts = formatdate(timeval=None, localtime=False, usegmt=True)
headers = format_headers(signature=signature, wid=worker.id) signature = hmac.new(key=worker.secret, msg=(body + ts).encode("utf8"), digestmod=hashlib.sha256).hexdigest()
headers = format_headers(signature=signature, wid=worker.id, ts=ts)
else: else:
headers = format_headers() headers = format_headers()
retries = 0 retries = 0
while retries < MAX_HTTP_RETRIES: while retries < MAX_HTTP_RETRIES:
try: try:
response = requests.post(self.url + endpoint, timeout=API_TIMEOUT, response = requests.post(self.url + endpoint, timeout=API_TIMEOUT,
headers=headers, data=body_bytes) headers=headers, data=body.encode("utf8"))
if response.status_code == 429: if response.status_code == 429:
delay = json.loads(response.text)["rate_limit_delay"] * 20 delay = json.loads(response.text)["rate_limit_delay"] * 20