3wordchant/capsul-flask
1
0
Fork 1
Code Issues Pull Requests 1 Releases Wiki Activity

got httpclient working, spoke heartbeat is working

Browse Source
This commit is contained in:
forest
2021-01-04 13:32:52 -06:00
parent d9c30e1ef8
commit 6764c5c97d
15 changed files with 281 additions and 112 deletions
Download Patch File Download Diff File Expand all files Collapse all files

22
capsulflask/__init__.py
View File

@ -1,8 +1,10 @@
import logging
from logging.config import dictConfig as logging_dict_config
import atexit
import os
import hashlib
import requests
import stripe
from dotenv import load_dotenv, find_dotenv
@ -11,11 +13,13 @@ from flask_mail import Mail
from flask import render_template
from flask import url_for
from flask import current_app
from apscheduler.schedulers.background import BackgroundScheduler
from capsulflask import hub_model, spoke_model, cli
from capsulflask.btcpay import client as btcpay
from capsulflask.http_client import MyHTTPClient
load_dotenv(find_dotenv())
app = Flask(__name__)
@ -32,7 +36,6 @@ app.config.from_mapping(
SPOKE_HOST_ID=os.environ.get("SPOKE_HOST_ID", default="default"),
SPOKE_HOST_TOKEN=os.environ.get("SPOKE_HOST_TOKEN", default="default"),
HUB_TOKEN=os.environ.get("HUB_TOKEN", default="default"),
HUB_URL=os.environ.get("HUB_URL", default="https://capsul.org"),
DATABASE_URL=os.environ.get("DATABASE_URL", default="sql://postgres:dev@localhost:5432/postgres"),
DATABASE_SCHEMA=os.environ.get("DATABASE_SCHEMA", default="public"),
@ -56,6 +59,8 @@ app.config.from_mapping(
BTCPAY_URL=os.environ.get("BTCPAY_URL", default="https://btcpay.cyberia.club")
)
app.config['HUB_URL'] = os.environ.get("HUB_URL", default=app.config['BASE_URL'])
logging_dict_config({
'version': 1,
'formatters': {'default': {
@ -89,6 +94,19 @@ if app.config['HUB_MODE_ENABLED']:
if app.config['HUB_MODEL'] == "capsul-flask":
app.config['HUB_MODEL'] = hub_model.CapsulFlaskHub()
# debug mode (flask reloader) runs two copies of the app. When running in debug mode,
# we only want to start the scheduler one time.
if not app.debug or os.environ.get('WERKZEUG_RUN_MAIN') == 'true':
scheduler = BackgroundScheduler()
heartbeat_task_url = f"{app.config['HUB_URL']}/hub/heartbeat-task"
heartbeat_task_headers = {'Authorization': f"Bearer {app.config['HUB_TOKEN']}"}
heartbeat_task = lambda: requests.post(heartbeat_task_url, headers=heartbeat_task_headers)
scheduler.add_job(func=heartbeat_task, trigger="interval", seconds=5)
scheduler.start()
atexit.register(lambda: scheduler.shutdown())
else:
app.config['HUB_MODEL'] = hub_model.MockHub()
@ -107,6 +125,8 @@ if app.config['HUB_MODE_ENABLED']:
app.add_url_rule("/", endpoint="index")
if app.config['SPOKE_MODE_ENABLED']:
if app.config['SPOKE_MODEL'] == "shell-scripts":

3
capsulflask/cli.py
View File

@ -11,7 +11,8 @@ from flask import current_app
from psycopg2 import ProgrammingError
from flask_mail import Message
from capsulflask.db import get_model, my_exec_info_message
from capsulflask.db import get_model
from capsulflask.shared import my_exec_info_message
from capsulflask.console import get_account_balance
bp = Blueprint('cli', __name__)

3
capsulflask/console.py
View File

@ -15,7 +15,8 @@ from nanoid import generate
from capsulflask.metrics import durations as metric_durations
from capsulflask.auth import account_required
from capsulflask.db import get_model, my_exec_info_message
from capsulflask.db import get_model
from capsulflask.shared import my_exec_info_message
from capsulflask.payment import poll_btcpay_session
from capsulflask import cli

5
capsulflask/db.py
View File

@ -40,7 +40,7 @@ def init_app(app):
hasSchemaVersionTable = False
actionWasTaken = False
schemaVersion = 0
desiredSchemaVersion = 8
desiredSchemaVersion = 9
cursor = connection.cursor()
@ -126,5 +126,4 @@ def close_db(e=None):
db_model.cursor.close()
current_app.config['PSYCOPG2_CONNECTION_POOL'].putconn(db_model.connection)
def my_exec_info_message(exec_info):
return "{}: {}".format(".".join([exec_info[0].__module__, exec_info[0].__name__]), exec_info[1])

15
capsulflask/db_model.py
View File

@ -4,12 +4,9 @@
from nanoid import generate
from flask import current_app
from typing import List
from capsulflask.hub_model import HTTPResult
class OnlineHost:
def __init__(self, id: str, url: str):
self.id = id
self.url = url
from capsulflask.shared import OnlineHost
class DBModel:
#def __init__(self, connection: Psycopg2Connection, cursor: Psycopg2Cursor):
@ -277,15 +274,19 @@ class DBModel:
# ------ HOSTS ---------
def authorized_for_host(self, id, token) -> bool:
self.cursor.execute("SELECT id FROM hosts WHERE id = %s token = %s", (id, token))
self.cursor.execute("SELECT id FROM hosts WHERE id = %s AND token = %s", (id, token))
return self.cursor.fetchone() != None
def host_heartbeat(self, id) -> None:
self.cursor.execute("UPDATE hosts SET last_health_check = NOW() WHERE id = %s", (id,))
self.connection.commit()
def get_all_hosts(self) -> List[OnlineHost]:
self.cursor.execute("SELECT id, https_url FROM hosts")
return list(map(lambda x: OnlineHost(id=x[0], url=x[1]), self.cursor.fetchall()))
def get_online_hosts(self) -> List[OnlineHost]:
self.cursor.execute("SELECT id, https_url FROM hosts WHERE last_health_check > NOW() - INTERVAL '10 seconds'")
self.cursor.execute("SELECT id, https_url FROM hosts WHERE last_health_check > NOW() - INTERVAL '20 seconds'")
return list(map(lambda x: OnlineHost(id=x[0], url=x[1]), self.cursor.fetchall()))
def create_operation(self, online_hosts: List[OnlineHost], email: str, payload: str) -> int:

111
capsulflask/http_client.py
View File

@ -1,12 +1,14 @@
import sys
import json
import itertools
import time
import threading
import aiohttp
import asyncio
from flask import current_app
from capsulflask.db import my_exec_info_message
from capsulflask.db_model import OnlineHost
from capsulflask.shared import OnlineHost, my_exec_info_message
from typing import List
class HTTPResult:
@ -16,14 +18,23 @@ class HTTPResult:
class MyHTTPClient:
def __init__(self, timeout_seconds = 5):
self.client_session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=timeout_seconds))
self.event_loop = asyncio.get_event_loop()
self.timeout_seconds = timeout_seconds
self.client_session = None
def make_requests_sync(self, online_hosts: List[OnlineHost], body: str) -> List(HTTPResult):
self.event_loop.run_until_complete(self.make_requests(online_hosts=online_hosts, body=body))
def make_requests_sync(self, online_hosts: List[OnlineHost], url_suffix: str, body: str, authorization_header=None) -> List[HTTPResult]:
future = run_coroutine(self.make_requests(online_hosts=online_hosts, url_suffix=url_suffix, body=body, authorization_header=authorization_header))
return future.result()
def post_json_sync(self, url: str, body: str, method="POST", authorization_header=None) -> HTTPResult:
self.event_loop.run_until_complete(self.post_json_sync(method=method, url=url, body=body))
future = run_coroutine(self.post_json(method=method, url=url, body=body, authorization_header=authorization_header))
return future.result()
def get_client_session(self):
if not self.client_session:
self.client_session = aiohttp.ClientSession(timeout=aiohttp.ClientTimeout(total=self.timeout_seconds))
return self.client_session
async def post_json(self, url: str, body: str, method="POST", authorization_header=None) -> HTTPResult:
response = None
@ -31,12 +42,11 @@ class MyHTTPClient:
headers = {}
if authorization_header != None:
headers['Authorization'] = authorization_header
response = await self.client_session.request(
response = await self.get_client_session().request(
method=method,
url=url,
json=body,
headers=headers,
auth=aiohttp.BasicAuth("hub", current_app.config['HUB_TOKEN']),
verify_ssl=True,
)
except:
@ -59,15 +69,92 @@ class MyHTTPClient:
return HTTPResult(response.status, response_body)
async def make_requests(self, online_hosts: List[OnlineHost], body: str) -> List(HTTPResult):
async def make_requests(self, online_hosts: List[OnlineHost], url_suffix: str, body: str, authorization_header=None) -> List[HTTPResult]:
tasks = []
# append to tasks in the same order as online_hosts
for host in online_hosts:
tasks.append(
self.post_json(url=host.url, body=body)
self.post_json(url=f"{host.url}/{url_suffix}", body=body, authorization_header=authorization_header)
)
# gather is like Promise.all from javascript, it returns a future which resolves to an array of results
# in the same order as the tasks that we passed in -- which were in the same order as online_hosts
results = await asyncio.gather(*tasks)
return results
return results
# i lifted this direct from https://stackoverflow.com/a/58616001
# this is the bridge between Flask's one-thread-per-request world
# and aiohttp's event-loop based world
class EventLoopThread(threading.Thread):
loop = None
_count = itertools.count(0)
def __init__(self):
name = f"{type(self).__name__}-{next(self._count)}"
super().__init__(name=name, daemon=True)
def __repr__(self):
loop, r, c, d = self.loop, False, True, False
if loop is not None:
r, c, d = loop.is_running(), loop.is_closed(), loop.get_debug()
return (
f"<{type(self).__name__} {self.name} id={self.ident} "
f"running={r} closed={c} debug={d}>"
)
def run(self):
self.loop = loop = asyncio.new_event_loop()
asyncio.set_event_loop(loop)
try:
loop.run_forever()
finally:
try:
shutdown_asyncgens = loop.shutdown_asyncgens()
except AttributeError:
pass
else:
loop.run_until_complete(shutdown_asyncgens)
loop.close()
asyncio.set_event_loop(None)
def stop(self):
loop, self.loop = self.loop, None
if loop is None:
return
loop.call_soon_threadsafe(loop.stop)
self.join()
_lock = threading.Lock()
_loop_thread = None
def get_event_loop():
global _loop_thread
if _loop_thread is None:
with _lock:
if _loop_thread is None:
_loop_thread = EventLoopThread()
_loop_thread.start()
# give the thread up to a second to produce a loop
deadline = time.time() + 1
while not _loop_thread.loop and time.time() < deadline:
time.sleep(0.001)
return _loop_thread.loop
def stop_event_loop():
global _loop_thread
with _lock:
if _loop_thread is not None:
_loop_thread.stop()
_loop_thread = None
def run_coroutine(coro):
"""Run the coroutine in the event loop running in a separate thread
Returns a Future, call Future.result() to get the output
"""
return asyncio.run_coroutine_threadsafe(coro, get_event_loop())

35
capsulflask/hub_api.py
View File

@ -4,23 +4,46 @@ from flask import current_app
from flask import request
from werkzeug.exceptions import abort
from capsulflask.db import get_model, my_exec_info_message
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):
auth_header_value = request.headers.get('Authorization').replace("Bearer ", "")
return get_model().authorized_for_host(id, auth_header_value)
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/<string:host_id>", methods=("POST"))
@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/<string:host_id>", methods=("POST",))
def heartbeat(host_id):
if authorized_for_host(host_id):
get_model().host_heartbeat(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/<int:operation_id>/<string:host_id>", methods=("POST"))
@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)

40
capsulflask/hub_model.py
View File

@ -12,33 +12,9 @@ from time import sleep
from os.path import join
from subprocess import run
from capsulflask.db_model import OnlineHost
from capsulflask.spoke_model import validate_capsul_id
from capsulflask.db import get_model, my_exec_info_message
from capsulflask.db import get_model
from capsulflask.http_client import HTTPResult
class VirtualMachine:
def __init__(self, id, host, ipv4=None, ipv6=None):
self.id = id
self.host = host
self.ipv4 = ipv4
self.ipv6 = ipv6
class VirtualizationInterface:
def capacity_avaliable(self, additional_ram_bytes: int) -> bool:
pass
def get(self, id: str) -> VirtualMachine:
pass
def list_ids(self) -> list:
pass
def create(self, email: str, id: str, template_image_file_name: str, vcpus: int, memory: int, ssh_public_keys: list):
pass
def destroy(self, email: str, id: str):
pass
from capsulflask.shared import VirtualizationInterface, VirtualMachine, OnlineHost, validate_capsul_id, my_exec_info_message
class MockHub(VirtualizationInterface):
def capacity_avaliable(self, additional_ram_bytes):
@ -61,11 +37,11 @@ class MockHub(VirtualizationInterface):
class CapsulFlaskHub(VirtualizationInterface):
def generic_operation(self, hosts: List[OnlineHost], payload: str, immediate_mode: bool) -> Tuple[int, List[HTTPResult]]:
operation_id = get_model().create_operation(hosts, payload)
authorization_header = f"Bearer {current_app.config['HUB_TOKEN']}"
results = current_app.config["HTTP_CLIENT"].make_requests_sync(hosts, payload, authorization_header=authorization_header)
results = current_app.config["HTTP_CLIENT"].make_requests_sync(hosts, "/spoke/operation", payload, authorization_header=authorization_header)
for i in range(len(hosts)):
host = hosts[i]
result = results[i]
@ -137,7 +113,7 @@ class CapsulFlaskHub(VirtualizationInterface):
host = get_model().host_of_capsul(id)
if host is not None:
payload = json.dumps(dict(type="get", id=id))
op = await self.generic_operation([host], payload, True)
op = self.generic_operation([host], payload, True)
results = op[1]
for result in results:
try:
@ -152,7 +128,7 @@ class CapsulFlaskHub(VirtualizationInterface):
def list_ids(self) -> list:
online_hosts = get_model().get_online_hosts()
payload = json.dumps(dict(type="list_ids"))
op = await self.generic_operation(online_hosts, payload, False)
op = self.generic_operation(online_hosts, payload, False)
operation_id = op[0]
results = op[1]
to_return = []
@ -197,7 +173,7 @@ class CapsulFlaskHub(VirtualizationInterface):
memory_mb=memory_mb,
ssh_public_keys=ssh_public_keys,
))
op = await self.generic_operation(online_hosts, payload, False)
op = self.generic_operation(online_hosts, payload, False)
operation_id = op[0]
results = op[1]
number_of_assigned = 0
@ -230,7 +206,7 @@ class CapsulFlaskHub(VirtualizationInterface):
host = get_model().host_of_capsul(id)
if host is not None:
payload = json.dumps(dict(type="destroy", id=id))
op = await self.generic_operation([host], payload, True)
op = self.generic_operation([host], payload, True)
results = op[1]
result_json_string = "<no response from host>"
for result in results:

3
capsulflask/payment.py
View File

@ -20,7 +20,8 @@ from werkzeug.exceptions import abort
from capsulflask.auth import account_required
from capsulflask.db import get_model, my_exec_info_message
from capsulflask.db import get_model
from capsulflask.shared import my_exec_info_message
bp = Blueprint("payment", __name__, url_prefix="/payment")

4
capsulflask/schema_migrations/09_up_introduce_hosts.sql
View File

@ -7,7 +7,7 @@ CREATE TABLE hosts (
token TEXT NOT NULL
);
INSERT INTO hosts (id, token) VALUES ('baikal', 'changeme');
INSERT INTO hosts (id, https_url, token) VALUES ('baikal', 'http://localhost:5000', 'changeme');
ALTER TABLE vms
ADD COLUMN host TEXT REFERENCES hosts(id) ON DELETE RESTRICT DEFAULT 'baikal';
@ -16,7 +16,7 @@ CREATE TABLE operations (
id SERIAL PRIMARY KEY ,
email TEXT REFERENCES accounts(email) ON DELETE RESTRICT,
created TIMESTAMP NOT NULL DEFAULT NOW(),
payload TEXT NOT NULL,
payload TEXT NOT NULL
);
CREATE TABLE host_operation (

44
capsulflask/shared.py Normal file
View File

@ -0,0 +1,44 @@
import re
from flask import current_app
class OnlineHost:
def __init__(self, id: str, url: str):
self.id = id
self.url = url
class VirtualMachine:
def __init__(self, id, host, ipv4=None, ipv6=None):
self.id = id
self.host = host
self.ipv4 = ipv4
self.ipv6 = ipv6
class VirtualizationInterface:
def capacity_avaliable(self, additional_ram_bytes: int) -> bool:
pass
def get(self, id: str) -> VirtualMachine:
pass
def list_ids(self) -> list:
pass
def create(self, email: str, id: str, template_image_file_name: str, vcpus: int, memory: int, ssh_public_keys: list):
pass
def destroy(self, email: str, id: str):
pass
def validate_capsul_id(id):
if not re.match(r"^(cvm|capsul)-[a-z0-9]{10}$", id):
raise ValueError(f"vm id \"{id}\" must match \"^capsul-[a-z0-9]{{10}}$\"")
def authorized_as_hub(headers):
if headers.get('Authorization'):
auth_header_value = headers.get('Authorization').replace("Bearer ", "")
return auth_header_value == current_app.config["HUB_TOKEN"]
return False
def my_exec_info_message(exec_info):
return "{}: {}".format(".".join([exec_info[0].__module__, exec_info[0].__name__]), exec_info[1])

14
capsulflask/spoke_api.py
View File

@ -7,17 +7,13 @@ from flask import request
from flask.json import jsonify
from werkzeug.exceptions import abort
from capsulflask.db import my_exec_info_message
from capsulflask.shared import my_exec_info_message, authorized_as_hub
bp = Blueprint("spoke", __name__, url_prefix="/spoke")
def authorized_as_hub(id):
auth_header_value = request.headers.get('Authorization').replace("Bearer ", "")
return auth_header_value == current_app.config["HUB_TOKEN"]
@bp.route("/heartbeat", methods=("POST"))
@bp.route("/heartbeat", methods=("POST",))
def heartbeat():
if authorized_as_hub(id):
if authorized_as_hub(request.headers):
url = f"{current_app.config['HUB_URL']}/hub/heartbeat/{current_app.config['SPOKE_HOST_ID']}"
authorization_header = f"Bearer {current_app.config['SPOKE_HOST_TOKEN']}"
result = current_app.config['HTTP_CLIENT'].post_json_sync(url, body=None, authorization_header=authorization_header)
@ -36,9 +32,9 @@ def heartbeat():
current_app.logger.info(f"/hosts/heartbeat returned 401: invalid hub token")
return abort(401, "invalid hub token")
@bp.route("/operation", methods=("POST"))
@bp.route("/operation", methods=("POST",))
def operation():
if authorized_as_hub(id):
if authorized_as_hub(request.headers):
request_body = request.json()
handlers = {
"capacity_avaliable": handle_capacity_avaliable,

5
capsulflask/spoke_model.py
View File

@ -8,11 +8,8 @@ from subprocess import run
from capsulflask.db import get_model
from capsulflask.hub_model import VirtualizationInterface, VirtualMachine
from capsulflask.shared import VirtualizationInterface, VirtualMachine, validate_capsul_id
def validate_capsul_id(id):
if not re.match(r"^(cvm|capsul)-[a-z0-9]{10}$", id):
raise ValueError(f"vm id \"{id}\" must match \"^capsul-[a-z0-9]{{10}}$\"")
class MockSpoke(VirtualizationInterface):
def capacity_avaliable(self, additional_ram_bytes):