From 1d65d7696c17eb0d1286968839331f7aeb12f3c3 Mon Sep 17 00:00:00 2001 From: Franck Cuny Date: Mon, 6 May 2024 14:07:34 -0700 Subject: just moving stuff around --- packages/robloxenv/robloxenv.py | 207 --------------------------------------- pyproject.toml | 8 +- src/cli/hashi_env.py | 208 ++++++++++++++++++++++++++++++++++++++++ src/cli/nomad_allocs.py | 51 ++++++++++ 4 files changed, 264 insertions(+), 210 deletions(-) delete mode 100755 packages/robloxenv/robloxenv.py create mode 100755 src/cli/hashi_env.py create mode 100755 src/cli/nomad_allocs.py diff --git a/packages/robloxenv/robloxenv.py b/packages/robloxenv/robloxenv.py deleted file mode 100755 index 79f9428..0000000 --- a/packages/robloxenv/robloxenv.py +++ /dev/null @@ -1,207 +0,0 @@ -#!/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 = [ - "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/pyproject.toml b/pyproject.toml index 0c944c3..4400828 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -65,8 +65,10 @@ combine-as-imports = true known-first-party = ["rbx_nomad"] [project.scripts] +hashi_env = "cli.hash_env:cli" +int2ip = "cli.ipconverter:int2ip" +ip2int = "cli.ipconverter:ip2int" +nomad_allocs = "cli.nomad_allocs:cli" pizza = "cli.pizza:main" -slocalc = "cli.slocalc:cli" seqstat = "cli.seqstat:cli" -ip2int = "cli.ipconverter:ip2int" -int2ip = "cli.ipconverter:int2ip" +slocalc = "cli.slocalc:cli" 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() -- cgit 1.4.1