Last active
May 2, 2024 03:22
-
-
Save vucong2409/d6da07f1256c559ba7692e501bee30f5 to your computer and use it in GitHub Desktop.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| #!/bin/python3 | |
| import click | |
| import boto3 | |
| from click import ClickException | |
| KEY_VALUE_LINE_PATTERN = "{}={}\n" | |
| @click.command() | |
| @click.option("--path", help="Path of SSM variable groups.") | |
| @click.option( | |
| "--output-type", | |
| help=""" | |
| Kind of output to inject environment from SSM into. | |
| `env` for inject into environment variable and `file` to output to file. | |
| """, | |
| ) | |
| @click.option( | |
| "--output-file", | |
| help="Name of the output file. Must set `--output-type` flag to `file`.", | |
| ) | |
| def inject_ssm_secrets(path, output_type, output_file): | |
| """ | |
| Get all secrets of PlainText type within path from SSM and inject it into environment variables/save it to a file. | |
| Please notice that this script will not loop recursively, only get top level parameter of the path. | |
| """ | |
| # Precheck all the parameters. | |
| _params_precheck(path, output_type, output_file) | |
| # Init SSM Boto3 Client. | |
| click.echo("Init boto3 client...") | |
| boto3_ssm_client = boto3.client("ssm") | |
| # Get secret from SSM | |
| click.echo("Getting secrets from SSM...") | |
| raw_secret_list = _get_secrets_by_path(boto3_ssm_client, path) | |
| # Remove prefix from secret key | |
| click.echo("Parsing response from SSM...") | |
| secret_list = _remove_prefix_from_secrets_key(raw_secret_list) | |
| # Save secret into file/global environment variables. | |
| click.echo("Saving...") | |
| _save_secret_list(secret_list, output_type, output_file) | |
| click.echo("Done!") | |
| def _params_precheck(path, output_type, output_file): | |
| """ | |
| Check if `output-type` flag is valid. | |
| Check if `output-type` is set to `file` then `output-file` must exist. | |
| """ | |
| valid_output_type = ["env", "file"] | |
| if output_type not in valid_output_type: | |
| raise ClickException( | |
| "Invalid output type! Valid output type are `file` and `env`." | |
| ) | |
| if output_type == "file" and output_file is None: | |
| raise ClickException("No output file!") | |
| def _get_secrets_by_path(boto3_ssm_client: any, path: str) -> dict: | |
| """ | |
| Get all SSM secrets within a path and put them in a dict: | |
| { | |
| "/path/key-1": "value-1", | |
| "/path/key-2": "value-2" | |
| } | |
| """ | |
| secret_list = {} | |
| raw_boto3_resp = boto3_ssm_client.get_parameters_by_path(Path=path) | |
| for param_information in raw_boto3_resp["Parameters"]: | |
| secret_list[param_information["Name"]] = param_information["Value"] | |
| # If there's NextToken in response, continue calling API to get all params. | |
| # We still need the part before the loop since Python does not support do-while loop. | |
| while "NextToken" in raw_boto3_resp: | |
| next_token = raw_boto3_resp["NextToken"] | |
| raw_boto3_resp = boto3_ssm_client.get_parameters_by_path( | |
| Path=path, NextToken=next_token | |
| ) | |
| for param_information in raw_boto3_resp["Parameters"]: | |
| secret_list[param_information["Name"]] = param_information["Value"] | |
| return secret_list | |
| def _remove_prefix_from_secrets_key(raw_secret_list: dict) -> dict: | |
| """ | |
| Remove all prefix path from key of secret list. | |
| For example: | |
| `/prefix/path/secretA` -> `secretA` | |
| """ | |
| secret_list = {} | |
| for raw_secret_key, secret_value in raw_secret_list.items(): | |
| secret_key = raw_secret_key.split("/")[-1] | |
| secret_list[secret_key] = secret_value | |
| return secret_list | |
| def _save_secret_list(secret_list: dict, output_type: str, output_file: str) -> None: | |
| """ | |
| Save secret list. | |
| """ | |
| if output_type == "env": | |
| _save_secret_list_into_global_env(secret_list) | |
| elif output_type == "file": | |
| _save_secret_list_into_file(secret_list, output_file) | |
| def _save_secret_list_into_global_env(secret_list: dict) -> None: | |
| """ | |
| Save secret list into environment variables. | |
| """ | |
| TOP_LINE_COMMENT = "\n# Added by SSM variables injection script.\n" | |
| with open("/etc/profile", "a") as target_file: | |
| target_file.writelines(TOP_LINE_COMMENT) | |
| for secret_key, secret_value in secret_list.items(): | |
| target_file.writelines( | |
| KEY_VALUE_LINE_PATTERN.format(secret_key, secret_value) | |
| ) | |
| click.echo( | |
| """Variables written into /etc/profile. Please run | |
| source /etc/profile | |
| or open a new shell to apply those changes.""" | |
| ) | |
| def _save_secret_list_into_file(secret_list: dict, output_file: str) -> None: | |
| """ | |
| Save secret list into file. | |
| File will contain a list of key=value pair. | |
| Example: | |
| ``` | |
| KEY1=VALUE1 | |
| KEY2=VALUE2 | |
| KEY3=VALUE3 | |
| ``` | |
| """ | |
| with open(output_file, "w") as target_file: | |
| for secret_key, secret_value in secret_list.items(): | |
| target_file.writelines( | |
| KEY_VALUE_LINE_PATTERN.format(secret_key, secret_value) | |
| ) | |
| if __name__ == "__main__": | |
| inject_ssm_secrets() |
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment