From 761557707dbe3691f8d8f14b705ac7c642c49f83 Mon Sep 17 00:00:00 2001
From: Franck Cuny <fcuny@roblox.com>
Date: Fri, 3 May 2024 13:10:12 -0700
Subject: add a comment to set environment variables

---
 packages/robloxenv/default.nix  |  1 +
 packages/robloxenv/robloxenv.py | 51 +++++++++++++++++++++++++++++++++++------
 2 files changed, 45 insertions(+), 7 deletions(-)

(limited to 'packages')

diff --git a/packages/robloxenv/default.nix b/packages/robloxenv/default.nix
index ca675e8..2b691e4 100644
--- a/packages/robloxenv/default.nix
+++ b/packages/robloxenv/default.nix
@@ -18,6 +18,7 @@ python3.pkgs.buildPythonApplication rec {
   installPhase = ''
     mkdir -p $out/bin
     cp $src $out/bin/${pname}
+    chmod +x $out/bin/${pname}
   '';
 
   meta = with lib; {
diff --git a/packages/robloxenv/robloxenv.py b/packages/robloxenv/robloxenv.py
index 222d946..79f9428 100755
--- a/packages/robloxenv/robloxenv.py
+++ b/packages/robloxenv/robloxenv.py
@@ -11,11 +11,15 @@ import requests
 # the ID of the vault in 1password
 op_vault = "v4mof5qwozyvob2utdk3crwxnu"
 
-vault_addr_chi1 = "https://chi1-vault.simulprod.com:8200"
+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", "chi11"]
+valid_dcs = ["ash1", "chi1"]
 valid_edges = [
     "ams1",
     "atl1",
@@ -140,17 +144,27 @@ def vault_read(path: str, addr: str, token: str, dc: str) -> str:
 def setpop(dc: str):
     """Print some information regarding hashi stack in a POP."""
     if dc not in valid_edges:
-        print("nop")
+        print(f"{dc} is not a valid edge location")
+        return
 
-    token = get_vault_token(vault_addr_chi1)
+    token = get_vault_token(vault_addrs["chi1"])
     consul_token = vault_read(
-        "secret/teams/neteng/traffic/consul", addr=vault_addr_chi1, token=token, dc=dc
+        "secret/teams/neteng/traffic/consul",
+        addr=vault_addrs["chi1"],
+        token=token,
+        dc=dc,
     )
     nomad_token = vault_read(
-        "secret/teams/neteng/traffic/nomad", addr=vault_addr_chi1, token=token, dc=dc
+        "secret/teams/neteng/traffic/nomad",
+        addr=vault_addrs["chi1"],
+        token=token,
+        dc=dc,
     )
     vault_token = vault_read(
-        "secret/teams/neteng/traffic/vault", addr=vault_addr_chi1, token=token, dc=dc
+        "secret/teams/neteng/traffic/vault",
+        addr=vault_addrs["chi1"],
+        token=token,
+        dc=dc,
     )
 
     print(f"consul token: {consul_token}")
@@ -159,12 +173,35 @@ def setpop(dc: str):
     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()
-- 
cgit 1.4.1