Dark-Alex-17/loki
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

feat: Support for secret injection into the global config file (API keys, for example)

Browse Source
This commit is contained in:
Alex Clarke
2025-10-16 12:30:18 -06:00
parent a10948614d
commit 950893f4a2
5 changed files with 3238 additions and 3186 deletions
Download Patch File Download Diff File Expand all files Collapse all files
config.example.yaml
+31 -31
View File
@@ -10,7 +10,7 @@ keybindings: emacs # Choose keybinding style (emacs, vi)
editor: null # Specifies the command used to edit input buffer or session. (e.g. vim, emacs, nano).
wrap: no # Controls text wrapping (no, auto, <max-width>)
wrap_code: false # Enables or disables wrapping of code blocks
vault_password_file: null # Path to a file containing the password for the Loki vault
vault_password_file: null # Path to a file containing the password for the Loki vault (cannot be a secret template)
# ---- function-calling ----
function_calling: true # Enables or disables function calling (Globally).
@@ -116,14 +116,14 @@ clients:
# See https://platform.openai.com/docs/quickstart
- type: openai
api_base: https://api.openai.com/v1 # Optional
api_key: xxx
api_key: '{{OPENAI_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
organization_id: org-xxx # Optional
# For any platform compatible with OpenAI's API
- type: openai-compatible
name: ollama
api_base: http://localhost:11434/v1
api_key: xxx # Optional
api_key: '{{OLLAMA_API_KEY}}' # Optional; You can either hard-code or inject secrets from the Loki vault
models:
- name: deepseek-r1
max_input_tokens: 131072
@@ -141,7 +141,7 @@ clients:
# See https://ai.google.dev/docs
- type: gemini
api_base: https://generativelanguage.googleapis.com/v1beta
api_key: xxx
api_key: '{{GEMINI_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
patch:
chat_completions:
'.*':
@@ -158,50 +158,50 @@ clients:
# See https://docs.anthropic.com/claude/reference/getting-started-with-the-api
- type: claude
api_base: https://api.anthropic.com/v1 # Optional
api_key: xxx
api_base: https://api.anthropic.com/v1 # Optional
api_key: '{{ANTHROPIC_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://docs.mistral.ai/
- type: openai-compatible
name: mistral
api_base: https://api.mistral.ai/v1
api_key: xxx
api_key: '{{MISTRAL_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://docs.x.ai/docs
- type: openai-compatible
name: xai
api_base: https://api.x.ai/v1
api_key: xxx
api_key: '{{XAI_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://docs.ai21.com/docs/quickstart
- type: openai-compatible
name: ai12
api_base: https://api.ai21.com/studio/v1
api_key: xxx
api_key: '{{AI21_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://docs.cohere.com/docs/the-cohere-platform
- type: cohere
api_base: https://api.cohere.ai/v2 # Optional
api_key: xxx
api_base: https://api.cohere.ai/v2 # Optional
api_key: '{{COHERE_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://docs.perplexity.ai/docs/getting-started
- type: openai-compatible
name: perplexity
api_base: https://api.perplexity.ai
api_key: xxx
api_key: '{{PERPLEXITY_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://console.groq.com/docs/quickstart
- type: openai-compatible
name: groq
api_base: https://api.groq.com/openai/v1
api_key: xxx
api_key: '{{GROQ_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://learn.microsoft.com/en-us/azure/ai-services/openai/chatgpt-quickstart
- type: azure-openai
api_base: https://{RESOURCE}.openai.azure.com
api_key: xxx
api_key: '{{AZURE_OPENAI_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
models:
- name: gpt-4o # Model deployment name
- name: gpt-4o # Model deployment name
max_input_tokens: 128000
supports_vision: true
supports_function_calling: true
@@ -230,76 +230,76 @@ clients:
# See https://docs.aws.amazon.com/bedrock/latest/userguide/
- type: bedrock
access_key_id: xxx
secret_access_key: xxx
access_key_id: '{{AWS_ACCESS_KEY_ID}}' # You can either hard-code or inject secrets from the Loki vault
secret_access_key: '{{AWS_SECRET_ACCESS_KEY}}' # You can either hard-code or inject secrets from the Loki vault
region: xxx
session_token: xxx # Optional, only needed for temporary credentials
session_token: xxx # Optional, only needed for temporary credentials
# See https://developers.cloudflare.com/workers-ai/
- type: openai-compatible
name: cloudflare
api_base: https://api.cloudflare.com/client/v4/accounts/{ACCOUNT_ID}/ai/v1
api_key: xxx
api_key: '{{CLOUDFLARE_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://cloud.baidu.com/doc/WENXINWORKSHOP/index.html
- type: openai-compatible
name: ernie
api_base: https://qianfan.baidubce.com/v2
api_key: xxx
api_key: '{{BAIDU_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://dashscope.aliyun.com/
- type: openai-compatible
name: qianwen
api_base: https://dashscope.aliyuncs.com/compatible-mode/v1
api_key: xxx
api_key: '{{ALIYUN_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://cloud.tencent.com/product/hunyuan
- type: openai-compatible
name: hunyuan
api_base: https://api.hunyuan.cloud.tencent.com/v1
api_key: xxx
api_key: '{{TENCENT_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://platform.moonshot.cn/docs/intro
- type: openai-compatible
name: moonshot
api_base: https://api.moonshot.cn/v1
api_key: xxx
api_key: '{{MOONSHOT_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://platform.deepseek.com/api-docs/
- type: openai-compatible
name: deepseek
api_base: https://api.deepseek.com
api_key: xxx
api_key: '{{DEEPSEEK_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://open.bigmodel.cn/dev/howuse/introduction
- type: openai-compatible
name: zhipuai
api_base: https://open.bigmodel.cn/api/paas/v4
api_key: xxx
api_key: '{{ZHIPUAI_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://platform.minimaxi.com/document/Fast%20access
- type: openai-compatible
name: minimax
api_base: https://api.minimax.chat/v1
api_key: xxx
api_key: '{{MINIMAX_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://openrouter.ai/docs#quick-start
- type: openai-compatible
name: openrouter
api_base: https://openrouter.ai/api/v1
api_key: xxx
api_key: '{{OPENROUTER_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://github.com/marketplace/models
- type: openai-compatible
name: github
api_base: https://models.inference.ai.azure.com
api_key: xxx
api_key: '{{GITHUB_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://deepinfra.com/docs
- type: openai-compatible
name: deepinfra
api_base: https://api.deepinfra.com/v1/openai
api_key: xxx
api_key: '{{DEEPINFRA_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# ----- RAG dedicated -----
@@ -308,10 +308,10 @@ clients:
- type: openai-compatible
name: jina
api_base: https://api.jina.ai/v1
api_key: xxx
api_key: '{{JINA_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
# See https://docs.voyageai.com/docs/introduction
- type: openai-compatible
name: voyageai
api_base: https://api.voyageai.com/v1
api_key: xxx
api_key: '{{VOYAGEAI_API_KEY}}' # You can either hard-code or inject secrets from the Loki vault
src/config/mod.rs
+2772 -2730
View File
File diff suppressed because it is too large Load Diff
src/mcp/mod.rs
+227 -237
View File
@@ -1,6 +1,6 @@
use crate::config::Config;
use crate::utils::{abortable_run_with_spinner, AbortSignal};
use crate::vault::SECRET_RE;
use crate::vault::interpolate_secrets;
use anyhow::{anyhow, Context, Result};
use futures_util::future::BoxFuture;
use futures_util::{stream, StreamExt, TryStreamExt};
@@ -26,303 +26,293 @@ type ConnectedServer = RunningService<RoleClient, ()>;
#[derive(Debug, Clone, Deserialize)]
struct McpServersConfig {
#[serde(rename = "mcpServers")]
mcp_servers: HashMap<String, McpServer>,
#[serde(rename = "mcpServers")]
mcp_servers: HashMap<String, McpServer>,
}
#[derive(Debug, Clone, Deserialize)]
struct McpServer {
command: String,
args: Option<Vec<String>>,
env: Option<HashMap<String, JsonField>>,
cwd: Option<String>,
command: String,
args: Option<Vec<String>>,
env: Option<HashMap<String, JsonField>>,
cwd: Option<String>,
}
#[derive(Debug, Clone, Deserialize)]
#[serde(untagged)]
enum JsonField {
Str(String),
Bool(bool),
Int(i64),
Str(String),
Bool(bool),
Int(i64),
}
#[derive(Debug, Clone, Default)]
pub struct McpRegistry {
log_path: Option<PathBuf>,
config: Option<McpServersConfig>,
servers: HashMap<String, Arc<RunningService<RoleClient, ()>>>,
log_path: Option<PathBuf>,
config: Option<McpServersConfig>,
servers: HashMap<String, Arc<RunningService<RoleClient, ()>>>,
}
impl McpRegistry {
pub async fn init(
log_path: Option<PathBuf>,
start_mcp_servers: bool,
use_mcp_servers: Option<String>,
abort_signal: AbortSignal,
config: &Config,
) -> Result<Self> {
let mut registry = Self {
log_path,
..Default::default()
};
if !Config::mcp_config_file().try_exists().with_context(|| {
format!(
"Failed to check MCP config file at {}",
Config::mcp_config_file().display()
)
})? {
debug!(
pub async fn init(
log_path: Option<PathBuf>,
start_mcp_servers: bool,
use_mcp_servers: Option<String>,
abort_signal: AbortSignal,
config: &Config,
) -> Result<Self> {
let mut registry = Self {
log_path,
..Default::default()
};
if !Config::mcp_config_file().try_exists().with_context(|| {
format!(
"Failed to check MCP config file at {}",
Config::mcp_config_file().display()
)
})? {
debug!(
"MCP config file does not exist at {}, skipping MCP initialization",
Config::mcp_config_file().display()
);
return Ok(registry);
}
let err = || {
format!(
"Failed to load MCP config file at {}",
Config::mcp_config_file().display()
)
};
let content = tokio::fs::read_to_string(Config::mcp_config_file())
.await
.with_context(err)?;
return Ok(registry);
}
let err = || {
format!(
"Failed to load MCP config file at {}",
Config::mcp_config_file().display()
)
};
let content = tokio::fs::read_to_string(Config::mcp_config_file())
.await
.with_context(err)?;
if content.trim().is_empty() {
debug!("MCP config file is empty, skipping MCP initialization");
return Ok(registry);
}
if content.trim().is_empty() {
debug!("MCP config file is empty, skipping MCP initialization");
return Ok(registry);
}
let mut missing_secrets = vec![];
let parsed_content = SECRET_RE.replace_all(&content, |caps: &fancy_regex::Captures<'_>| {
let secret = config.vault.get_secret(&caps[1], false);
match secret {
Ok(s) => s,
Err(_) => {
missing_secrets.push(caps[1].to_string());
"".to_string()
}
}
});
let (parsed_content, missing_secrets) = interpolate_secrets(&content, &config.vault);
if !missing_secrets.is_empty() {
return Err(anyhow!(formatdoc!(
if !missing_secrets.is_empty() {
return Err(anyhow!(formatdoc!(
"
MCP config file references secrets that are missing from the vault: {:?}
Please add these secrets to the vault and try again.",
missing_secrets
)));
}
}
let mcp_servers_config: McpServersConfig =
serde_json::from_str(&parsed_content).with_context(err)?;
registry.config = Some(mcp_servers_config);
let mcp_servers_config: McpServersConfig =
serde_json::from_str(&parsed_content).with_context(err)?;
registry.config = Some(mcp_servers_config);
if start_mcp_servers && config.mcp_servers {
abortable_run_with_spinner(
registry.start_select_mcp_servers(use_mcp_servers),
"Loading MCP servers",
abort_signal,
)
.await?;
}
if start_mcp_servers && config.mcp_servers {
abortable_run_with_spinner(
registry.start_select_mcp_servers(use_mcp_servers),
"Loading MCP servers",
abort_signal,
)
.await?;
}
Ok(registry)
}
Ok(registry)
}
pub async fn reinit(
registry: McpRegistry,
use_mcp_servers: Option<String>,
abort_signal: AbortSignal,
) -> Result<Self> {
debug!("Reinitializing MCP registry");
debug!("Stopping all MCP servers");
let mut new_registry = abortable_run_with_spinner(
registry.stop_all_servers(),
"Stopping MCP servers",
abort_signal.clone(),
)
.await?;
pub async fn reinit(
registry: McpRegistry,
use_mcp_servers: Option<String>,
abort_signal: AbortSignal,
) -> Result<Self> {
debug!("Reinitializing MCP registry");
debug!("Stopping all MCP servers");
let mut new_registry = abortable_run_with_spinner(
registry.stop_all_servers(),
"Stopping MCP servers",
abort_signal.clone(),
)
.await?;
abortable_run_with_spinner(
new_registry.start_select_mcp_servers(use_mcp_servers),
"Loading MCP servers",
abort_signal,
)
.await?;
abortable_run_with_spinner(
new_registry.start_select_mcp_servers(use_mcp_servers),
"Loading MCP servers",
abort_signal,
)
.await?;
Ok(new_registry)
}
Ok(new_registry)
}
async fn start_select_mcp_servers(&mut self, use_mcp_servers: Option<String>) -> Result<()> {
if self.config.is_none() {
debug!("MCP config is not present; assuming MCP servers are disabled globally. Skipping MCP initialization");
return Ok(());
}
async fn start_select_mcp_servers(&mut self, use_mcp_servers: Option<String>) -> Result<()> {
if self.config.is_none() {
debug!("MCP config is not present; assuming MCP servers are disabled globally. Skipping MCP initialization");
return Ok(());
}
if let Some(servers) = use_mcp_servers {
debug!("Starting selected MCP servers: {:?}", servers);
let config = self
.config
.as_ref()
.with_context(|| "MCP Config not defined. Cannot start servers")?;
let mcp_servers = config.mcp_servers.clone();
if let Some(servers) = use_mcp_servers {
debug!("Starting selected MCP servers: {:?}", servers);
let config = self
.config
.as_ref()
.with_context(|| "MCP Config not defined. Cannot start servers")?;
let mcp_servers = config.mcp_servers.clone();
let enabled_servers: HashSet<String> =
servers.split(',').map(|s| s.trim().to_string()).collect();
let server_ids: Vec<String> = if servers == "all" {
mcp_servers.into_keys().collect()
} else {
mcp_servers
.into_keys()
.filter(|id| enabled_servers.contains(id))
.collect()
};
let enabled_servers: HashSet<String> =
servers.split(',').map(|s| s.trim().to_string()).collect();
let server_ids: Vec<String> = if servers == "all" {
mcp_servers.into_keys().collect()
} else {
mcp_servers
.into_keys()
.filter(|id| enabled_servers.contains(id))
.collect()
};
let results: Vec<(String, Arc<_>)> = stream::iter(
server_ids
.into_iter()
.map(|id| async { self.start_server(id).await }),
)
.buffer_unordered(num_cpus::get())
.try_collect()
.await?;
let results: Vec<(String, Arc<_>)> = stream::iter(
server_ids
.into_iter()
.map(|id| async { self.start_server(id).await }),
)
.buffer_unordered(num_cpus::get())
.try_collect()
.await?;
self.servers = results.into_iter().collect();
}
self.servers = results.into_iter().collect();
}
Ok(())
}
Ok(())
}
async fn start_server(&self, id: String) -> Result<(String, Arc<ConnectedServer>)> {
let server = self
.config
.as_ref()
.and_then(|c| c.mcp_servers.get(&id))
.with_context(|| format!("MCP server not found in config: {id}"))?;
let mut cmd = Command::new(&server.command);
if let Some(args) = &server.args {
cmd.args(args);
}
if let Some(env) = &server.env {
let env: HashMap<String, String> = env
.iter()
.map(|(k, v)| match v {
JsonField::Str(s) => (k.clone(), s.clone()),
JsonField::Bool(b) => (k.clone(), b.to_string()),
JsonField::Int(i) => (k.clone(), i.to_string()),
})
.collect();
cmd.envs(env);
}
if let Some(cwd) = &server.cwd {
cmd.current_dir(cwd);
}
async fn start_server(&self, id: String) -> Result<(String, Arc<ConnectedServer>)> {
let server = self
.config
.as_ref()
.and_then(|c| c.mcp_servers.get(&id))
.with_context(|| format!("MCP server not found in config: {id}"))?;
let mut cmd = Command::new(&server.command);
if let Some(args) = &server.args {
cmd.args(args);
}
if let Some(env) = &server.env {
let env: HashMap<String, String> = env
.iter()
.map(|(k, v)| match v {
JsonField::Str(s) => (k.clone(), s.clone()),
JsonField::Bool(b) => (k.clone(), b.to_string()),
JsonField::Int(i) => (k.clone(), i.to_string()),
})
.collect();
cmd.envs(env);
}
if let Some(cwd) = &server.cwd {
cmd.current_dir(cwd);
}
let transport = if let Some(log_path) = self.log_path.as_ref() {
cmd.stdin(Stdio::piped()).stdout(Stdio::piped());
let transport = if let Some(log_path) = self.log_path.as_ref() {
cmd.stdin(Stdio::piped()).stdout(Stdio::piped());
let log_file = OpenOptions::new()
.create(true)
.append(true)
.open(log_path)?;
let (transport, _) = TokioChildProcess::builder(cmd).stderr(log_file).spawn()?;
transport
} else {
TokioChildProcess::new(cmd)?
};
let log_file = OpenOptions::new()
.create(true)
.append(true)
.open(log_path)?;
let (transport, _) = TokioChildProcess::builder(cmd).stderr(log_file).spawn()?;
transport
} else {
TokioChildProcess::new(cmd)?
};
let service = Arc::new(
().serve(transport)
.await
.with_context(|| format!("Failed to start MCP server: {}", &server.command))?,
);
debug!(
let service = Arc::new(
().serve(transport)
.await
.with_context(|| format!("Failed to start MCP server: {}", &server.command))?,
);
debug!(
"Available tools for MCP server {id}: {:?}",
service.list_tools(None).await?
);
info!("Started MCP server: {id}");
info!("Started MCP server: {id}");
Ok((id.to_string(), service))
}
Ok((id.to_string(), service))
}
pub async fn stop_all_servers(mut self) -> Result<Self> {
for (id, server) in self.servers {
Arc::try_unwrap(server)
.map_err(|_| anyhow!("Failed to unwrap Arc for MCP server: {id}"))?
.cancel()
.await
.with_context(|| format!("Failed to stop MCP server: {id}"))?;
info!("Stopped MCP server: {id}");
}
pub async fn stop_all_servers(mut self) -> Result<Self> {
for (id, server) in self.servers {
Arc::try_unwrap(server)
.map_err(|_| anyhow!("Failed to unwrap Arc for MCP server: {id}"))?
.cancel()
.await
.with_context(|| format!("Failed to stop MCP server: {id}"))?;
info!("Stopped MCP server: {id}");
}
self.servers = HashMap::new();
self.servers = HashMap::new();
Ok(self)
}
Ok(self)
}
pub fn list_started_servers(&self) -> Vec<String> {
self.servers.keys().cloned().collect()
}
pub fn list_started_servers(&self) -> Vec<String> {
self.servers.keys().cloned().collect()
}
pub fn list_configured_servers(&self) -> Vec<String> {
if let Some(config) = &self.config {
config.mcp_servers.keys().cloned().collect()
} else {
vec![]
}
}
pub fn list_configured_servers(&self) -> Vec<String> {
if let Some(config) = &self.config {
config.mcp_servers.keys().cloned().collect()
} else {
vec![]
}
}
pub fn catalog(&self) -> BoxFuture<'static, Result<Value>> {
let servers: Vec<(String, Arc<ConnectedServer>)> = self
.servers
.iter()
.map(|(id, s)| (id.clone(), s.clone()))
.collect();
pub fn catalog(&self) -> BoxFuture<'static, Result<Value>> {
let servers: Vec<(String, Arc<ConnectedServer>)> = self
.servers
.iter()
.map(|(id, s)| (id.clone(), s.clone()))
.collect();
Box::pin(async move {
let mut out = Vec::with_capacity(servers.len());
for (id, server) in servers {
let tools = server.list_tools(None).await?;
let resources = server.list_resources(None).await.unwrap_or_default();
// TODO implement prompt sampling for MCP servers
// let prompts = server.service.list_prompts(None).await.unwrap_or_default();
out.push(json!({
Box::pin(async move {
let mut out = Vec::with_capacity(servers.len());
for (id, server) in servers {
let tools = server.list_tools(None).await?;
let resources = server.list_resources(None).await.unwrap_or_default();
// TODO implement prompt sampling for MCP servers
// let prompts = server.service.list_prompts(None).await.unwrap_or_default();
out.push(json!({
"server": id,
"tools": tools,
"resources": resources,
}));
}
Ok(Value::Array(out))
})
}
}
Ok(Value::Array(out))
})
}
pub fn invoke(
&self,
server: &str,
tool: &str,
arguments: Value,
) -> BoxFuture<'static, Result<CallToolResult>> {
let server = self
.servers
.get(server)
.cloned()
.with_context(|| format!("Invoked MCP server does not exist: {server}"));
pub fn invoke(
&self,
server: &str,
tool: &str,
arguments: Value,
) -> BoxFuture<'static, Result<CallToolResult>> {
let server = self
.servers
.get(server)
.cloned()
.with_context(|| format!("Invoked MCP server does not exist: {server}"));
let tool = tool.to_owned();
Box::pin(async move {
let server = server?;
let call_tool_request = CallToolRequestParam {
name: Cow::Owned(tool.to_owned()),
arguments: arguments.as_object().cloned(),
};
let tool = tool.to_owned();
Box::pin(async move {
let server = server?;
let call_tool_request = CallToolRequestParam {
name: Cow::Owned(tool.to_owned()),
arguments: arguments.as_object().cloned(),
};
let result = server.call_tool(call_tool_request).await?;
Ok(result)
})
}
let result = server.call_tool(call_tool_request).await?;
Ok(result)
})
}
pub fn is_empty(&self) -> bool {
self.servers.is_empty()
}
pub fn is_empty(&self) -> bool {
self.servers.is_empty()
}
}
src/vault/mod.rs
+93 -91
View File
@@ -1,5 +1,7 @@
mod utils;
pub use utils::interpolate_secrets;
use crate::cli::Cli;
use crate::config::Config;
use crate::vault::utils::ensure_password_file_initialized;
@@ -11,122 +13,122 @@ use inquire::{required, Password, PasswordDisplayMode};
use std::sync::LazyLock;
use tokio::runtime::Handle;
pub static SECRET_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{\{(.+)}}").unwrap());
static SECRET_RE: LazyLock<Regex> = LazyLock::new(|| Regex::new(r"\{\{(.+)}}").unwrap());
#[derive(Debug, Default, Clone)]
pub struct Vault {
local_provider: LocalProvider,
local_provider: LocalProvider,
}
impl Vault {
pub fn init(config: &Config) -> Self {
let vault_password_file = config.vault_password_file();
let mut local_provider = LocalProvider {
password_file: Some(vault_password_file),
git_branch: None,
..LocalProvider::default()
};
pub fn init(config: &Config) -> Self {
let vault_password_file = config.vault_password_file();
let mut local_provider = LocalProvider {
password_file: Some(vault_password_file),
git_branch: None,
..LocalProvider::default()
};
ensure_password_file_initialized(&mut local_provider)
.expect("Failed to initialize password file");
ensure_password_file_initialized(&mut local_provider)
.expect("Failed to initialize password file");
Self { local_provider }
}
Self { local_provider }
}
pub fn add_secret(&self, secret_name: &str) -> Result<()> {
let secret_value = Password::new("Enter the secret value:")
.with_validator(required!())
.with_display_mode(PasswordDisplayMode::Masked)
.prompt()
.with_context(|| "unable to read secret from input")?;
pub fn add_secret(&self, secret_name: &str) -> Result<()> {
let secret_value = Password::new("Enter the secret value:")
.with_validator(required!())
.with_display_mode(PasswordDisplayMode::Masked)
.prompt()
.with_context(|| "unable to read secret from input")?;
let h = Handle::current();
tokio::task::block_in_place(|| {
h.block_on(self.local_provider.set_secret(secret_name, &secret_value))
})?;
println!("✓ Secret '{secret_name}' added to the vault.");
let h = Handle::current();
tokio::task::block_in_place(|| {
h.block_on(self.local_provider.set_secret(secret_name, &secret_value))
})?;
println!("✓ Secret '{secret_name}' added to the vault.");
Ok(())
}
Ok(())
}
pub fn get_secret(&self, secret_name: &str, display_output: bool) -> Result<String> {
let h = Handle::current();
let secret = tokio::task::block_in_place(|| {
h.block_on(self.local_provider.get_secret(secret_name))
})?;
pub fn get_secret(&self, secret_name: &str, display_output: bool) -> Result<String> {
let h = Handle::current();
let secret = tokio::task::block_in_place(|| {
h.block_on(self.local_provider.get_secret(secret_name))
})?;
if display_output {
println!("{}", secret);
}
if display_output {
println!("{}", secret);
}
Ok(secret)
}
Ok(secret)
}
pub fn update_secret(&self, secret_name: &str) -> Result<()> {
let secret_value = Password::new("Enter the secret value:")
.with_validator(required!())
.with_display_mode(PasswordDisplayMode::Masked)
.prompt()
.with_context(|| "unable to read secret from input")?;
let h = Handle::current();
tokio::task::block_in_place(|| {
h.block_on(
self.local_provider
.update_secret(secret_name, &secret_value),
)
})?;
println!("✓ Secret '{secret_name}' updated in the vault.");
pub fn update_secret(&self, secret_name: &str) -> Result<()> {
let secret_value = Password::new("Enter the secret value:")
.with_validator(required!())
.with_display_mode(PasswordDisplayMode::Masked)
.prompt()
.with_context(|| "unable to read secret from input")?;
let h = Handle::current();
tokio::task::block_in_place(|| {
h.block_on(
self.local_provider
.update_secret(secret_name, &secret_value),
)
})?;
println!("✓ Secret '{secret_name}' updated in the vault.");
Ok(())
}
Ok(())
}
pub fn delete_secret(&self, secret_name: &str) -> Result<()> {
let h = Handle::current();
tokio::task::block_in_place(|| h.block_on(self.local_provider.delete_secret(secret_name)))?;
println!("✓ Secret '{secret_name}' deleted from the vault.");
pub fn delete_secret(&self, secret_name: &str) -> Result<()> {
let h = Handle::current();
tokio::task::block_in_place(|| h.block_on(self.local_provider.delete_secret(secret_name)))?;
println!("✓ Secret '{secret_name}' deleted from the vault.");
Ok(())
}
Ok(())
}
pub fn list_secrets(&self, display_output: bool) -> Result<Vec<String>> {
let h = Handle::current();
let secrets =
tokio::task::block_in_place(|| h.block_on(self.local_provider.list_secrets()))?;
pub fn list_secrets(&self, display_output: bool) -> Result<Vec<String>> {
let h = Handle::current();
let secrets =
tokio::task::block_in_place(|| h.block_on(self.local_provider.list_secrets()))?;
if display_output {
if secrets.is_empty() {
println!("The vault is empty.");
} else {
for key in &secrets {
println!("{}", key);
}
}
}
if display_output {
if secrets.is_empty() {
println!("The vault is empty.");
} else {
for key in &secrets {
println!("{}", key);
}
}
}
Ok(secrets)
}
Ok(secrets)
}
pub fn handle_vault_flags(cli: Cli, config: Config) -> Result<()> {
if let Some(secret_name) = cli.add_secret {
config.vault.add_secret(&secret_name)?;
}
pub fn handle_vault_flags(cli: Cli, config: Config) -> Result<()> {
if let Some(secret_name) = cli.add_secret {
config.vault.add_secret(&secret_name)?;
}
if let Some(secret_name) = cli.get_secret {
config.vault.get_secret(&secret_name, true)?;
}
if let Some(secret_name) = cli.get_secret {
config.vault.get_secret(&secret_name, true)?;
}
if let Some(secret_name) = cli.update_secret {
config.vault.update_secret(&secret_name)?;
}
if let Some(secret_name) = cli.update_secret {
config.vault.update_secret(&secret_name)?;
}
if let Some(secret_name) = cli.delete_secret {
config.vault.delete_secret(&secret_name)?;
}
if let Some(secret_name) = cli.delete_secret {
config.vault.delete_secret(&secret_name)?;
}
if cli.list_secrets {
config.vault.list_secrets(true)?;
}
if cli.list_secrets {
config.vault.list_secrets(true)?;
}
Ok(())
}
Ok(())
}
}
src/vault/utils.rs
+115 -97
View File
@@ -1,97 +1,99 @@
use crate::config::ensure_parent_exists;
use crate::vault::{Vault, SECRET_RE};
use anyhow::anyhow;
use anyhow::Result;
use gman::providers::local::LocalProvider;
use indoc::formatdoc;
use inquire::validator::Validation;
use inquire::{min_length, required, Confirm, Password, PasswordDisplayMode, Text};
use std::borrow::Cow;
use std::path::PathBuf;
pub fn ensure_password_file_initialized(local_provider: &mut LocalProvider) -> Result<()> {
let vault_password_file = local_provider
.password_file
.clone()
.ok_or_else(|| anyhow!("Password file is not configured"))?;
let vault_password_file = local_provider
.password_file
.clone()
.ok_or_else(|| anyhow!("Password file is not configured"))?;
if vault_password_file.exists() {
{
let file_contents = std::fs::read_to_string(&vault_password_file)?;
if !file_contents.trim().is_empty() {
return Ok(());
}
}
if vault_password_file.exists() {
{
let file_contents = std::fs::read_to_string(&vault_password_file)?;
if !file_contents.trim().is_empty() {
return Ok(());
}
}
let ans = Confirm::new(
format!(
"The configured password file '{}' is empty. Create a password?",
vault_password_file.display()
)
.as_str(),
)
.with_default(true)
.prompt()?;
let ans = Confirm::new(
format!(
"The configured password file '{}' is empty. Create a password?",
vault_password_file.display()
)
.as_str(),
)
.with_default(true)
.prompt()?;
if !ans {
return Err(anyhow!("The configured password file '{}' is empty. Please populate it with a password and try again.", vault_password_file.display()));
}
if !ans {
return Err(anyhow!("The configured password file '{}' is empty. Please populate it with a password and try again.", vault_password_file.display()));
}
let password = Password::new("Enter a password to encrypt all vault secrets:")
.with_validator(required!())
.with_validator(min_length!(10))
.with_display_mode(PasswordDisplayMode::Masked)
.prompt();
let password = Password::new("Enter a password to encrypt all vault secrets:")
.with_validator(required!())
.with_validator(min_length!(10))
.with_display_mode(PasswordDisplayMode::Masked)
.prompt();
match password {
Ok(pw) => {
std::fs::write(&vault_password_file, pw.as_bytes())?;
println!(
"✓ Password file '{}' updated.",
vault_password_file.display()
);
}
Err(_) => {
return Err(anyhow!(
match password {
Ok(pw) => {
std::fs::write(&vault_password_file, pw.as_bytes())?;
println!(
"✓ Password file '{}' updated.",
vault_password_file.display()
);
}
Err(_) => {
return Err(anyhow!(
"Failed to read password from input. Password file not updated."
));
}
}
} else {
let ans = Confirm::new("No password file configured. Do you want to create one now?")
.with_default(true)
.prompt()?;
}
}
} else {
let ans = Confirm::new("No password file configured. Do you want to create one now?")
.with_default(true)
.prompt()?;
if !ans {
return Err(anyhow!("A password file is required to utilize the Loki vault. Please configure a password file in your config file and try again."));
}
if !ans {
return Err(anyhow!("A password file is required to utilize the Loki vault. Please configure a password file in your config file and try again."));
}
let password_file: PathBuf = Text::new("Enter the path to the password file to create:")
.with_default(&vault_password_file.display().to_string())
.with_validator(required!("Password file path is required"))
.with_validator(|input: &str| {
let path = PathBuf::from(input);
if path.exists() {
Ok(Validation::Invalid(
"File already exists. Please choose a different path.".into(),
))
} else if let Some(parent) = path.parent() {
if !parent.exists() {
Ok(Validation::Invalid(
"Parent directory does not exist.".into(),
))
} else {
Ok(Validation::Valid)
}
} else {
Ok(Validation::Valid)
}
})
.prompt()?
.into();
let password_file: PathBuf = Text::new("Enter the path to the password file to create:")
.with_default(&vault_password_file.display().to_string())
.with_validator(required!("Password file path is required"))
.with_validator(|input: &str| {
let path = PathBuf::from(input);
if path.exists() {
Ok(Validation::Invalid(
"File already exists. Please choose a different path.".into(),
))
} else if let Some(parent) = path.parent() {
if !parent.exists() {
Ok(Validation::Invalid(
"Parent directory does not exist.".into(),
))
} else {
Ok(Validation::Valid)
}
} else {
Ok(Validation::Valid)
}
})
.prompt()?
.into();
if password_file != vault_password_file {
println!(
"{}",
formatdoc!(
if password_file != vault_password_file {
println!(
"{}",
formatdoc!(
"
Note: The default password file path is '{}'.
You have chosen to create a different path: '{}'.
@@ -100,33 +102,49 @@ pub fn ensure_password_file_initialized(local_provider: &mut LocalProvider) -> R
vault_password_file.display(),
password_file.display()
)
);
}
);
}
ensure_parent_exists(&password_file)?;
ensure_parent_exists(&password_file)?;
let password = Password::new("Enter a password to encrypt all vault secrets:")
.with_display_mode(PasswordDisplayMode::Masked)
.with_validator(required!())
.with_validator(min_length!(10))
.prompt();
let password = Password::new("Enter a password to encrypt all vault secrets:")
.with_display_mode(PasswordDisplayMode::Masked)
.with_validator(required!())
.with_validator(min_length!(10))
.prompt();
match password {
Ok(pw) => {
std::fs::write(&password_file, pw.as_bytes())?;
local_provider.password_file = Some(password_file);
println!(
"✓ Password file '{}' created.",
vault_password_file.display()
);
}
Err(_) => {
return Err(anyhow!(
match password {
Ok(pw) => {
std::fs::write(&password_file, pw.as_bytes())?;
local_provider.password_file = Some(password_file);
println!(
"✓ Password file '{}' created.",
vault_password_file.display()
);
}
Err(_) => {
return Err(anyhow!(
"Failed to read password from input. Password file not created."
));
}
}
}
}
}
}
Ok(())
Ok(())
}
pub fn interpolate_secrets<'a>(content: &'a str, vault: &Vault) -> (Cow<'a, str>, Vec<String>) {
let mut missing_secrets = vec![];
let parsed_content = SECRET_RE.replace_all(content, |caps: &fancy_regex::Captures<'_>| {
let secret = vault.get_secret(caps[1].trim(), false);
match secret {
Ok(s) => s,
Err(_) => {
missing_secrets.push(caps[1].to_string());
"".to_string()
}
}
});
(parsed_content, missing_secrets)
}