capsul-flask/capsulflask/hub_api.py

175 lines
6.5 KiB
Python

import json
import ipaddress
from flask import Blueprint
from flask import current_app
from flask import request, make_response
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):
payload_json = get_model().get_payload_json_from_host_operation(operation_id, host_id)
if payload_json is None:
error_message = f"{host_id} can't claim operation {operation_id} because host_operation row not found"
current_app.logger.error(error_message)
return abort(404, error_message)
# right now if there is a can_claim_handler there needs to be a corresponding on_claimed_handler.
# this is sole-ly due to laziness in error handling.
can_claim_handlers = {
"create": can_claim_create,
}
on_claimed_handlers = {
"create": on_create_claimed,
}
error_message = ""
payload = None
payload_is_dict = False
payload_has_type = False
payload_has_valid_type = False
try:
payload = json.loads(payload_json)
payload_is_dict = isinstance(payload, dict)
payload_has_type = payload_is_dict and 'type' in payload
payload_has_valid_type = payload_has_type and payload['type'] in can_claim_handlers
if not payload_is_dict:
error_message = "invalid json: expected an object"
elif not payload_has_type:
error_message = "invalid json: 'type' field is required"
elif not payload_has_valid_type:
error_message = f"invalid json: expected type \"{payload['type']}\" to be one of [{', '.join(can_claim_handlers.keys())}]"
except:
error_message = "could not parse payload as json"
if error_message != "":
error_message = f"{host_id} can't claim operation {operation_id} because {error_message}"
current_app.logger.error(error_message)
return abort(400, error_message)
# we will only return this payload as json if claiming succeeds, so might as well do this now...
payload['assignment_status'] = 'assigned'
# invoke the appropriate can_claim_handler for this operation type
result_tuple = can_claim_handlers[payload['type']](payload, host_id)
modified_payload = result_tuple[0]
error_message = result_tuple[1]
if error_message != "":
error_message = f"{host_id} can't claim operation {operation_id} because {error_message}"
current_app.logger.error(error_message)
return abort(400, error_message)
claimed = get_model().claim_operation(operation_id, host_id)
if claimed:
modified_payload_json = json.dumps(modified_payload)
get_model().update_operation(operation_id, modified_payload_json)
on_claimed_handlers[payload['type']](modified_payload, host_id)
response = make_response(modified_payload_json)
response.headers.set("Content-Type", "application/json")
return response
else:
return abort(409, f"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")
def can_claim_create(payload, host_id) -> (str, str):
hosts = get_model().list_hosts_with_networks(host_id)
if host_id not in hosts:
return "", f"the host \"{host_id}\" does not appear to have any networks."
networks = hosts[host_id]['networks']
vms_by_host_and_network = get_model().non_deleted_vms_by_host_and_network(host_id)
vms_by_network = dict()
if host_id in vms_by_host_and_network:
vms_by_network = vms_by_host_and_network[host_id]
allocated_ipv4_address = None
allocated_network_name = None
for network in networks:
vms = []
if network["network_name"] in vms_by_network:
vms = vms_by_network[network["network_name"]]
claimed_ipv4s = dict()
for vm in vms:
claimed_ipv4s[vm['public_ipv4']] = True
ipv4_network = ipaddress.ip_network(network["public_ipv4_cidr_block"], False)
i = 0
for ipv4_address in ipv4_network:
i += 1
if i > 2 and str(ipv4_address) not in claimed_ipv4s:
allocated_ipv4_address = str(ipv4_address)
break
if allocated_ipv4_address is not None:
allocated_network_name = network["network_name"]
break
if allocated_network_name is None or allocated_ipv4_address is None:
return "", f"host \"{host_id}\" does not have any avaliable IP addresses on any of its networks."
payload["network_name"] = allocated_network_name
payload["public_ipv4"] = allocated_ipv4_address
return payload, ""
def on_create_claimed(payload, host_id):
get_model().create_vm(
email=payload['email'],
id=payload['id'],
size=payload['size'],
os=payload['os'],
host=host_id,
network_name=payload['network_name'],
public_ipv4=payload['public_ipv4'],
ssh_authorized_keys=list(map(lambda x: x["name"], payload['ssh_authorized_keys'])),
)