Initial code commit

setup-scripts/CommandRunner.py

@@ -0,0 +1,93 @@
+import os, sys, subprocess, select
+from typing import List, Tuple
+
+class CommandRunner:
+    @staticmethod
+    def run_command(command: str | List[str], check=True) -> Tuple[int, str, str]:
+        """Run a command on the system and return the output
+
+        Args:
+            command (str | List[str]): The command to run
+            check (bool, optional): If the command should raise an exception if it fails (`check=` for `subprocess.run`). Defaults to True.
+        
+        Returns:
+            Tuple[int, str, str]: The return code of the command, the output of the command (on standard out), and the error output of the command (on standard error)
+        """
+        
+        # Copy of the environment variables
+        new_env = os.environ.copy()
+
+        if not check:
+            # If the check isn't set, leave it out of the `subprocess.run`` call
+            # This means no exception will be raised if the command fails
+            result = subprocess.run(command, env=new_env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True)
+            return result.returncode, result.stdout.decode('utf-8').strip(), result.stderr.decode('utf-8')
+        else:
+            # Run the command and raise an exception if it fails
+            try:
+                result = subprocess.run(command, env=new_env, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, check=True)
+                return result.returncode, result.stdout.decode('utf-8').strip(), result.stderr.decode('utf-8').strip()
+            except subprocess.CalledProcessError as e:
+                # Log the error and re-raise or handle as appropriate
+                raise RuntimeError(f"Command '{command}' failed with error: {e.stderr}") from e
+    
+    @staticmethod
+    def run_command_in_real_time(command: str | List[str]) -> Tuple[int, str, str]:
+        """Similar to `run_command` but prints the output of the command in real-time
+
+        Args:
+            command (str | List[str]): The command to run
+        Returns:
+            tuple[int, str, str]: The return code of the command, the output of the command (on standard out), and the error output of the command (on standard error)
+        """
+
+        # Variables to store the output of the command
+        stdout_lines: list[str] = []
+        stderr_lines: list[str] = []
+
+        # Copy of the environment variables
+        my_env = os.environ.copy()
+
+        # Run the command
+        # We use a `with` statement to ensure the process is closed properly
+        with subprocess.Popen(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, env=my_env, shell=True) as proc:
+            # Call `poll()` initially because we want to emulate a do-while loop
+            return_code = proc.poll()
+            
+            # Loop until the process is finished
+            # In theory this will loop about every second (timeout on the select call)
+            # Or slightly faster if the process outputs something
+            while(return_code is None):
+                # Because `readline` block until it gets a line,
+                # But the executed command ISN'T guaranteed to output a line every time
+                # We use `select` to check if there's something to read
+                # It Waits for 1 second or for the process to output something
+                rlist, _, _ = select.select([proc.stdout.fileno(), proc.stderr.fileno()], [], [], 1)
+
+                # There was something to read from the process's stdout
+                if proc.stdout.fileno() in rlist:
+                    # Read the line from the process
+                    stdout_line = proc.stdout.readline()
+
+                    # Add the line to the cumulative output
+                    stdout_lines.append(stdout_line.decode('utf-8').strip())
+                    
+                    # Print the output in real-time
+                    print(stdout_line.decode('utf-8').strip())
+                
+                # There was something to read from the process's stderr
+                if proc.stderr.fileno() in rlist:
+                    # Read the line from the process
+                    stderr_line = proc.stderr.readline()
+                    
+                    # Add the line to the cumulative output
+                    stderr_lines.append(stderr_line.decode('utf-8').strip())
+
+                    # Print the error output of the command in real-time to stderr
+                    print(stderr_line.decode('utf-8').strip(), file=sys.stderr)
+
+                # Update the return code (to see if the process is finished)
+                return_code = proc.poll()
+            
+        # Return the error code AND the full output of the command as a string
+        return return_code, '\n'.join(stdout_lines), '\n'.join(stderr_lines)

setup-scripts/app-role-access.py

