capsul-flask/capsulflask/console.py

import re
import sys
import json
from datetime import datetime, timedelta
from flask import Blueprint
from flask import flash
from flask import current_app
from flask import g
from flask import request
from flask import session
from flask import render_template
from flask import redirect
from flask import url_for
from werkzeug.exceptions import abort
from nanoid import generate

from capsulflask.metrics import durations as metric_durations
from capsulflask.auth import account_required
from capsulflask.db import get_model
from capsulflask.shared import my_exec_info_message, get_vm_months_float, get_account_balance, average_number_of_days_in_a_month
from capsulflask.payment import poll_btcpay_session, check_if_shortterm_flag_can_be_unset
from capsulflask import cli

bp = Blueprint("console", __name__, url_prefix="/console")

def make_capsul_id():
  letters_n_nummers = generate(alphabet="1234567890qwertyuiopasdfghjklzxcvbnm", size=10)
  return f"capsul-{letters_n_nummers}"

def double_check_capsul_address(id, existing_ipv4, get_ssh_host_keys):
  try:
    result = current_app.config["HUB_MODEL"].get(id, get_ssh_host_keys)

    if result != None and result.ipv4 != "" and (existing_ipv4 == None or existing_ipv4 == ""):
      get_model().update_vm_ip(email=session["account"], id=id, ipv4=result.ipv4)

    if result != None and result.ssh_host_keys != None and get_ssh_host_keys:
      get_model().update_vm_ssh_host_keys(email=session["account"], id=id, ssh_host_keys=result.ssh_host_keys)
  except:
    current_app.logger.error(f"""
      the virtualization model threw an error in double_check_capsul_address of {id}:
      {my_exec_info_message(sys.exc_info())}"""
    )
    return None

  return result

@bp.route("/")
@account_required
def index():
  vms = get_vms()
  vms = list(filter(lambda x: not x['deleted'], vms))
  created = request.args.get('created')

  # this is here to prevent xss
  if created and not re.match(r"^(cvm|capsul)-[a-z0-9]{10}$", created):
    created = '___________'

  # for now we are going to check the IP according to the virt model
  # on every request. this could be done by a background job and cached later on...
  for vm in vms:
    result = double_check_capsul_address(vm["id"], vm["ipv4"], False)
    if result is not None:
      vm["ipv4"] = result.ipv4
      vm["state"] = result.state
    else:
      vm["state"] = "unknown"


  mappedVms = []
  for vm in vms:
    ip_display = vm['ipv4']
    if not ip_display:
      if vm["state"] == "running":
        ip_display = "..booting.."
      else:
        ip_display = "unknown"

    ip_display_class = "ok"
    if not vm['ipv4']:
      if vm["state"] == "running":
        ip_display_class = "waiting-pulse"
      else:
        ip_display_class = "yellow"

    mappedVms.append(dict(
      id=vm['id'], 
      size=vm['size'], 
      state=vm['state'],
      ipv4=ip_display, 
      ipv4_status=ip_display_class, 
      os=vm['os'], 
      created=vm['created'].strftime("%b %d %Y")
    ))

  return render_template("capsuls.html", vms=mappedVms, has_vms=len(vms) > 0, created=created)

