update the version for PoC

README.md

@@ -1,3 +1,11 @@
 # vault-access-tool
 
 Allow Open Web-UI to access Hashicorp Vault POC
+
+This is a proof of concept for allowing the Open Web-UI to access Hashicorp Vault.
+
+You should instruct your Model on how to use the tool like this:
+
+"""
+You have access to a tool called "HashiCorp Vault Tool" that can read secrets from our hosted Hashicorp Vault.  If a user asks for a secret, assume it is stored within Vault.  You will need the path of where the secret is stored and an access token to access secrets.  If the user makes a request for a secret attempt to retrieve it via the tool's built-in token, but if it fails, you can ask the user to provide you a Token to use to access the secret. Always attempt to read the secret when asked, and show the current secret unless otherwise instructed. If you do not know the secret from running the tool, let the user know you are unable to get the secret, and show them the error as to why. An example of how to use the tool: get_vault_secret( "kv1/path/example", "hvs.sk38fm29rnhmd9833k3jke")
+"""

vault_reader.py

@@ -3,9 +3,19 @@ import subprocess
 import logging
 import urllib3
 import os
+from typing import Optional, Any
+import asyncio
+import contextlib
+
+
+class Tools:
+    def __init__(self) -> None:
+        pass
 
     # Configure logging
-logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+    logging.basicConfig(
+        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
+    )
     logger = logging.getLogger(__name__)
 
     # Try to import hvac, install if not present
@@ -16,6 +26,7 @@ except ImportError:
         try:
             subprocess.check_call([sys.executable, "-m", "pip", "install", "hvac"])
             import hvac
+
             logger.info("hvac package installed successfully")
         except subprocess.CalledProcessError as e:
             logger.error(f"Failed to install hvac package: {str(e)}")
@@ -25,46 +36,49 @@ import argparse
 
     urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
 
-def get_vault_secret(token, path, vault_addr='http://127.0.0.1:8200', verify=True):
+    @staticmethod
+    def get_vault_secret(
+        token: str,
+        path: str,
+        vault_addr: str = "http://127.0.0.1:8200",
+        verify: bool = True,
+    ) -> Optional[Any]:
         try:
             # Initialize the Vault client
-        client = hvac.Client(
+            client = hvac.Client(url=vault_addr, token=token, verify=verify)
-            url=vault_addr,
+            with contextlib.closing(client):
-            token=token,
-            verify=verify
-        )
 
                 # Check if client is authenticated
                 if not client.is_authenticated():
-            logger.error("Failed to authenticate with Vault")
+                    print("Failed to authenticate with Vault")
                     return None
 
             # Split path to separate the key if it exists
-        path_parts = path.rsplit('.', 1)
+            path_parts = path.rsplit(".", 1)
             secret_path = path_parts[0]
-        key = path_parts[1] if len(path_parts) > 1 else 'value'
+            key = path_parts[1] if len(path_parts) > 1 else "value"
 
             # Try KV v2 first
             try:
                 # For KV v2, try with and without /data/ in the path
                 try:
-                if '/data/' not in secret_path:
+                    if "/data/" not in secret_path:
-                    v2_path = secret_path.replace('//', '/').strip('/')
+                        v2_path = secret_path.replace("//", "/").strip("/")
-                    mount_point = v2_path.split('/')[0]
+                        mount_point = v2_path.split("/")[0]
-                    v2_path = '/'.join(v2_path.split('/')[1:])
+                        v2_path = "/".join(v2_path.split("/")[1:])
                     else:
                         # Remove /data/ for the API call
-                    v2_path = secret_path.replace('/data/', '/')
+                        v2_path = secret_path.replace("/data/", "/")
-                    mount_point = v2_path.split('/')[0]
+                        mount_point = v2_path.split("/")[0]
-                    v2_path = '/'.join(v2_path.split('/')[1:])
+                        v2_path = "/".join(v2_path.split("/")[1:])
 
                     secret = client.secrets.kv.v2.read_secret_version(
                         path=v2_path,
                         mount_point=mount_point,
                         raise_on_deleted_version=False,
                     )
-                if secret and 'data' in secret and 'data' in secret['data']:
+                    if secret and "data" in secret and "data" in secret["data"]:
-                    secret_data = secret['data']['data']
+                        secret_data = secret["data"]["data"]
                         if key in secret_data:
                             return secret_data[key]
                         logger.warning(f"Key '{key}' not found in KV v2 secret")
