infra-as-code/service.py

289 lines
9.2 KiB
Python
Executable File

#!/usr/bin/env python3
from __future__ import annotations
import argparse
import dataclasses
import sys
from os import path
from pathlib import Path
import os
import time
from typing import Any, DefaultDict, Iterable, Mapping
from jinja2 import Template
import itertools
import subprocess as sub
import shutil
import string
import secrets
from dataclasses import dataclass
TEMPLATE_DIR = "./templates"
#region util
def gen_pass() -> str:
alphabet = string.ascii_letters + string.digits
password = "".join(secrets.choice(alphabet) for _ in range(20))
return password
def check_positive(value: str):
ivalue = int(value)
if ivalue <= 0:
raise Exception("Supplied number must be >= 0")
return ivalue
def encode_member(member: str, mapping: Mapping[str, Any]) -> str:
return member + " " + " ".join([f"{k}={v}" for k, v in mapping.items()])
def iter_ips(ip_format: str, start_octet: int):
ip_int = start_octet
ip_fmt = ip_format
while ip_int < 255:
yield ip_fmt.format(ip_int)
ip_int += 1
def copy_template(src: str, dest: str, mapping: Mapping[str, Any] = {}):
c = Path(src).read_text()
t: Template = Template(c)
r = t.render(mapping)
Path(dest).write_text(r)
#endregion
#region models
@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))
class InventoryWriter:
def __init__(self, location: str) -> None:
self._file_handle = Path(location)
self._groups: dict[str, set[str]] = DefaultDict(set)
def add(self, name: str, members: Iterable[str]):
self._groups[name] |= set(members)
def _build_group(self, name: str, members: set[str]):
fmt = f"[{name}]\n" + "\n".join(members)
return fmt
def flush(self):
txt = ""
for name, members in self._groups.items():
txt += self._build_group(name, members) + "\n\n"
self._file_handle.write_text(txt, encoding="utf8")
#endregion
#region CLI positional flows
def list_envs(args: argparse.Namespace):
try:
customer_path = path.join("customers", args.customer_name, "envs")
print(" ".join(os.listdir(customer_path)))
except FileNotFoundError:
raise Exception(f"Customer `{args.customer_name}` does not exist.")
def delete_env(args: argparse.Namespace):
for env in args.env_names:
env_path = path.join("customers", args.customer_name, "envs", env)
sub.call(["vagrant", "destroy", "-f"], cwd=env_path)
shutil.rmtree(env_path)
print(f"Deleted `{env}` from customer `{ args.customer_name}`")
def recover_env(args: argparse.Namespace):
for env in args.env_names:
env_path = path.join("customers", args.customer_name, "envs", env)
sub.call(["vagrant", "up"], cwd=env_path)
# Artificial sleep waiting for sshd on VM to correctly start
print("Waiting on virtual machines...")
time.sleep(5)
sub.call(["ansible-playbook", "../../../../site.yml"], cwd=env_path)
print(f"Recovered `{env}` from customer `{ args.customer_name}`")
def create_env(args: argparse.Namespace):
if (args.num_nginx_web + args.num_nginx_lb + args.num_postgres) == 0:
raise Exception("At least one item should be deployed")
env_path = path.join("customers", args.customer_name, "envs", args.env_name)
Path(env_path).mkdir(exist_ok=True, parents=True)
vagrant_mapping: dict[str, Any] = {
"webserver_specs": None,
"loadbalancers_specs": None,
"postgres_specs": None,
"num_webserver": args.num_nginx_web,
"num_loadbalancers": args.num_nginx_lb,
"num_postgres": args.num_postgres,
"env": args.env_name,
"customer_name": args.customer_name,
"ip_int": args.ip_int,
"ip_format": args.ip_format.replace("{}", "%d"),
}
if args.num_nginx_web > 0:
print("\nNginx webserver resources:")
web_specs = MachineResources.from_prompt()
vagrant_mapping["webserver_specs"] = dataclasses.asdict(web_specs)
if args.num_nginx_lb > 0:
print("\nNginx loadbalancer resources: ")
lb_specs = MachineResources.from_prompt()
vagrant_mapping["loadbalancers_specs"] = dataclasses.asdict(lb_specs)
if args.num_postgres > 0:
print("\nPostgresql machine resources: ")
psql_specs = MachineResources.from_prompt()
vagrant_mapping["postgres_specs"] = dataclasses.asdict(psql_specs)
# Template `ansible.cfg`
src = path.join(TEMPLATE_DIR, "ansible.cfg.template")
dest = path.join(env_path, "ansible.cfg")
copy_template(src=src, dest=dest)
# Create inventory file
inv_path = path.join(env_path, "inventory")
iw = InventoryWriter(inv_path)
ip_generator = iter_ips(args.ip_format, args.ip_int)
web_ips = itertools.islice(ip_generator, args.num_nginx_web)
iw.add("webserver", web_ips)
lb_ips = itertools.islice(ip_generator, args.num_nginx_lb)
iw.add("loadbalancer", lb_ips)
def psql_gen_pass(x: str) -> str:
return encode_member(x, {"psql_pass": gen_pass()})
psql_ips = list(itertools.islice(ip_generator, args.num_postgres))
psql_ips = map(psql_gen_pass, psql_ips)
iw.add("postgresql", psql_ips)
iw.flush()
# Template `Vagrantfile`
src = path.join(TEMPLATE_DIR, "Vagrantfile.template")
dest = path.join(env_path, "Vagrantfile")
should_reload = Path(dest).exists()
copy_template(src=src, dest=dest, mapping=vagrant_mapping)
# Generate .ssh if it doesn't exist already
ssh_dir = path.join(env_path, ".ssh")
Path(ssh_dir).mkdir(exist_ok=True)
rsa_path = path.join(ssh_dir, "id_rsa")
if not Path(rsa_path).exists():
print(end="\n")
sub.call([
"ssh-keygen",
"-t",
"rsa",
"-b",
"2048",
"-f",
rsa_path,
])
# Provision and configure machines
# Create VM's that do not exist yet
sub.call(["vagrant", "up"], cwd=env_path)
# Update VM's that already existed
if should_reload:
sub.call(["vagrant", "reload", "--provision"], cwd=env_path)
# Artificial sleep waiting for sshd on VM to correctly start
print("Waiting on virtual machines...")
time.sleep(5)
sub.call(["ansible-playbook", "../../../../site.yml"], cwd=env_path)
#endregion
def main() -> int:
parser = argparse.ArgumentParser()
sub_parser = parser.add_subparsers()
# CLI definition for positional arg "list"
list_parser = sub_parser.add_parser("list", help="list customer-owned environments")
list_parser.add_argument("customer_name", type=str, help="name of the customer")
list_parser.set_defaults(func=list_envs)
# CLI definition for positional arg "create"
cenv_parser = sub_parser.add_parser("create", help="create a new environment")
cenv_parser.add_argument("customer_name", type=str, help="name of the customer")
cenv_parser.add_argument("env_name", type=str, help="name of the environment")
cenv_parser.add_argument(
"--num-postgres",
type=check_positive,
help="number of postgres databases",
default=0,
)
cenv_parser.add_argument(
"--num-nginx-web",
type=check_positive,
help="number of nginx webservers",
default=0,
)
cenv_parser.add_argument(
"--num-nginx-lb",
type=check_positive,
help="number of nginx loadbalancers",
default=0,
)
cenv_parser.add_argument(
"--ip-format", type=str, help="format of ip", default="192.168.56.{}"
)
cenv_parser.add_argument(
"--ip-int", type=check_positive, help="4th octet to start at", default=10
)
cenv_parser.set_defaults(func=create_env)
# CLI definition for positional arg "delete"
denv_parser = sub_parser.add_parser("delete", help="delete an environment")
denv_parser.add_argument("customer_name", type=str, help="name of the customer")
denv_parser.add_argument(
"env_names", type=str, nargs="+", help="name of one or more environments"
)
denv_parser.set_defaults(func=delete_env)
# CLI definition for positional arg "recover"
denv_parser = sub_parser.add_parser("recover", help="attempts to recover an env")
denv_parser.add_argument("customer_name", type=str, help="name of the customer")
denv_parser.add_argument(
"env_names", type=str, nargs="+", help="name of one or more environments"
)
denv_parser.set_defaults(func=recover_env)
# Parse args
args = parser.parse_args(sys.argv[1:])
args.func(args)
return 0
if __name__ == "__main__":
raise SystemExit(main())