@@ -0,0 +1,211 @@
+#!/usr/bin/env python3
+
+# |*******************************************************************|
+# | Setup script                                                      |
+# |                                                                   |
+# | Author: Alan Bridgeman                                            |
+# | Created: 2024-03-30                                               |
+# |                                                                   |
+# | COPYRIGHT © 2024 Bridgeman Accessible/Alan Bridgeman.             |
+# |                                                                   |
+# | This work is presented AS IS, with no warranty of any kind.       |
+# | Any modification or use of this script is at the user's own risk. |
+# |*******************************************************************|
+
+import os, sys, subprocess, logging
+
+from CommandRunner import CommandRunner
+
+#def run_command(command):
+#    try:
+#        result = subprocess.run(command, stdout=subprocess.PIPE, stderr=subprocess.PIPE, shell=True, check=True)
+#        return result.stdout.decode('utf-8').strip()
+#    except subprocess.CalledProcessError as e:
+#        # Log the error and re-raise or handle as appropriate
+#        raise RuntimeError(f"Command '{' '.join(command)}' failed with error: {e.stderr}") from e
+
+def create_policy(policy_name: str, policy_capabilities: list[str], policy_path: str = 'secret/*'):
+    """Create a policy in vault
+
+    Args:
+        policy_name (str): The name of the policy to create
+        policy_capabilities (list[str]): The capabilities of the policy (ex. read, write, etc...)
+        policy_path (str, optional): The path the policy should apply to. Defaults to all secrets (`secret/*`).
+    """
+
+    policy_caps = '["' + '","'.join(policy_capabilities) + '"]'
+    policy_content = 'path "' + policy_path + '" {\n  capabilities = ' + policy_caps + '\n}'
+    policy = '<<EOF\n' + policy_content + '\nEOF'
+    policy_return_code, policy_output, policy_err = CommandRunner.run_command('vault policy write ' + policy_name + ' - ' + policy, False)
+    if policy_return_code == 2:
+        logging.error('Failed to create the policy')
+        logging.error('Policy Output: ' + policy_output)
+        logging.error('Policy Error: ' + policy_err)
+        raise RuntimeError('Failed to create the policy')
+
+def get_role_id(role_name: str) -> str:
+    """Get the `role_id` for a given role
+
+    Args:
+        role_name (str): The name of the role to get the `role_id` for
+
+    Returns:
+        str: The `role_id` for the given role
+    """
+
+    # Get the role_id from vault
+    role_read_path = '/'.join(['auth', 'approle', 'role', role_name, 'role-id'])
+    role_return_code, role_id, role_err = CommandRunner.run_command('vault read ' + role_read_path + ' | grep role_id | awk \'{print $2}\'')
+        
+    logging.debug('Role ID: ' + role_id)
+
+    return role_id
+
+def get_secret_id(role_name: str) -> str:
+    """Get a/the `secret_id` for a given role
+
+    Args:
+        role_name (str): The name of the role to get the `secret_id` for
+
+    Returns:
+        str: The `secret_id` for the given role
+    """
+
+    # Get the secret_id from vault
+    secret_write_path = '/'.join(['auth', 'approle', 'role', role_name, 'secret-id'])
+    secret_return_code, secret_id, secret_err = CommandRunner.run_command('vault write -f ' + secret_write_path + ' | grep "secret_id " | awk \'{print $2}\'')
+    
+    logging.debug('Secret ID: ' + secret_id)
+
+    return secret_id
+
+def create_app_role(role_name: str, policy_name: str) -> tuple[str, str]:
+    """Create an approle role and return the role_id and secret_id
+
+    Args:
+        role_name (str): The name of the role to create
+        policy_name (str): The name of the policy to associate with the role
+
+    Returns:
+        tuple[str, str]: The `role_id` and `secret_id` of the newly created role
+    """
+
+    # Create a role
+    role_write_path = '/'.join(['auth', 'approle', 'role', role_name])
+    role_write_return_code, role_write_output, role_write_err = CommandRunner.run_command('vault write ' + role_write_path + ' token_policies="' + policy_name + '"')
+    
+    logging.debug(role_write_output)
+
+    # Get the role_id
+    role_id = get_role_id(role_name)
+    
+    # Get the secret_id
+    secret_id = get_secret_id(role_name)
+    
+    return role_id, secret_id
+
+def save_role_vars_to_backup_file(role_name: str, role_id: str, secret_id: str):
+    """Save the role_id and secret_id to a backup file (in the `/vault/creds` directory which is mounted to the host)
+
+    This is used as a sort of backup in rare case where the `/role_vars/.env` doesn't persist across restarts/instances.
+
+    Args:
+        role_name (str): The name of the role (used as the filename inside of the `/vault/creds` directory)
+        role_id (str): The `role_id` to save
+        secret_id (str): The `secret_id` to save
+    """
+    
+    file_name = f'/vault/creds/{role_name}'
+    
+    logging.debug('File Name: ' + file_name)
+    
+    # Create the file if it doesn't exist
+    if not os.path.isfile(file_name):
+        CommandRunner.run_command('touch "' + file_name + '"')
+    
+    # Write the role_id and secret_id to the file
+    with open(file_name, 'w+') as f:
+        f.write('='.join([os.environ['ROLE_ID_SECRET_NAME'], role_id]) + '\n')
+        f.write('='.join([os.environ['SECRET_ID_SECRET_NAME'], secret_id]))
+
+def save_role_vars_to_file(role_id: str, secret_id: str):
+    """Save the role_id and secret_id to a file (`/role_vars/.env`)
+
+    This file can then be loaded/mounted/etc... into other containers or by other scripts to load the values so that the app can load it
+
+    Args:
+        role_id (str): The `role_id` to save
+        secret_id (str): The `secret_id` to save
+    """
+
+    # Create the directory if it doesn't exist
+    if not os.path.isdir('/role_vars'):
+        os.mkdir('/role_vars')
+    
+    file_name = '/role_vars/.env'
+    
+    logging.debug('File Name: ' + file_name)
+    
+    # Create the file if it doesn't exist
+    if not os.path.isfile(file_name):
+        CommandRunner.run_command('touch "' + file_name + '"')
+    
+    # Write the role_id and secret_id to the file
+    with open(file_name, 'w+') as f:
+        f.write('='.join([os.environ['ROLE_ID_SECRET_NAME'], role_id]) + '\n')
+        f.write('='.join([os.environ['SECRET_ID_SECRET_NAME'], secret_id]))
+
+def main(token: str):
+    #logging.basicConfig(
+    #    filename='/var/log/entrypoint.log',
+    #    level=logging.DEBUG,
+    #    format='%(asctime)s - %(levelname)s - %(message)s'
+    #)
+
+    # Login to vault
+    CommandRunner.run_command(f'vault login {token}')
+
+    # Enable approle auth method
+    CommandRunner.run_command('vault auth enable approle')
+    
+    # -- Create a policy for the app --
+    
+    # The policy name can be set via the `POLICY_NAME` environment variable
+    # Or defaults to `node-app-policy` if not set
+    if 'POLICY_NAME' in os.environ:
+        policy_name = os.environ['POLICY_NAME']
+    else:
+        policy_name = 'node-app-policy'
+
+    logging.debug(f'Policy Name: {policy_name}')
+    
+    # The policy capabilities can be set via the `POLICY_CAPABILITIES` environment variable
+    # Or defaults to `['read', 'create', 'update']` if not set
+    if 'POLICY_CAPABILITIES' in os.environ:
+        policy_capabilities = [cap.strip() for cap in os.environ['POLICY_CAPABILITIES'].split(',')] 
+    else:
+        policy_capabilities = ['read', 'create', 'update']
+    
+    logging.debug(f'Policy Capabilities: {', '.join(policy_capabilities)}')
+
+    create_policy(policy_name, policy_capabilities)
+
+    # -- create an approle role --
+
+    # Create a role
+    if 'ROLE_NAME' in os.environ:
+        role_name = os.environ['ROLE_NAME']
+    else:
+        role_name = 'node-app'
+    
+    role_id, secret_id = create_app_role(role_name, policy_name)
+    
+    # Save the role_id and secret_id to a backup file
+    save_role_vars_to_backup_file(role_name, role_id, secret_id)
+
+    # Save the role_id and secret_id to a file
+    save_role_vars_to_file(role_id, secret_id)
+
+if __name__ == '__main__':
+    token = sys.argv[1]
+    main(token)

