From a330bb17d1bea1dea68d7fba4955fdf7dd04af39 Mon Sep 17 00:00:00 2001
From: strNophix <nvdpoel01@gmail.com>
Date: Tue, 15 Mar 2022 23:08:54 +0100
Subject: [PATCH] Added resource mngmnt option for machine types

---
 service.py                     | 51 +++++++++++++++++++++++++++++++++-
 templates/Vagrantfile.template |  9 ++++--
 2 files changed, 56 insertions(+), 4 deletions(-)

diff --git a/service.py b/service.py
index 0444eb5..3d92f23 100755
--- a/service.py
+++ b/service.py
@@ -1,6 +1,7 @@
 #!/usr/bin/env python3
 from __future__ import annotations
 import argparse
+import dataclasses
 import sys
 from os import path
 from pathlib import Path
@@ -13,10 +14,38 @@ import subprocess as sub
 import shutil
 import string
 import secrets
+from dataclasses import dataclass
+
 
 TEMPLATE_DIR = "./templates"
 
 
+@dataclass
+class MachineResources:
+    cpus: int
+    mem: int
+
+    @staticmethod
+    def from_prompt() -> "MachineResources":
+        cpus = input("How many processors would you like to assign (default=1): ")
+        if not cpus:
+            cpus = "1"
+
+        if not cpus.isdigit() or int(cpus) < 0:
+            raise Exception("Expected a postive amount of processors")
+
+        mem = input(
+            "How many megabytes of RAM would you like to assign (default=1024): "
+        )
+        if not mem:
+            mem = "1024"
+
+        if not mem.isdigit() or int(mem) < 0:
+            raise Exception("Expected a postive amount of memory")
+
+        return MachineResources(cpus=int(cpus), mem=int(mem))
+
+
 def gen_pass() -> str:
     alphabet = string.ascii_letters + string.digits
     password = "".join(secrets.choice(alphabet) for _ in range(20))
@@ -91,6 +120,21 @@ def create_env(args: argparse.Namespace):
     env_path = path.join("customers", args.customer_name, "envs", args.env_name)
     Path(env_path).mkdir(exist_ok=True, parents=True)
 
+    web_specs = None
+    if args.num_nginx_web > 0:
+        print("\nNginx webserver resources:")
+        web_specs = MachineResources.from_prompt()
+
+    lb_specs = None
+    if args.num_nginx_lb > 0:
+        print("\nNginx loadbalancer resources: ")
+        lb_specs = MachineResources.from_prompt()
+
+    psql_specs = None
+    if args.num_postgres > 0:
+        print("\nPostgresql machine resources: ")
+        psql_specs = MachineResources.from_prompt()
+
     # Template `ansible.cfg`
     src = path.join(TEMPLATE_DIR, "ansible.cfg.template")
     dest = path.join(env_path, "ansible.cfg")
@@ -129,6 +173,9 @@ def create_env(args: argparse.Namespace):
         "num_webserver": args.num_nginx_web,
         "num_loadbalancers": args.num_nginx_lb,
         "num_postgres": args.num_postgres,
+        "webserver_specs": dataclasses.asdict(web_specs),
+        "loadbalancers_specs": dataclasses.asdict(lb_specs),
+        "postgres_specs": dataclasses.asdict(psql_specs),
     }
     copy_template(src=src, dest=dest, mapping=mapping)
 
@@ -137,6 +184,7 @@ def create_env(args: argparse.Namespace):
     Path(ssh_dir).mkdir(exist_ok=True)
     rsa_path = path.join(ssh_dir, "id_rsa")
     if not Path(rsa_path).exists():
+        print(end="\n")
         ssh_key_cmd = [
             "ssh-keygen",
             "-t",
@@ -150,7 +198,8 @@ def create_env(args: argparse.Namespace):
 
     # Provision and configure machines
     sub.call(["vagrant", "up"], cwd=env_path)
-    time.sleep(1)
+    print("Waiting on virtual machines...")
+    time.sleep(5)
     sub.call(["ansible-playbook", "../../../../site.yml"], cwd=env_path)
 
 
diff --git a/templates/Vagrantfile.template b/templates/Vagrantfile.template
index 2fef70e..e21e558 100644
--- a/templates/Vagrantfile.template
+++ b/templates/Vagrantfile.template
@@ -26,7 +26,8 @@ Vagrant.configure("2") do |config|
       web.vm.provision "file", source: "./.ssh/id_rsa.pub", destination: "~/.ssh/authorized_keys"
 
       web.vm.provider "virtualbox" do |vb|
-        vb.memory = "1024"
+        vb.memory = {{ webserver_specs["mem"] }}
+        vb.cpus = {{ webserver_specs["cpus"] }}
         vb.gui = false
         vb.name = machine_id
       end
@@ -45,7 +46,8 @@ Vagrant.configure("2") do |config|
       web.vm.provision "file", source: "./.ssh/id_rsa.pub", destination: "~/.ssh/authorized_keys"
 
       web.vm.provider "virtualbox" do |vb|
-        vb.memory = "1024"
+        vb.memory = {{ loadbalancers_specs["mem"] }}
+        vb.cpus = {{ loadbalancers_specs["cpus"] }}
         vb.gui = false
         vb.name = machine_id
       end
@@ -64,7 +66,8 @@ Vagrant.configure("2") do |config|
       web.vm.provision "file", source: "./.ssh/id_rsa.pub", destination: "~/.ssh/authorized_keys"
 
       web.vm.provider "virtualbox" do |vb|
-        vb.memory = "1024"
+        vb.memory = {{ postgres_specs["mem"] }}
+        vb.cpus = {{ postgres_specs["cpus"] }}
         vb.gui = false
         vb.name = machine_id
       end