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.info(f"pinging {len(all_hosts)} hosts...") authorization_header = f"Bearer {current_app.config['HUB_TOKEN']}" results = current_app.config["HTTP_CLIENT"].make_requests_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.info(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.info(f"/hub/heartbeat-task returned 401: invalid hub token") return abort(401, "invalid hub token") @bp.route("/heartbeat/", methods=("POST",)) def heartbeat(host_id): if authorized_for_host(host_id): return "ok" else: current_app.logger.info(f"/hub/heartbeat/{host_id} returned 401: invalid token") return abort(401, "invalid host id or token") @bp.route("/claim-operation//", 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.info(f"/hub/claim-operation/{operation_id}/{host_id} returned 401: invalid token") return abort(401, "invalid host id or token")