add `robloxenv` to manage hashi credentials

I never know how to get the credentials for some hashi components.
This script helps with getting the credentials I need at the edge.

5 files changed, 207 insertions, 3 deletions

diff --git a/nix/flake/packages.nix b/nix/flake/packages.nix
index dddf4b5..4038107 100644
--- a/nix/flake/packages.nix
+++ b/nix/flake/packages.nix
@@ -11,11 +11,12 @@
     };
 
     packages = {
-      seqstat = pkgs.callPackage "${self}/packages/seqstat" { };
       git-blame-stats = pkgs.callPackage "${self}/packages/git-blame-stats" { };
       git-broom = pkgs.callPackage "${self}/packages/git-broom" { };
       ipconverter = pkgs.callPackage "${self}/packages/ipconverter" { };
       pizza = pkgs.callPackage "${self}/packages/pizza" { };
+      robloxenv = pkgs.callPackage "${self}/packages/robloxenv" { };
+      seqstat = pkgs.callPackage "${self}/packages/seqstat" { };
       slocalc = pkgs.callPackage "${self}/packages/slocalc" { };
     };
   };
diff --git a/nix/profiles/home-manager/shell.nix b/nix/profiles/home-manager/shell.nix
index bf63775..2c22880 100644
--- a/nix/profiles/home-manager/shell.nix
+++ b/nix/profiles/home-manager/shell.nix
@@ -30,9 +30,9 @@
     self.packages.${pkgs.system}.git-blame-stats
     self.packages.${pkgs.system}.git-broom
     self.packages.${pkgs.system}.ipconverter
+    self.packages.${pkgs.system}.pizza
     self.packages.${pkgs.system}.seqstat
     self.packages.${pkgs.system}.slocalc
-    self.packages.${pkgs.system}.pizza
   ];
 
   xdg = {
diff --git a/nix/profiles/home-manager/work.nix b/nix/profiles/home-manager/work.nix
index 0643426..56b4bb3 100644
--- a/nix/profiles/home-manager/work.nix
+++ b/nix/profiles/home-manager/work.nix
@@ -1,4 +1,4 @@
-{ ... }: {
+{ pkgs, self, ... }: {
 
   home.stateVersion = "23.05";
 
@@ -12,6 +12,10 @@
     ./zsh.nix
   ];
 
+  home.packages = with pkgs; [
+    self.packages.${pkgs.system}.robloxenv
+  ];
+
   programs.git = {
     userEmail = "fcuny@roblox.com";
   };
diff --git a/packages/robloxenv/default.nix b/packages/robloxenv/default.nix
new file mode 100644
index 0000000..ca675e8
--- /dev/null
+++ b/packages/robloxenv/default.nix
@@ -0,0 +1,29 @@
+{ lib, python3 }:
+
+python3.pkgs.buildPythonApplication rec {
+  pname = "robloxenv";
+  src = ./robloxenv.py;
+  version = "0.1.0";
+  format = "other";
+
+  buildInputs = [ python3 ];
+  propagatedBuildInputs = with python3.pkgs; [
+    requests
+    click
+  ];
+
+  dontUnpack = true;
+  dontBuild = true;
+
+  installPhase = ''
+    mkdir -p $out/bin
+    cp $src $out/bin/${pname}
+  '';
+
+  meta = with lib; {
+    description = "This is a tool.";
+    license = with licenses; [ mit ];
+    platforms = platforms.unix;
+    maintainers = with maintainers; [ fcuny ];
+  };
+}
diff --git a/packages/robloxenv/robloxenv.py b/packages/robloxenv/robloxenv.py
new file mode 100755
index 0000000..222d946
--- /dev/null
+++ b/packages/robloxenv/robloxenv.py
@@ -0,0 +1,170 @@
+#!/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_addr_chi1 = "https://chi1-vault.simulprod.com:8200"
+creds_cache = os.path.join(os.getenv("HOME") or "", ".local/state/rbxenv")
+
+
+valid_dcs = ["ash1", "chi11"]
+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("nop")
+
+    token = get_vault_token(vault_addr_chi1)
+    consul_token = vault_read(
+        "secret/teams/neteng/traffic/consul", addr=vault_addr_chi1, token=token, dc=dc
+    )
+    nomad_token = vault_read(
+        "secret/teams/neteng/traffic/nomad", addr=vault_addr_chi1, token=token, dc=dc
+    )
+    vault_token = vault_read(
+        "secret/teams/neteng/traffic/vault", addr=vault_addr_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.group(help="CLI tool to manage environment variables for hashi things at Roblox")
+def cli():
+    pass
+
+
+cli.add_command(setpop)
+
+if __name__ == "__main__":
+    cli()