update the version for PoC
This commit is contained in:
parent
a8068f6a36
commit
21e481e85f
@ -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")
|
||||
"""
|
||||
120
vault_reader.py
120
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(
|
||||
url=vault_addr,
|
||||
token=token,
|
||||
verify=verify
|
||||
)
|
||||
client = hvac.Client(url=vault_addr, token=token, verify=verify)
|
||||
with contextlib.closing(client):
|
||||
|
||||
# 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:
|
||||
v2_path = secret_path.replace('//', '/').strip('/')
|
||||
mount_point = v2_path.split('/')[0]
|
||||
v2_path = '/'.join(v2_path.split('/')[1:])
|
||||
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:])
|
||||
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 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")
|
||||
@ -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:
|
||||
secret_data = secret['data']
|
||||
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)}")
|
||||
|
||||
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():
|
||||
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='http://127.0.0.1:8200', help='Vault server address')
|
||||
parser.add_argument('--no-verify', action='store_true', help='Disable TLS verification')
|
||||
@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 = 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")
|
||||
sys.exit(1)
|
||||
|
||||
secret = get_vault_secret(
|
||||
print(
|
||||
"No token provided. Please set VAULT_TOKEN environment variable or use --token"
|
||||
)
|
||||
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()
|
||||
|
||||
183
vault_tool.py
183
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)}")
|
||||
sys.exit(1)
|
||||
print(f"Failed to install hvac package: {str(e)}")
|
||||
|
||||
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 get_vault_secret(
|
||||
token: str,
|
||||
path: str,
|
||||
vault_addr: str = "http://127.0.0.1:8200",
|
||||
verify: bool = True,
|
||||
) -> Optional[Any]:
|
||||
def __init__(self):
|
||||
self.valves = self.Valves()
|
||||
pass
|
||||
|
||||
def get_vault_secret(self, path: str, token: str) -> str:
|
||||
"""
|
||||
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
|
||||
client = hvac.Client(url=vault_addr, token=token, verify=verify)
|
||||
|
||||
# Check if client is authenticated
|
||||
if self.valves.TLS_VERIFY == True:
|
||||
client = hvac.Client(
|
||||
url=vault_addr, token=token, verify=self.valves.TLS_VERIFY
|
||||
)
|
||||
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 None
|
||||
return "Failed to authenticate with Vault"
|
||||
|
||||
# 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)}"
|
||||
|
||||
Loading…
Reference in New Issue
Block a user