@bp.route("/<string:id>", methods=("GET", "POST"))
@account_required
def detail(id):

  duration=request.args.get('duration')
  if not duration:
    duration = "5m"

  vm = get_model().get_vm_detail(email=session["account"], id=id)

  if vm is None:
    return abort(404, f"{id} doesn't exist.")

  if vm['deleted']:
    return render_template("capsul-detail.html", vm=vm, delete=True, deleted=True)

  vm["created"] = vm['created'].strftime("%b %d %Y %H:%M")
  vm["ssh_authorized_keys"] = ", ".join(vm["ssh_authorized_keys"]) if len(vm["ssh_authorized_keys"]) > 0 else "<missing>"

  if request.method == "POST":
    if "csrf-token" not in request.form or request.form['csrf-token'] != session['csrf-token']:
      return abort(418, f"u want tea")

    if 'action' not in request.form:
      return abort(400, "action is required")

    if request.form['action'] == "start":
      get_model().set_desired_state(email=session["account"], vm_id=id, desired_state="running")
      current_app.config["HUB_MODEL"].vm_state_command(email=session['account'], id=id, command="start")

      vm["state"] = "starting"
      return render_template("capsul-detail.html", vm=vm)
    elif request.form['action'] == "delete":
      if 'are_you_sure' not in request.form or not request.form['are_you_sure']:
        return render_template(
          "capsul-detail.html", 
          csrf_token = session["csrf-token"],
          vm=vm, 
          delete=True
        )
      else:
        current_app.logger.info(f"deleting {vm['id']} per user request ({session['account']})")
        current_app.config["HUB_MODEL"].destroy(email=session['account'], id=id)
        get_model().delete_vm(email=session['account'], id=id)
        
        # now that the user has 1 less vm, check if thier account can fund all thier vms for 
        # 1 month and if so, set all thier vms to longterm
        check_if_shortterm_flag_can_be_unset(session['account'])

        return render_template("capsul-detail.html", vm=vm, deleted=True)
    elif request.form['action'] == "force-stop":
      if 'are_you_sure' not in request.form or not request.form['are_you_sure']:
        return render_template(
          "capsul-detail.html", 
          csrf_token = session["csrf-token"],
          vm=vm, 
          force_stop=True,
        )
      else:
        current_app.logger.info(f"force stopping {vm['id']} per user request ({session['account']})")
        get_model().set_desired_state(email=session["account"], vm_id=id, desired_state="shut off")
        current_app.config["HUB_MODEL"].vm_state_command(email=session['account'], id=id, command="force-stop")

        vm["state"] = "stopped"

        flash("please note that your account will still be billed for this capsul while it is in a stopped state")

        return render_template(
          "capsul-detail.html",
          csrf_token = session["csrf-token"],
          vm=vm,
          durations=list(map(lambda x: x.strip("_"), metric_durations.keys())),
          duration=duration
        )
    else:
      return abort(400, "action must be either delete, force-stop, or start")


  else:
    needs_ssh_host_keys = "ssh_host_keys" not in vm or len(vm["ssh_host_keys"]) == 0

    vm_from_virt_model = double_check_capsul_address(vm["id"], vm["ipv4"], needs_ssh_host_keys)
    
    if vm_from_virt_model is not None:
      vm["ipv4"] = vm_from_virt_model.ipv4
      vm["state"] = vm_from_virt_model.state
      if needs_ssh_host_keys:
        vm["ssh_host_keys"] = vm_from_virt_model.ssh_host_keys
    else:
      vm["state"] = "unknown"

    if vm["state"] == "running" and not vm["ipv4"]:
      vm["state"] = "starting"

    return render_template(
      "capsul-detail.html", 
      csrf_token = session["csrf-token"],
      vm=vm, 
      durations=list(map(lambda x: x.strip("_"), metric_durations.keys())),
      duration=duration
    )