setup-scripts/prod-setup.py

@@ -0,0 +1,144 @@
+import os, json
+from CommandRunner import CommandRunner
+
+class Initializer:
+    def create_unseal_keys_file(self, file = '/vault/creds/unseal-keys'):
+        """Write the vault's unseal keys to a file
+
+        Args:
+            file (str, optional): The path of the file to output. Defaults to '/vault/creds/unseal-keys'.
+        """
+
+        with open(file, 'w+') as f:
+            f.write('\n'.join(self.unseal_keys))
+    
+    def create_root_token_file(self, file = '/vault/creds/root-token'):
+        """Write the vault's root token to a file
+
+        Args:
+            file (str, optional): The path of the file to output. Defaults to '/vault/creds/root-token'.
+        """
+
+        with open(file, 'w+') as f:
+            f.write(self.root_token)
+
+    def init_vault(self):
+        """Initialize vault
+        
+        This includes creating the root token and unseal keys.
+        Which we want to store in case we need them later
+        """
+
+        print('*----------------------*')
+        print('| Initialization Vault |')
+        print('*----------------------*')
+
+        # Initialize the vault
+        return_code, init_output, init_err = CommandRunner.run_command('vault operator init -format=json')
+
+        # Parse the unseal keys and root token from the initialization response
+        self.unseal_keys = json.loads(init_output)['unseal_keys_b64']
+        self.root_token = json.loads(init_output)['root_token']
+
+        # UPDATE: Is mounted as a volume instead
+        #CommandRunner.run_command('mkdir /vault/creds')
+
+        self.create_unseal_keys_file()
+        self.create_root_token_file()
+
+    def is_vault_unsealed(self) -> bool:
+        """Check if the vault is sealed or not
+
+        Returns:
+            bool: If the vault is unsealed or not
+        """
+
+        # Get the status of the vault
+        # Note, because it returns a non-zero exit code when the vault is sealed, we set check to False
+        # Which is also why we need to check the return code manually
+        seal_status_returncode, seal_status_raw, seal_status_err = CommandRunner.run_command('vault status -format=json', False)
+
+        # Verify the return code is either 0 (unsealed) or 2 (sealed)
+        if seal_status_returncode != 0 and seal_status_returncode != 2:
+            raise RuntimeError('Failed to get the status of the vault')
+        
+        # Print the raw status
+        print(seal_status_raw)
+        
+        # Parse the seal stat from the status
+        seal_status = json.loads(seal_status_raw)['sealed']
+
+        print(f'Is Sealed: {seal_status}')
+
+        return seal_status
+
+    def unseal_vault(self):
+        """Unseal the vault"""
+
+        print('*-----------------*')
+        print('| Unsealing Vault |')
+        print('*-----------------*')
+
+        # Use each key to unseal the vault
+        for key in self.unseal_keys:
+            return_code, unseal_output, unseal_err = CommandRunner.run_command(f'vault operator unseal {key}')
+
+            print(unseal_output)
+
+            # If the vault is now unsealed break/escape from the loop
+            if not self.is_vault_unsealed():
+                print('Vault is unsealed')
+                break
+
+    def setup_secrets_engine(self):
+        """Setup the secrets engine"""
+
+        print('*---------------------------*')
+        print('| Setting up secrets engine |')
+        print('*---------------------------*')
+        
+        login_return_code, login_output, login_err = CommandRunner.run_command(f'vault login {self.root_token}')
+        
+        print(login_output)
+
+        engin_enable_return_code, engine_enable_output, engine_enable_err = CommandRunner.run_command('vault secrets enable -path secret kv')
+
+        print(engine_enable_output)
+    
+    def setup_audit_device(self):
+        print('*---------------------------*')
+        print('| Setting up Audit Device   |')
+        print('*---------------------------*')
+        
+        audit_return_code, audit_output, audit_err = CommandRunner.run_command('vault audit enable file file_path=/vault/logs/vault-audit.log')
+
+        print(audit_output)
+
+    def setup_app_role_access(self):
+        """Run the app role creation script"""
+        
+        print('*----------------------------*')
+        print('| Setting up App Role access |')
+        print('*----------------------------*')
+
+        print(f'Policy Capabilities: {os.getenv("POLICY_CAPABILITIES")}')
+
+        # Run the custom entrypoint Python script
+        CommandRunner.run_command_in_real_time(f'python3 /setup-scripts/app-role-access.py {self.root_token}')
+
+def main():
+    initializer = Initializer()
+    # Check if the root-token file and unseal keys files exist
+    #if os.path.exists('/vault/creds/root-token') and os.path.exists('/vault/creds/unseal-keys'):
+    if not initializer.is_vault_unsealed():
+        print('Vault already setup. Skipping...')
+        # QUESTION: Should there be code here to get the Role ID and Secret ID in case the originally created .env file doesn't exist for some reason 
+    else:
+        initializer.init_vault()
+        initializer.unseal_vault()
+        initializer.setup_secrets_engine()
+        initializer.setup_audit_device()
+        initializer.setup_app_role_access()
+
+if __name__ == '__main__':
+    main()