capsul-flask/capsulflask/virt_model.py

import subprocess
import re

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

class VirtualMachine:
  def __init__(self, id, ipv4=None, ipv6=None):
    self.id = id
    self.ipv4 = ipv4
    self.ipv6 = ipv6

class VirtualizationInterface:
  def get(self, id: str) -> VirtualMachine:
    pass

  def listIds(self) -> list:
    pass

  def create(self, id: str, template_file_name: str, memory: int, vcpus: int, ssh_public_keys: list) -> VirtualMachine:
    pass

  def destroy(self, id: str):
    pass

class ShellScriptVirtualization(VirtualizationInterface):

  def validateId(self, id):
    if not re.match(r"^capsul-[a-z0-9]{10}$", id):
      raise ValueError(f"vm id \"{id}\" must match \"^capsul-[a-z0-9]{{10}}$\"")

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

  def get(self, id):
    self.validateId(id)
    completedProcess = run([join(current_app.root_path, 'shell_scripts/get.sh'), id], capture_output=True)
    self.validateCompletedProcess(completedProcess)
    lines = completedProcess.stdout.splitlines()

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

    return VirtualMachine(id, ipv4=lines[0])

  def listIds(self) -> list:
    completedProcess = run([join(current_app.root_path, 'shell_scripts/listIds.sh')], capture_output=True)
    self.validateCompletedProcess(completedProcess)
    return completedProcess.stdout.splitlines()

  def create(self, id: str, template_file_name: str, vcpus: int, memory_mb: int, ssh_public_keys: list):
    self.validateId(id)

    if not re.match(r"^[a-zA-Z0-9_.-]$", template_file_name):
      raise ValueError(f"template_file_name \"{template_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_file_name,
      str(vcpus),
      str(memory_mb),
      ssh_keys_string
    ], capture_output=True)

    self.validateCompletedProcess(completedProcess)
    lines = completedProcess.stdout.splitlines()

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

    if not lines[len(lines)-1] == "success":
      raise ValueError(f"""failed to create vm with:
        {vmSettings}
        stdout:
        {completedProcess.stdout}
        stderr:
        {completedProcess.stderr}
      """)
    
    for _ in range(0, 10):
      sleep(6)
      result = self.get(id)
      if result != None:
        return result

    for _ in range(0, 10):
      sleep(60)
      result = self.get(id)
      if result != None:
        return result
    
    raise TimeoutError(f"""timed out waiting for vm to obtain an IP address:
        {vmSettings}
    """)

  def destroy(self, id: str):
    self.validateId(id)
    completedProcess = run([join(current_app.root_path, 'shell_scripts/destroy.sh'), id], capture_output=True)
    self.validateCompletedProcess(completedProcess)
    lines = completedProcess.stdout.splitlines()

    if not lines[len(lines)-1] == "success":
      raise ValueError(f"""failed to destroy vm "{id}":
        stdout:
        {completedProcess.stdout}
        stderr:
        {completedProcess.stderr}
      """)