@bp.route("/create", methods=("GET", "POST"))
@account_required
def create():
  vm_sizes = get_model().vm_sizes_dict()
  operating_systems = get_model().operating_systems_dict()
  public_keys_for_account = get_model().list_ssh_public_keys_for_account(session["account"])
  account_balance = get_account_balance(get_vms(), get_payments(), datetime.utcnow())
  account_balance_in_one_month = get_account_balance(get_vms(), get_payments(), datetime.utcnow() + timedelta(days=average_number_of_days_in_a_month))
  capacity_avaliable = current_app.config["HUB_MODEL"].capacity_avaliable(512*1024*1024)
  errors = list()

  month_funded_vm_sizes = dict()
  hour_funded_vm_sizes = dict()
  for key, vm_size in vm_sizes.items():
    # if a user deposits $7.50 and then creates an f1-s vm which costs 7.50 a month, 
    # then they have to delete the vm and re-create it, they will not be able to, they will have to pay again.
    # so for UX it makes a lot of sense to give a small margin of 25 cents for usability sake
    if vm_size["dollars_per_month"] <= account_balance_in_one_month+0.25:
      month_funded_vm_sizes[key] = vm_size

    one_month_in_hours = float(730.5)
    two_hours_as_fraction_of_month = float(2)/one_month_in_hours
    if float(vm_size["dollars_per_month"])*two_hours_as_fraction_of_month <= account_balance:
      hour_funded_vm_sizes[key] = vm_size

  if request.method == "POST":
    if "csrf-token" not in request.form or request.form['csrf-token'] != session['csrf-token']:
      return abort(418, f"u want tea")

    size = request.form["size"]
    os = request.form["os"]
    if not size:
      errors.append("Size is required")
    elif size not in vm_sizes:
      errors.append(f"Invalid size {size}")
    elif size not in hour_funded_vm_sizes:
      errors.append(f"Your account must have enough credit to run an {size} for 2 hours before you will be allowed to create it")

    if not os:
      errors.append("OS is required")
    elif os not in operating_systems:
      errors.append(f"Invalid os {os}")

    posted_keys_count = int(request.form["ssh_authorized_key_count"])
    posted_keys = list()

    if posted_keys_count > 1000:
      errors.append("something went wrong with ssh keys")
    else:
      for i in range(0, posted_keys_count):
        if f"ssh_key_{i}" in request.form:
          posted_name = request.form[f"ssh_key_{i}"]
          key = None
          for x in public_keys_for_account:
            if x['name'] == posted_name:
              key = x
          if key:
            posted_keys.append(key)
          else:
            errors.append(f"SSH Key \"{posted_name}\" doesn't exist")

    if len(posted_keys) == 0:
      errors.append("At least one SSH Public Key is required")

    capacity_avaliable = current_app.config["HUB_MODEL"].capacity_avaliable(vm_sizes[size]['memory_mb']*1024*1024)

    if not capacity_avaliable:
      errors.append("""
        host(s) at capacity. no capsuls can be created at this time. sorry. 
      """)

    if len(errors) == 0:
      id = make_capsul_id()
      # we can't create the vm record in the DB yet because its IP address needs to be allocated first.
      # so it will be created when the allocation happens inside the hub_api.
      current_app.config["HUB_MODEL"].create(
        email = session["account"],
        id=id,
        os=os,
        size=size,
        shortterm=(size not in month_funded_vm_sizes),
        template_image_file_name=operating_systems[os]['template_image_file_name'],
        vcpus=vm_sizes[size]['vcpus'],
        memory_mb=vm_sizes[size]['memory_mb'],
        ssh_authorized_keys=list(map(lambda x: dict(name=x['name'], content=x['content']), posted_keys))
      )
      
      return redirect(f"{url_for('console.index')}?created={id}")
  

  for error in errors:
    flash(error)

  if not capacity_avaliable:
    current_app.logger.warning(f"when capsul capacity is restored, send an email to {session['account']}")
  
  return render_template(
    "create-capsul.html",
    csrf_token = session["csrf-token"],
    capacity_avaliable=capacity_avaliable,
    account_balance=format(account_balance, '.2f'),
    account_balance_in_one_month=format(account_balance_in_one_month, '.2f'),
    ssh_authorized_keys=public_keys_for_account,
    ssh_authorized_key_count=len(public_keys_for_account),
    no_ssh_public_keys=len(public_keys_for_account) == 0,
    operating_systems=operating_systems,
    cant_afford=len(hour_funded_vm_sizes) == 0,
    vm_sizes=hour_funded_vm_sizes,
    month_funded_vm_sizes=month_funded_vm_sizes
  )

