dashboard/infrastructure/management/commands/vm.py

# Copyright © 2022 Aravinth Manivannan <realaravinth@batsense.net>
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as
# published by the Free Software Foundation, either version 3 of the
# License, or (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
from enum import Enum, unique

from django.core.management.base import BaseCommand
from django.core.exceptions import ValidationError
from django.conf import settings
from django.contrib.auth import get_user_model

from oauth2_provider.models import get_application_model
from oauth2_provider.generators import generate_client_id, generate_client_secret

from dash.models import InstanceConfiguration, Instance
from dash.utils import create_instance, VmException, VmErrors
from infrastructure.utils import create_vm_if_not_exists


def translate_sizes(flavor: str):
    if flavor == "small":
        size = "s1-2"
    elif flavor == "medium":
        size = "s1-4"
    elif flavor == "large":
        size = "s1-8"
    else:
        print("flavour no match")
        size = flavor
    return size


@unique
class Actions(Enum):
    CREATE = "create"
    DELETE = "delete"

    def __str__(self) -> str:
        return self.name


class Command(BaseCommand):
    help = "Get user ID from username"
    action_key = "action"
    vm_name_key = "vm_name"
    flavor_key = "flavor"
    owner_key = "owner"

    def add_arguments(self, parser):
        parser.add_argument(
            self.action_key,
            type=Actions,
            help="VM action: create/delete",
        )

        parser.add_argument(
            self.vm_name_key,
            type=str,
            help="Name of the VM",
        )

        parser.add_argument(
            f"--{self.owner_key}",
            type=str,
            help="Owner username",
        )

        parser.add_argument(
            f"--{self.flavor_key}",
            type=str,
            help="Name of the VM flavor: small, medium, large",
        )

    def create_vm(self, *args, **options):
        owner = options[self.owner_key]
        flavor = options[self.flavor_key]
        vm_name = options[self.vm_name_key]

        size = translate_sizes(flavor)

        user = get_user_model().objects.get(username=owner)
        try:
            instance = create_instance(
                vm_name=vm_name, configuration_name=size, user=user
            )
            (gitea_password, _commit) = create_vm_if_not_exists(instance)
            print("Instance created")
            print(f"Gitea admin password: {gitea_password}")
        except VmException as e:
            if e.code == VmErrors.NAME_EXISTS:
                instance = Instance.objects.get(name=vm_name)
                if instance.configuration_id.name != size:
                    instance.configuration_id = InstanceConfiguration.objects.get(
                        name=size
                    )
                    instance.save()
                    (gitea_password, _commit) = create_vm_if_not_exists(instance)
                    print("Instance created")
                    print(f"Gitea admin password: {gitea_password}")

            else:
                self.stderr.write(self.style.ERROR(f"error: {str(e)}"))
        except Exception as e:
            self.stderr.write(self.style.ERROR(f"error: {str(e)}"))

    def delete_vm(self, *args, **options):
        from infrastructure.utils import delete_vm

        vm_name = options[self.vm_name_key]
        if Instance.objects.filter(name=vm_name).exists():
            instance = Instance.objects.get(name=vm_name)
            delete_vm(instance=instance)

    def handle(self, *args, **options):
        for i in [self.action_key, self.vm_name_key]:
            if i not in options:
                self.stdout.write(self.style.ERROR(f"Please provide {i}"))
                return
        if options[self.action_key] == Actions.CREATE:
            for i in [self.flavor_key, self.owner_key]:
                if i not in options:
                    self.stdout.write(self.style.ERROR(f"Please provide {i}"))
                    return

            self.create_vm(*args, **options)
        elif options[self.action_key] == Actions.DELETE:
            self.delete_vm(*args, **options)

        else:
            self.stdout.write(
                self.style.ERROR(f"Unknown action: {options[self.action_key]}")
            )
            return