Initial code commit

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()