Changes for task_tracker@26dee89

This commit is contained in:
simon987 2019-03-30 10:17:35 -04:00
parent ef07c059ad
commit 3aa187c2c4

View File

@ -1,10 +1,8 @@
import base64 import base64
import hashlib
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
@ -71,7 +69,7 @@ class Worker:
self.id: int = wid self.id: int = wid
self.alias: str = alias self.alias: str = alias
self.secret: bytes = base64.b64decode(secret) self.secret: bytes = base64.b64decode(secret)
self._secret_b64 = secret self.secret_b64 = secret
self._api: TaskTrackerApi = api self._api: TaskTrackerApi = api
def fetch_task(self, project_id): def fetch_task(self, project_id):
@ -99,7 +97,7 @@ class Worker:
json.dump({ json.dump({
"id": self.id, "id": self.id,
"alias": self.alias, "alias": self.alias,
"secret": self._secret_b64 "secret": self.secret_b64
}, out) }, out)
@staticmethod @staticmethod
@ -112,7 +110,7 @@ class Worker:
return None return None
def format_headers(ts: str = None, ua: str = None, wid: int = None, signature: str = None): def format_headers(ua: str = None, wid: int = None, secret: str = None):
headers = dict() headers = dict()
if ua is None: if ua is None:
@ -121,8 +119,7 @@ def format_headers(ts: str = None, ua: str = None, wid: int = None, signature: s
headers["User-Agent"] = ua headers["User-Agent"] = ua
headers["X-Worker-Id"] = str(wid) headers["X-Worker-Id"] = str(wid)
headers["X-Signature"] = str(signature) headers["X-Secret"] = str(secret)
headers["Timestamp"] = str(ts)
return headers return headers
@ -170,13 +167,17 @@ class TaskTrackerApi:
def log(self, worker: Worker, level: int, message: str, timestamp: int, scope: str): def log(self, worker: Worker, level: int, message: str, timestamp: int, scope: str):
if level == LOG_TRACE: if level == LOG_TRACE:
return self._http_post("/log/trace", {"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker) return self._http_post("/log/trace",
{"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker)
if level == LOG_INFO: if level == LOG_INFO:
return self._http_post("/log/info", {"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker) return self._http_post("/log/info",
{"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker)
if level == LOG_WARN: if level == LOG_WARN:
return self._http_post("/log/warn", {"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker) return self._http_post("/log/warn",
{"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker)
if level == LOG_ERROR: if level == LOG_ERROR:
return self._http_post("/log/error", {"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker) return self._http_post("/log/error",
{"level": level, "message": message, "timestamp": timestamp, "scope": scope}, worker)
print("Invalid log level") print("Invalid log level")
@ -201,9 +202,7 @@ 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:
ts = formatdate(timeval=None, localtime=False, usegmt=True) headers = format_headers(secret=worker.secret_b64, 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
@ -229,9 +228,7 @@ class TaskTrackerApi:
body = json.dumps(body) body = json.dumps(body)
if worker is not None: if worker is not None:
ts = formatdate(timeval=None, localtime=False, usegmt=True) headers = format_headers(secret=worker.secret_b64, 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
@ -251,4 +248,3 @@ class TaskTrackerApi:
retries += 1 retries += 1
pass pass
return None return None