61 lines
2.3 KiB
Python
61 lines
2.3 KiB
Python
|
|
from flask import Blueprint
|
|
from flask import current_app
|
|
from flask import request
|
|
from werkzeug.exceptions import abort
|
|
|
|
from capsulflask.db import get_model
|
|
from capsulflask.shared import my_exec_info_message, authorized_as_hub
|
|
|
|
bp = Blueprint("hub", __name__, url_prefix="/hub")
|
|
|
|
def authorized_for_host(id):
|
|
if request.headers.get('Authorization'):
|
|
auth_header_value = request.headers.get('Authorization').replace("Bearer ", "")
|
|
return get_model().authorized_for_host(id, auth_header_value)
|
|
return False
|
|
|
|
@bp.route("/heartbeat-task", methods=("POST",))
|
|
def ping_all_hosts_task():
|
|
if authorized_as_hub(request.headers):
|
|
all_hosts = get_model().get_all_hosts()
|
|
current_app.logger.debug(f"pinging {len(all_hosts)} hosts...")
|
|
authorization_header = f"Bearer {current_app.config['HUB_TOKEN']}"
|
|
results = current_app.config["HTTP_CLIENT"].do_multi_http_sync(all_hosts, "/spoke/heartbeat", None, authorization_header=authorization_header)
|
|
for i in range(len(all_hosts)):
|
|
host = all_hosts[i]
|
|
result = results[i]
|
|
current_app.logger.debug(f"response from {host.id} ({host.url}): {result.status_code} {result.body}")
|
|
if result.status_code == 200:
|
|
get_model().host_heartbeat(host.id)
|
|
|
|
return "ok"
|
|
else:
|
|
current_app.logger.warning(f"/hub/heartbeat-task returned 401: invalid hub token")
|
|
return abort(401, "invalid hub token")
|
|
|
|
|
|
@bp.route("/heartbeat/<string:host_id>", methods=("POST",))
|
|
def heartbeat(host_id):
|
|
if authorized_for_host(host_id):
|
|
return "ok"
|
|
else:
|
|
current_app.logger.warning(f"/hub/heartbeat/{host_id} returned 401: invalid token")
|
|
return abort(401, "invalid host id or token")
|
|
|
|
@bp.route("/claim-operation/<int:operation_id>/<string:host_id>", methods=("POST",))
|
|
def claim_operation(operation_id: int, host_id: str):
|
|
if authorized_for_host(host_id):
|
|
exists = get_model().host_operation_exists(operation_id, host_id)
|
|
if not exists:
|
|
return abort(404, "host operation not found")
|
|
claimed = get_model().claim_operation(operation_id, host_id)
|
|
if claimed:
|
|
return "ok"
|
|
else:
|
|
return abort(409, "operation was already assigned to another host")
|
|
else:
|
|
current_app.logger.warning(f"/hub/claim-operation/{operation_id}/{host_id} returned 401: invalid token")
|
|
return abort(401, "invalid host id or token")
|
|
|