capsul-flask/capsulflask/operation_model.py

import subprocess
import re
import requests
import json
import asyncio
from typing import List

from aiohttp import ClientSession
from flask import current_app
from time import sleep
from os.path import join
from subprocess import run

from capsulflask.db_model import OnlineHost
from capsulflask.db import get_model

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 VirtualMachine:
  def __init__(self, id, ipv4=None, ipv6=None):
    self.id = id
    self.ipv4 = ipv4
    self.ipv6 = ipv6

class OperationInterface:
  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

class MockVirtualization(OperationInterface):
  def capacity_avaliable(self, additional_ram_bytes):
    return True

  def get(self, id):
    validate_capsul_id(id)
    return VirtualMachine(id, ipv4="1.1.1.1")

  def list_ids(self) -> list:
    return get_model().all_non_deleted_vm_ids()

  def create(self, email: str, id: str, template_image_file_name: str, vcpus: int, memory_mb: int, ssh_public_keys: list):
    validate_capsul_id(id)
    current_app.logger.info(f"mock create: {id} for {email}")
    sleep(1)

  def destroy(self, email: str, id: str):
    current_app.logger.info(f"mock destroy: {id} for {email}")


class GoshtOperation(OperationInterface):

  
  async def post_json(self, method: str, url: str, body: str, session: ClientSession, **kwargs) -> str:
      resp = await session.request(method=method, url=url, json=body, **kwargs)
      resp.raise_for_status()
      return await resp.text()

  async def make_requests(self, online_hosts: List[OnlineHost], body: str, **kwargs) -> None:
      async with ClientSession() as session:
          tasks = []
          for host in online_hosts:
              tasks.append(
                  self.post_json(method="POST", url=host.url, body=body, session=session, **kwargs)
              )
          results = await asyncio.gather(*tasks)
          # do something with results

  def validate_completed_process(self, completedProcess, email=None):
    emailPart = ""
    if email != None:
      emailPart = f"for {email}"

    if completedProcess.returncode != 0:
      raise RuntimeError(f"""{" ".join(completedProcess.args)} failed {emailPart} with exit code {completedProcess.returncode}
        stdout:
        {completedProcess.stdout}
        stderr:
        {completedProcess.stderr}
      """)



  def capacity_avaliable(self, additional_ram_bytes):

    online_hosts = get_model().get_online_hosts()
    payload = json.dumps(dict(type="capacity_avaliable", additional_ram_bytes=""))
    get_model().create_operation(online_hosts, None, payload)


    
    

  def get(self, id):
    validate_capsul_id(id)
    completedProcess = run([join(current_app.root_path, 'shell_scripts/get.sh'), id], capture_output=True)
    self.validate_completed_process(completedProcess)
    lines = completedProcess.stdout.splitlines()
    ipaddr = lines[0].decode("utf-8")

    if not re.match(r"^([0-9]{1,3}\.){3}[0-9]{1,3}$", ipaddr):
      return None

    return VirtualMachine(id, ipv4=ipaddr)

  def list_ids(self) -> list:
    completedProcess = run([join(current_app.root_path, 'shell_scripts/list-ids.sh')], capture_output=True)
    self.validate_completed_process(completedProcess)
    return list(map(lambda x: x.decode("utf-8"), completedProcess.stdout.splitlines() ))

  def create(self, email: str, id: str, template_image_file_name: str, vcpus: int, memory_mb: int, ssh_public_keys: list):
    validate_capsul_id(id)

    if not re.match(r"^[a-zA-Z0-9/_.-]+$", template_image_file_name):
      raise ValueError(f"template_image_file_name \"{template_image_file_name}\" must match \"^[a-zA-Z0-9/_.-]+$\"")

    for ssh_public_key in ssh_public_keys:
      if not re.match(r"^(ssh|ecdsa)-[0-9A-Za-z+/_=@. -]+$", ssh_public_key):
        raise ValueError(f"ssh_public_key \"{ssh_public_key}\" must match \"^(ssh|ecdsa)-[0-9A-Za-z+/_=@. -]+$\"")

    if vcpus < 1 or vcpus > 8:
      raise ValueError(f"vcpus ({vcpus}) must match 1 <= vcpus <= 8")

    if memory_mb < 512 or memory_mb > 16384:
      raise ValueError(f"memory_mb ({memory_mb}) must match 512 <= memory_mb <= 16384")

    ssh_keys_string = "\n".join(ssh_public_keys)

    completedProcess = run([
      join(current_app.root_path, 'shell_scripts/create.sh'),
      id,
      template_image_file_name,
      str(vcpus),
      str(memory_mb),
      ssh_keys_string
    ], capture_output=True)

    self.validate_completed_process(completedProcess, email)
    lines = completedProcess.stdout.splitlines()
    status = lines[len(lines)-1].decode("utf-8")

    vmSettings = f"""
      id={id}
      template_image_file_name={template_image_file_name}
      vcpus={str(vcpus)}
      memory={str(memory_mb)}
      ssh_public_keys={ssh_keys_string}
    """

    if not status == "success":
      raise ValueError(f"""failed to create vm for {email} with:
        {vmSettings}
        stdout:
        {completedProcess.stdout}
        stderr:
        {completedProcess.stderr}
      """)

  def destroy(self, email: str, id: str):
    validate_capsul_id(id)
    completedProcess = run([join(current_app.root_path, 'shell_scripts/destroy.sh'), id], capture_output=True)
    self.validate_completed_process(completedProcess, email)
    lines = completedProcess.stdout.splitlines()
    status = lines[len(lines)-1].decode("utf-8")

    if not status == "success":
      raise ValueError(f"""failed to destroy vm "{id}" for {email}:
        stdout:
        {completedProcess.stdout}
        stderr:
        {completedProcess.stderr}
      """)