|
1 | 1 | """Defines helper functions used by Manifester."""
|
2 | 2 | from collections import UserDict
|
| 3 | +import json |
| 4 | +import os |
3 | 5 | from pathlib import Path
|
4 | 6 | import random
|
| 7 | +import re |
| 8 | +import subprocess |
| 9 | +import sys |
5 | 10 | import time
|
6 | 11 |
|
7 | 12 | from logzero import logger
|
8 | 13 | from requests import HTTPError
|
9 | 14 | import yaml
|
10 | 15 |
|
| 16 | +from manifester.logger import setup_logzero |
11 | 17 | from manifester.settings import settings
|
12 | 18 |
|
| 19 | +setup_logzero(level="info") |
| 20 | + |
| 21 | + |
13 | 22 | RESULTS_LIMIT = 10000
|
14 | 23 |
|
15 | 24 |
|
@@ -226,3 +235,139 @@ def __getitem__(self, key):
|
226 | 235 | def __call__(self, *args, **kwargs):
|
227 | 236 | """Allow MockStub to be used like a function."""
|
228 | 237 | return self
|
| 238 | + |
| 239 | + |
| 240 | +class InvalidVaultURLForOIDC(Exception): |
| 241 | + """Raised if the vault doesn't allow OIDC login.""" |
| 242 | + |
| 243 | + |
| 244 | +class Vault: |
| 245 | + """Helper class for retrieving secrets from HashiCorp Vault.""" |
| 246 | + |
| 247 | + HELP_TEXT = ( |
| 248 | + "The Vault CLI in not installed on this system." |
| 249 | + "Please follow https://learn.hashicorp.com/tutorials/vault/getting-started-install to " |
| 250 | + "install the Vault CLI." |
| 251 | + ) |
| 252 | + |
| 253 | + def __init__(self, env_file=".env"): |
| 254 | + manifester_directory = Path() |
| 255 | + |
| 256 | + if "MANIFESTER_DIRECTORY" in os.environ: |
| 257 | + envar_location = Path(os.environ["MANIFESTER_DIRECTORY"]) |
| 258 | + if envar_location.is_dir(): |
| 259 | + manifester_directory = envar_location |
| 260 | + self.env_path = manifester_directory.joinpath(env_file) |
| 261 | + self.envdata = None |
| 262 | + self.vault_enabled = None |
| 263 | + |
| 264 | + def setup(self): |
| 265 | + """Read environment variables from .env.""" |
| 266 | + if self.env_path.exists(): |
| 267 | + self.envdata = self.env_path.read_text() |
| 268 | + is_enabled = re.findall("^(?:.*\n)*VAULT_ENABLED_FOR_DYNACONF=(.*)", self.envdata) |
| 269 | + if is_enabled: |
| 270 | + self.vault_enabled = is_enabled[0] |
| 271 | + self.export_vault_addr() |
| 272 | + |
| 273 | + def teardown(self): |
| 274 | + """Remove VAULT_ADDR environment variable if present.""" |
| 275 | + if os.environ.get("VAULT_ADDR") is not None: |
| 276 | + del os.environ["VAULT_ADDR"] |
| 277 | + |
| 278 | + def export_vault_addr(self): |
| 279 | + """Set the URL of the Vault server and ensure that the URL is not localhost.""" |
| 280 | + vaulturl = re.findall("VAULT_URL_FOR_DYNACONF=(.*)", self.envdata)[0] |
| 281 | + |
| 282 | + # Set Vault CLI Env Var |
| 283 | + os.environ["VAULT_ADDR"] = vaulturl |
| 284 | + |
| 285 | + # Dynaconf Vault Env Vars |
| 286 | + if ( |
| 287 | + self.vault_enabled |
| 288 | + and self.vault_enabled in ["True", "true"] |
| 289 | + and "localhost:8200" in vaulturl |
| 290 | + ): |
| 291 | + raise InvalidVaultURLForOIDC( |
| 292 | + f"{vaulturl} does not support OIDC login." |
| 293 | + "Please set the correct vault URL vault the .env file." |
| 294 | + ) |
| 295 | + |
| 296 | + def exec_vault_command(self, command: str, **kwargs): |
| 297 | + """Wrap Vault CLI commands for execution. |
| 298 | +
|
| 299 | + :param comamnd str: The vault CLI command |
| 300 | + :param kwargs dict: Arguments to the subprocess run command to customize the run behavior |
| 301 | + """ |
| 302 | + COMMAND_NOT_FOUND_EXIT_CODE = 127 |
| 303 | + vcommand = subprocess.run(command.split(), capture_output=True, **kwargs) |
| 304 | + if vcommand.returncode != 0: |
| 305 | + verror = str(vcommand.stderr) |
| 306 | + if vcommand.returncode == COMMAND_NOT_FOUND_EXIT_CODE: |
| 307 | + logger.error(f"Error! {self.HELP_TEXT}") |
| 308 | + sys.exit(1) |
| 309 | + if vcommand.stderr: |
| 310 | + if "Error revoking token" in verror: |
| 311 | + logger.info("Token is already revoked") |
| 312 | + elif "Error looking up token" in verror: |
| 313 | + logger.info("Vault is not logged in") |
| 314 | + else: |
| 315 | + logger.error(f"Error: {verror}") |
| 316 | + return vcommand |
| 317 | + |
| 318 | + def login(self, **kwargs): |
| 319 | + """Authenticate to Vault server and add auth token to .env file.""" |
| 320 | + if ( |
| 321 | + self.vault_enabled |
| 322 | + and self.vault_enabled in ["True", "true"] |
| 323 | + and "VAULT_SECRET_ID_FOR_DYNACONF" not in os.environ |
| 324 | + and self.status(**kwargs).returncode != 0 |
| 325 | + ): |
| 326 | + logger.info( |
| 327 | + "Warning: A browser tab will open for Vault OIDC login. " |
| 328 | + "Please close the tab once the sign-in is complete" |
| 329 | + ) |
| 330 | + if ( |
| 331 | + self.exec_vault_command(command="vault login -method=oidc", **kwargs).returncode |
| 332 | + == 0 |
| 333 | + ): |
| 334 | + self.exec_vault_command(command="vault token renew -i 10h", **kwargs) |
| 335 | + logger.info("Success! Vault OIDC Logged-In and extended for 10 hours!") |
| 336 | + # Fetch token |
| 337 | + token = self.exec_vault_command("vault token lookup --format json").stdout |
| 338 | + token = json.loads(str(token.decode("UTF-8")))["data"]["id"] |
| 339 | + # Set new token in .env file |
| 340 | + _envdata = re.sub( |
| 341 | + ".*VAULT_TOKEN_FOR_DYNACONF=.*", |
| 342 | + f"VAULT_TOKEN_FOR_DYNACONF={token}", |
| 343 | + self.envdata, |
| 344 | + ) |
| 345 | + self.env_path.write_text(_envdata) |
| 346 | + logger.info("New OIDC token succesfully added to .env file") |
| 347 | + |
| 348 | + def logout(self): |
| 349 | + """Revoke Vault auth token and remove it from .env file.""" |
| 350 | + # Teardown - Setting dummy token in env file |
| 351 | + _envdata = re.sub( |
| 352 | + ".*VAULT_TOKEN_FOR_DYNACONF=.*", "# VAULT_TOKEN_FOR_DYNACONF=myroot", self.envdata |
| 353 | + ) |
| 354 | + self.env_path.write_text(_envdata) |
| 355 | + vstatus = self.exec_vault_command("vault token revoke -self") |
| 356 | + if vstatus.returncode == 0: |
| 357 | + logger.info("OIDC token successfully removed from .env file") |
| 358 | + |
| 359 | + def status(self, **kwargs): |
| 360 | + """Check status of Vault auth token.""" |
| 361 | + vstatus = self.exec_vault_command("vault token lookup", **kwargs) |
| 362 | + if vstatus.returncode == 0: |
| 363 | + logger.info(str(vstatus.stdout.decode("UTF-8"))) |
| 364 | + return vstatus |
| 365 | + |
| 366 | + def __enter__(self): |
| 367 | + """Set up Vault context manager.""" |
| 368 | + self.setup() |
| 369 | + return self |
| 370 | + |
| 371 | + def __exit__(self, exc_type, exc_val, exc_tb): |
| 372 | + """Tear down Vault context manager.""" |
| 373 | + self.teardown() |
0 commit comments