@@ -74,50 +88,76 @@ def get_vault_secret(token, path, vault_addr='http://127.0.0.1:8200', verify=Tru
                 # Try KV v1
                 try:
                     secret = client.read(secret_path)
-                if secret and 'data' in secret:
+                    if secret and "data" in secret:
-                    secret_data = secret['data']
+                        secret_data = secret["data"]
                         if key in secret_data:
                             return secret_data[key]
                         logger.warning(f"Key '{key}' not found in KV v1 secret")
                 except Exception as e:
                     logger.debug(f"KV v1 attempt failed: {str(e)}")
 
-            logger.error(f"Secret not found at path: {path}")
+                print(f"Secret not found at path: {path}")
                 return None
 
             except Exception as e:
-            logger.error(f"Error reading secret: {str(e)}")
+                print(f"Error reading secret: {str(e)}")
                 return None
 
         except Exception as e:
-        logger.error(f"Error connecting to Vault: {str(e)}")
+            print(f"Error connecting to Vault: {str(e)}")
             return None
 
-def main():
+    @staticmethod
-    parser = argparse.ArgumentParser(description='Retrieve secrets from HashiCorp Vault')
+    def vault_access() -> None:
-    parser.add_argument('--token', default=os.environ.get('VAULT_TOKEN'), 
+        """
-                       help='Vault authentication token (defaults to VAULT_TOKEN environment variable)')
+        Query Vault for secrets based on path and optionally a key
-    parser.add_argument('--path', required=True, help='Path to the secret in Vault (with optional .key)')
+        User needs to provide a path, optionally a key, and the Vault token if not set in the environment.
-    parser.add_argument('--vault-addr', default='http://127.0.0.1:8200', help='Vault server address')
+        """
-    parser.add_argument('--no-verify', action='store_true', help='Disable TLS verification')
+        parser = argparse.ArgumentParser(
+            description="Retrieve secrets from HashiCorp Vault"
+        )
+        parser.add_argument(
+            "--token",
+            default=os.environ.get("VAULT_TOKEN"),
+            help="Vault authentication token (defaults to VAULT_TOKEN environment variable)",
+        )
+        parser.add_argument(
+            "--path",
+            required=True,
+            help="Path to the secret in Vault (with optional .key)",
+        )
+        parser.add_argument(
+            "--vault-addr",
+            default="https://192.168.1.8:8200",
+            help="Vault server address",
+        )
+        parser.add_argument(
+            "--no-verify", action="store_true", help="Disable TLS verification"
+        )
 
-    args = parser.parse_args()
+        args: argparse.Namespace = parser.parse_args()
 
         if not args.token:
-        logger.error("No token provided. Please set VAULT_TOKEN environment variable or use --token")
+            print(
-        sys.exit(1)
+                "No token provided. Please set VAULT_TOKEN environment variable or use --token"
-
+            )
-    secret = get_vault_secret(
+        try: 
+            secret: Optional[Any] = Tools.get_vault_secret(
                 token=args.token,
                 path=args.path,
                 vault_addr=args.vault_addr,
-        verify=not args.no_verify
+                verify=not args.no_verify,
             )
 
             if secret is not None:
                 print(secret)
             else:
-        logger.error("Failed to retrieve secret")
+                print("Failed to retrieve secret")
+        except Exception as e:
+            print(f"Error: {e}")
+        finally:
+            if asyncio.get_event_loop().is_running():
+                pending = asyncio.all_tasks(asyncio.get_event_loop())
+                for task in pending:
+                    task.cancel()
 
-if __name__ == '__main__':
-    main()

vault_tool.py

@@ -1,153 +1,90 @@
-import sys
+"""
+title: Hashicorp Vault Tool
+author: Josh Knapp
+version: 0.4.5
+"""
+
 import subprocess
-import logging
-import urllib3
 import os
-from typing import Optional, Any
+import hvac
+import urllib3
+from pydantic import BaseModel, Field
 
 
-class Tools:
-    def __init__(self) -> None:
-        pass
-
-    # Configure logging
-    logging.basicConfig(
-        level=logging.INFO, format="%(asctime)s - %(levelname)s - %(message)s"
-    )
-    logger = logging.getLogger(__name__)
-
 # Try to import hvac, install if not present
 try:
     import hvac
 except ImportError:
-        logger.info("hvac package not found. Attempting to install...")
+    print("hvac package not found. Attempting to install...")
     try:
         subprocess.check_call([sys.executable, "-m", "pip", "install", "hvac"])
         import hvac
 
-            logger.info("hvac package installed successfully")
+        print("hvac package installed successfully")
     except subprocess.CalledProcessError as e:
-            logger.error(f"Failed to install hvac package: {str(e)}")
+        print(f"Failed to install hvac package: {str(e)}")
-            sys.exit(1)
 
-    import argparse
 
-    urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+class Tools:
+    class Valves(BaseModel):
+        VAULT_ADDR: str = Field(
+            default="",
+            description="The Web Address for the Vault Server",
+        )
+        TLS_VERIFY: bool = Field(
+            default="False",
+            description="Check the TLS Certificate for the Vault Server",
+        )
 
-    @staticmethod
+    def __init__(self):
-    def get_vault_secret(
+        self.valves = self.Valves()
-        token: str,
+        pass
-        path: str,
+
-        vault_addr: str = "http://127.0.0.1:8200",
+    def get_vault_secret(self, path: str, token: str) -> str:
-        verify: bool = True,
+        """
-    ) -> Optional[Any]:
+        Read secrets stored within HashiCorp Vault based on the path provided by the user.
+        :param token: Vault authentication token
+        :param path: Path to the secret in Vault
+        :return: The value of the secret or the error message if one is produced
+
+        """
+        # Check if Vault Address has been set
+        if not self.valves.VAULT_ADDR:
+            return "The Vault Address has not been set. Please define it in the Tool's Valves"
+        # Check if a Token is set
+        if not token:
+            return f"No token defined, please either provide one in the prompt or via an Environment Variable {token}"
+        vault_addr = self.valves.VAULT_ADDR
+        # Attempt to connect to HashiCorp Vault
         try:
-            # Initialize the Vault client
+            if self.valves.TLS_VERIFY == True:
-            client = hvac.Client(url=vault_addr, token=token, verify=verify)
+                client = hvac.Client(
-
+                    url=vault_addr, token=token, verify=self.valves.TLS_VERIFY
-            # Check if client is authenticated
+                )
+            else:
+                urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)
+                client = hvac.Client(
+                    url=vault_addr, token=token, verify=self.valves.TLS_VERIFY
+                )
+            # Check Authentication
             if not client.is_authenticated():
-                logger.error("Failed to authenticate with Vault")
+                return "Failed to authenticate with Vault"
-                return None
 
             # Split path to separate the key if it exists
             path_parts = path.rsplit(".", 1)
             secret_path = path_parts[0]
             key = path_parts[1] if len(path_parts) > 1 else "value"
-
-            # Try KV v2 first
-            try:
-                # For KV v2, try with and without /data/ in the path
-                try:
-                    if "/data/" not in secret_path:
-                        v2_path = secret_path.replace("//", "/").strip("/")
-                        mount_point = v2_path.split("/")[0]
-                        v2_path = "/".join(v2_path.split("/")[1:])
-                    else:
-                        # Remove /data/ for the API call
-                        v2_path = secret_path.replace("/data/", "/")
-                        mount_point = v2_path.split("/")[0]
-                        v2_path = "/".join(v2_path.split("/")[1:])
-
-                    secret = client.secrets.kv.v2.read_secret_version(
-                        path=v2_path,
-                        mount_point=mount_point,
-                        raise_on_deleted_version=False,
-                    )
-                    if secret and "data" in secret and "data" in secret["data"]:
-                        secret_data = secret["data"]["data"]
-                        if key in secret_data:
-                            return secret_data[key]
-                        logger.warning(f"Key '{key}' not found in KV v2 secret")
-                except Exception as e:
-                    logger.debug(f"KV v2 attempt failed: {str(e)}")
-
-                # Try KV v1
             try:
                 secret = client.read(secret_path)
                 if secret and "data" in secret:
                     secret_data = secret["data"]
                     if key in secret_data:
                         return secret_data[key]
-                        logger.warning(f"Key '{key}' not found in KV v1 secret")
-                except Exception as e:
-                    logger.debug(f"KV v1 attempt failed: {str(e)}")
-
-                print(f"Secret not found at path: {path}")
-                return None
-
-            except Exception as e:
-                print(f"Error reading secret: {str(e)}")
-                return None
-
-        except Exception as e:
-            print(f"Error connecting to Vault: {str(e)}")
-            return None
-
-    @staticmethod
-    def vault_access() -> None:
-        """
-        Query Vault for secrets based on path and optionally a key
-        User needs to provide a path, optionally a key, and the Vault token if not set in the environment.
-        """
-        parser = argparse.ArgumentParser(
-            description="Retrieve secrets from HashiCorp Vault"
-        )
-        parser.add_argument(
-            "--token",
-            default=os.environ.get("VAULT_TOKEN"),
-            help="Vault authentication token (defaults to VAULT_TOKEN environment variable)",
-        )
-        parser.add_argument(
-            "--path",
-            required=True,
-            help="Path to the secret in Vault (with optional .key)",
-        )
-        parser.add_argument(
-            "--vault-addr",
-            default="https://192.168.1.8:8200",
-            help="Vault server address",
-        )
-        parser.add_argument(
-            "--no-verify", action="store_true", help="Disable TLS verification"
-        )
-
-        args: argparse.Namespace = parser.parse_args()
-
-        if not args.token:
-            print(
-                "No token provided. Please set VAULT_TOKEN environment variable or use --token"
-            )
-            sys.exit(1)
-
-        secret: Optional[Any] = Tools.get_vault_secret(
-            token=args.token,
-            path=args.path,
-            vault_addr=args.vault_addr,
-            verify=not args.no_verify,
-        )
-
-        if secret is not None:
-            print(secret)
                     else:
-            print("Failed to retrieve secret")
+                        return f"Key '{key}' not found in KV v1 secret path. Check your secret path and ensure the token used has the appropriate permissions."
+
+            except Exception as e:
+                print(f"KV v1 attempt failed: {str(e)}")
+                return f"Secret not found at path: {path}"
+        except Exception as e:
+            return f"Error connecting to Vault: {str(e)}"