@bp.route("/ssh", methods=("GET", "POST"))
@account_required
def ssh_public_keys():
  errors = list()

  if request.method == "POST":
    if "csrf-token" not in request.form or request.form['csrf-token'] != session['csrf-token']:
      return abort(418, f"u want tea")
      
    method = request.form["method"]
    content = None
    if method == "POST":
      content = request.form["content"].replace("\r", " ").replace("\n", " ").strip()
      
    name = request.form["name"]
    if not name or len(name.strip()) < 1:
      if method == "POST":
        parts = re.split(" +", content)
        if len(parts) > 2 and len(parts[2].strip()) > 0:
          name = parts[2].strip()
        else:
          name = parts[0].strip()
      else:
        errors.append("Name is required")
    if not re.match(r"^[0-9A-Za-z_@:. -]+$", name):
      errors.append(f"Key name '{name}' must match \"^[0-9A-Za-z_@:. -]+$\"")

    if method == "POST":
      if not content or len(content.strip()) < 1:
        errors.append("Content is required")
      else:
        if not re.match(r"^(ssh|ecdsa)-[0-9A-Za-z+/_=@:. -]+$", content):
          errors.append("Content must match \"^(ssh|ecdsa)-[0-9A-Za-z+/_=@:. -]+$\"")

      if get_model().ssh_public_key_name_exists(session["account"], name):
        errors.append("A key with that name already exists")

      if len(errors) == 0:
        get_model().create_ssh_public_key(session["account"], name, content)

    elif method == "DELETE":

      if len(errors) == 0:
        get_model().delete_ssh_public_key(session["account"], name)

  for error in errors:
    flash(error)

  keys_list=list(map(
    lambda x: dict(name=x['name'], content=f"{x['content'][:20]}...{x['content'][len(x['content'])-20:]}"), 
    get_model().list_ssh_public_keys_for_account(session["account"])
  ))

  return render_template(
    "ssh-public-keys.html", 
    csrf_token = session["csrf-token"],
    ssh_public_keys=keys_list, 
    has_ssh_public_keys=len(keys_list) > 0
  )

def get_vms():
  if 'user_vms' not in g:
    g.user_vms = get_model().list_vms_for_account(session["account"])
  return g.user_vms

def get_payments():
  if 'user_payments' not in g:
    g.user_payments = get_model().list_payments_for_account(session["account"])
  return g.user_payments



@bp.route("/account-balance")
@account_required
def account_balance():

  payment_sessions = get_model().list_payment_sessions_for_account(session['account'])
  for payment_session in payment_sessions:
    if payment_session['type'] == 'btcpay':
      poll_btcpay_session(payment_session['id'])

  payments = get_payments()
  vms = get_vms()
  balance_1w = get_account_balance(vms, payments, datetime.utcnow() + timedelta(days=7)) 
  balance_1d = get_account_balance(vms, payments, datetime.utcnow() + timedelta(days=1)) 
  balance_now = get_account_balance(vms, payments, datetime.utcnow())

  warning_index = -1
  warning_text = ""
  warnings = cli.get_warnings_list()

  for i in range(0, len(warnings)):
    if warnings[i]['get_active'](balance_1w, balance_1d, balance_now):
      warning_index = i
  if warning_index > -1:
    pluralize_capsul = "s" if len(vms) > 1 else ""
    warning_id = warnings[warning_index]['id']
    warning_text = cli.get_warning_headline(warning_id, pluralize_capsul)

  vms_billed = list()

  for vm in get_vms():
    vm_months = get_vm_months_float(vm, datetime.utcnow())
    vms_billed.append(dict(
      id=vm["id"], 
      dollars_per_month=vm["dollars_per_month"],
      created=vm["created"].strftime("%b %d %Y"),
      deleted=vm["deleted"].strftime("%b %d %Y") if vm["deleted"] else "N/A",
      months=format(vm_months, '.3f'),
      dollars=format(vm_months * float(vm["dollars_per_month"]), '.2f')
    ))

  return render_template(
    "account-balance.html", 
    has_vms=len(vms_billed)>0, 
    vms_billed=vms_billed,
    warning_text=warning_text,
    btcpay_enabled=current_app.config["BTCPAY_URL"] != "",
    btcpay_is_working="BTCPAY_CLIENT" in current_app.config,
    payments=list(map(
      lambda x: dict(
        dollars=x["dollars"], 
        class_name="invalidated" if x["invalidated"] else "", 
        created=x["created"].strftime("%b %d %Y")
      ), 
      payments
    )), 
    has_payments=len(payments)>0, 
    account_balance=format(balance_now, '.2f')
  )