diff options
Diffstat (limited to 'src')
| -rwxr-xr-x | src/cli/hashi_env.py | 208 | ||||
| -rwxr-xr-x | src/cli/nomad_allocs.py | 51 |
2 files changed, 259 insertions, 0 deletions
diff --git a/src/cli/hashi_env.py b/src/cli/hashi_env.py new file mode 100755 index 0000000..aeac5ef --- /dev/null +++ b/src/cli/hashi_env.py @@ -0,0 +1,208 @@ +#!/usr/bin/env python3 + +import datetime +import hashlib +import json +import os + +import click +import requests + +# the ID of the vault in 1password +op_vault = "v4mof5qwozyvob2utdk3crwxnu" + +vault_addrs = { + "chi1": "https://chi1-vault.simulprod.com:8200", + "ash1": "https://ash1-vault.simulprod.com:8200", +} + +creds_cache = os.path.join(os.getenv("HOME") or "", ".local/state/rbxenv") + + +valid_dcs = ["ash1", "chi1"] +valid_edges = [ + "bom1", + "ams1", + "atl1", + "dwf1", + "fra2", + "hkg1", + "iad1", + "lax1", + "lga1", + "lhr1", + "mia1", + "nrt1", + "ord1", + "ord2", + "sjfc1", +] + + +def _path_cached_file(val: str) -> str: + """The path to the cache file. + The full path is created using the URL to the vault server. + """ + m = hashlib.sha256() + m.update(bytes(val, "utf-8")) + val = m.hexdigest() + return os.path.join(creds_cache, f"{val}.json") + + +def get_ldap_password() -> str: + """Return the LDAP password. + The password is expected to be in 1password, under `LDAP'. + """ + return ( + os.popen(f"/usr/local/bin/op read op://{op_vault}/LDAP/password") + .read() + .rstrip() + ) + + +def _get_token_from_cache(addr: str) -> str | None: + """Return the token from the cache if it is still valid.""" + cached_path = _path_cached_file(addr) + if not os.path.isfile(cached_path): + return None + + with open(cached_path) as f: + d = json.load(f) + expires_after = datetime.datetime.fromtimestamp(int(d["until"])) + if datetime.datetime.now() > expires_after: + return None + return d["token"] + + return None + + +def _set_token_to_cache(addr: str, token: str, expires_after: datetime.datetime): + """Set the token in the cache. + The cache also contains the time after which the token won't be valid anymore.""" + if not os.path.isdir(creds_cache): + os.makedirs(creds_cache) + + cache = { + "until": int(expires_after.timestamp()), + "token": token, + } + + cached_path = _path_cached_file(addr) + + with open(cached_path, "w") as f: + json.dump(cache, f) + + +def _vault_login(addr: str) -> str: + """Log into vault to get a token. + If we get a token, we store it in the cache so we don't need to request it again until it expires. + """ + ldap_username = os.getenv("USER") + ldap_password = get_ldap_password() + + url = "{}/v1/auth/ldap/login/{}".format(addr, ldap_username) + obj = {"method": "push", "password": ldap_password} + + try: + resp = requests.post(url, json=obj) + resp.raise_for_status() + except Exception as e: + print("{} returned {}".format(url, str(e))) + + expires_after = datetime.datetime.now() + datetime.timedelta( + seconds=resp.json()["auth"]["lease_duration"] + ) + token = resp.json()["auth"]["client_token"] + + _set_token_to_cache(addr, token, expires_after) + + return token + + +def get_vault_token(addr: str) -> str: + """Get the token for vault.""" + token = _get_token_from_cache(addr) + if token is None: + token = _vault_login(addr) + return token + + +def vault_read(path: str, addr: str, token: str, dc: str) -> str: + """Read some values from a path in vault.""" + url = "{}/v1/{}".format(addr, path) + headers = {"x-vault-token": token} + try: + resp = requests.get(url, headers=headers) + resp.raise_for_status() + except Exception as e: + print("{} returned {}".format(url, str(e))) + + return resp.json()["data"][dc] + + +@click.command() +@click.argument("dc") +def setpop(dc: str): + """Print some information regarding hashi stack in a POP.""" + if dc not in valid_edges: + print(f"{dc} is not a valid edge location") + return + + token = get_vault_token(vault_addrs["chi1"]) + consul_token = vault_read( + "secret/teams/neteng/traffic/consul", + addr=vault_addrs["chi1"], + token=token, + dc=dc, + ) + nomad_token = vault_read( + "secret/teams/neteng/traffic/nomad", + addr=vault_addrs["chi1"], + token=token, + dc=dc, + ) + vault_token = vault_read( + "secret/teams/neteng/traffic/vault", + addr=vault_addrs["chi1"], + token=token, + dc=dc, + ) + + print(f"consul token: {consul_token}") + print(f"nomad token: {nomad_token}") + print(f"vault token: {vault_token}") + print(f"https://{dc}-vault.simulprod.com/ui/vault/auth?with=token") + + +@click.command() +@click.argument("dc") +def setenv(dc: str): + """Print environment variables for the URL and tokens to various components. + This command can be passed to `export` in order to export variables. + For example: + ``` + export (robloxenv setenv chi1) + ``` + """ + if dc not in valid_dcs: + print(f"{dc} is not a valid data center") + return + + if dc not in vault_addrs: + print(f"{dc} doesn't have an associated vault address") + + token = get_vault_token(vault_addrs[dc]) + print(f"VAULT_ADDR={vault_addrs[dc]}") + print(f"VAULT_TOKEN={token}") + + +@click.group(help="CLI tool to manage environment variables for hashi things at Roblox") +def cli(): + pass + + +cli.add_command(setpop) +cli.add_command(setenv) + +if __name__ == "__main__": + cli() diff --git a/src/cli/nomad_allocs.py b/src/cli/nomad_allocs.py new file mode 100755 index 0000000..f2369bb --- /dev/null +++ b/src/cli/nomad_allocs.py @@ -0,0 +1,51 @@ +#!/usr/bin/env python3 + +import sys + +import click +import requests + + +@click.command() +@click.argument("job") +@click.option("--dc", default="bom1", help="Name of the data center") +@click.option("--token", help="Token for Nomad") +def cli(job, dc, token): + if token is None: + print("you need to pass a valid token") + sys.exit(1) + + headers = {"Authorization": f"Bearer {token}"} + url = f"https://{dc}-nomad.simulprod.com/v1/job/{job}/allocations" + try: + resp = requests.get(url, headers=headers) + resp.raise_for_status + except Exception as e: + print("return {}".format(str(e))) + + running_tasks = [] + terminated_tasks = [] + for task in resp.json(): + task_name = list(task["TaskStates"].keys())[0] + if task["TaskStates"][task_name]["State"] == "running": + running_tasks.append( + f"https://{dc}-nomad.simulprod.com/ui/allocations/{task['ID']}/{task_name}/logs" + ) + else: + terminated_tasks.append( + f"https://{dc}-nomad.simulprod.com/ui/allocations/{task['ID']}/{task_name}/logs" + ) + + if len(running_tasks) > 0: + print("running tasks") + for t in running_tasks: + print(f"→ {t}") + + if len(terminated_tasks) > 0: + print("terminated tasks") + for t in terminated_tasks: + print(f"→ {t}") + + +if __name__ == "__main__": + cli() |