175 lines
6.5 KiB
Python
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'])),
|
|
) |