use crate::error::Result;
use crate::types::{CommandId, HostId, SessionId};
use chrono::{DateTime, Utc};
use rusqlite::{Connection, OptionalExtension, params};
use std::path::Path;
use uuid::Uuid;
#[derive(Debug, Clone)]
pub struct Host {
pub id: HostId,
pub hostname: String,
pub created_at: DateTime<Utc>,
}
#[derive(Debug, Clone)]
pub struct Session {
pub id: SessionId,
pub host_id: HostId,
pub started_at: DateTime<Utc>,
pub ended_at: Option<DateTime<Utc>>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct CommandEntry {
pub id: CommandId,
pub session_id: SessionId,
pub command: String,
pub timestamp: DateTime<Utc>,
pub directory: String,
pub redacted: bool,
pub exit_code: Option<i32>,
}
#[derive(Debug, Clone)]
pub struct Token {
pub id: i64,
pub command_id: CommandId,
pub token_type: String, pub placeholder: String,
pub original_value: String,
pub created_at: DateTime<Utc>,
}
#[derive(Debug, Clone, serde::Serialize)]
pub struct Alias {
pub alias: String,
pub command: String,
pub description: String,
pub date_created: DateTime<Utc>,
pub date_updated: DateTime<Utc>,
}
#[derive(Debug, Clone, Default)]
pub struct DatabaseStats {
pub total_commands: usize,
pub total_sessions: usize,
pub total_hosts: usize,
pub redacted_commands: usize,
pub stored_tokens: usize,
pub oldest_entry: Option<DateTime<Utc>>,
pub newest_entry: Option<DateTime<Utc>>,
}
pub struct Database {
conn: Connection,
current_host_id: HostId,
current_session_id: Option<SessionId>,
}
impl Database {
#[must_use = "Database connection must be used"]
pub fn new(db_path: &Path) -> Result<Self> {
if let Some(parent) = db_path.parent() {
std::fs::create_dir_all(parent)?;
}
let conn = Connection::open(db_path)?;
conn.execute("PRAGMA foreign_keys = ON", [])?;
let mut db = Self {
conn,
current_host_id: HostId::new(0),
current_session_id: None,
};
db.initialize_schema()?;
db.ensure_current_host()?;
Ok(db)
}
fn initialize_schema(&self) -> Result<()> {
self.conn.execute(
"CREATE TABLE IF NOT EXISTS hosts (
id INTEGER PRIMARY KEY AUTOINCREMENT,
hostname TEXT NOT NULL UNIQUE,
created_at TEXT NOT NULL
)",
[],
)?;
self.conn.execute(
"CREATE TABLE IF NOT EXISTS sessions (
id TEXT PRIMARY KEY,
host_id INTEGER NOT NULL,
started_at TEXT NOT NULL,
ended_at TEXT,
FOREIGN KEY (host_id) REFERENCES hosts(id) ON DELETE CASCADE
)",
[],
)?;
self.conn.execute(
"CREATE TABLE IF NOT EXISTS commands (
id INTEGER PRIMARY KEY AUTOINCREMENT,
session_id TEXT NOT NULL,
command TEXT NOT NULL,
timestamp TEXT NOT NULL,
directory TEXT NOT NULL,
redacted INTEGER NOT NULL DEFAULT 0,
exit_code INTEGER,
FOREIGN KEY (session_id) REFERENCES sessions(id) ON DELETE CASCADE
)",
[],
)?;
self.conn.execute(
"CREATE TABLE IF NOT EXISTS tokens (
id INTEGER PRIMARY KEY AUTOINCREMENT,
command_id INTEGER NOT NULL,
token_type TEXT NOT NULL,
placeholder TEXT NOT NULL,
original_value TEXT NOT NULL,
created_at TEXT NOT NULL,
FOREIGN KEY (command_id) REFERENCES commands(id) ON DELETE CASCADE
)",
[],
)?;
self.conn.execute(
"CREATE TABLE IF NOT EXISTS aliases (
alias TEXT PRIMARY KEY,
command TEXT NOT NULL,
description TEXT NOT NULL,
date_created TEXT NOT NULL,
date_updated TEXT NOT NULL
)",
[],
)?;
self.conn.execute(
"CREATE INDEX IF NOT EXISTS idx_commands_timestamp ON commands(timestamp DESC)",
[],
)?;
self.conn.execute(
"CREATE INDEX IF NOT EXISTS idx_commands_session ON commands(session_id)",
[],
)?;
self.conn.execute(
"CREATE INDEX IF NOT EXISTS idx_commands_directory ON commands(directory)",
[],
)?;
self.conn.execute(
"CREATE INDEX IF NOT EXISTS idx_tokens_command ON tokens(command_id)",
[],
)?;
self.conn.execute(
"CREATE INDEX IF NOT EXISTS idx_sessions_host ON sessions(host_id)",
[],
)?;
Ok(())
}
fn ensure_current_host(&mut self) -> Result<()> {
let hostname = hostname::get()
.map(|h| h.to_string_lossy().to_string())
.unwrap_or_else(|_| "unknown".to_string());
let host_id: Option<i64> = self
.conn
.query_row(
"SELECT id FROM hosts WHERE hostname = ?1",
params![hostname],
|row| row.get(0),
)
.optional()?;
self.current_host_id = if let Some(id) = host_id {
HostId::new(id)
} else {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO hosts (hostname, created_at) VALUES (?1, ?2)",
params![hostname, now],
)?;
HostId::new(self.conn.last_insert_rowid())
};
Ok(())
}
pub fn start_session(&mut self) -> Result<String> {
let session_id = Uuid::new_v4().to_string();
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO sessions (id, host_id, started_at) VALUES (?1, ?2, ?3)",
params![session_id, self.current_host_id.as_i64(), now],
)?;
self.current_session_id = Some(SessionId::new(session_id.clone()));
Ok(session_id)
}
pub fn end_session(&mut self, session_id: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"UPDATE sessions SET ended_at = ?1 WHERE id = ?2",
params![now, session_id],
)?;
if self.current_session_id.as_deref() == Some(session_id) {
self.current_session_id = None;
}
Ok(())
}
pub fn ensure_session(&mut self) -> Result<String> {
if let Some(ref session_id) = self.current_session_id {
Ok(session_id.as_str().to_string())
} else {
self.start_session()
}
}
pub fn add_command(
&mut self,
command: &str,
directory: &str,
timestamp: DateTime<Utc>,
redacted: bool,
exit_code: Option<i32>,
) -> Result<i64> {
let session_id = self.ensure_session()?;
let timestamp_str = timestamp.to_rfc3339();
self.conn.execute(
"INSERT INTO commands (session_id, command, timestamp, directory, redacted, exit_code)
VALUES (?1, ?2, ?3, ?4, ?5, ?6)",
params![
session_id,
command,
timestamp_str,
directory,
redacted as i32,
exit_code
],
)?;
Ok(self.conn.last_insert_rowid())
}
pub fn store_token(
&self,
command_id: i64,
token_type: &str,
placeholder: &str,
original_value: &str,
) -> Result<i64> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO tokens (command_id, token_type, placeholder, original_value, created_at)
VALUES (?1, ?2, ?3, ?4, ?5)",
params![command_id, token_type, placeholder, original_value, now],
)?;
Ok(self.conn.last_insert_rowid())
}
#[must_use = "Token query results should be used"]
pub fn get_tokens_for_command(&self, command_id: CommandId) -> Result<Vec<Token>> {
let mut stmt = self.conn.prepare(
"SELECT id, command_id, token_type, placeholder, original_value, created_at
FROM tokens WHERE command_id = ?1",
)?;
let tokens = stmt
.query_map(params![command_id.as_i64()], |row| {
Ok(Token {
id: row.get(0)?,
command_id: CommandId::new(row.get(1)?),
token_type: row.get(2)?,
placeholder: row.get(3)?,
original_value: row.get(4)?,
created_at: row
.get::<_, String>(5)?
.parse()
.unwrap_or_else(|_| Utc::now()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(tokens)
}
pub fn get_tokens_by_session(&self, session_id: &str) -> Result<Vec<Token>> {
let mut stmt = self.conn.prepare(
"SELECT t.id, t.command_id, t.token_type, t.placeholder, t.original_value, t.created_at
FROM tokens t
JOIN commands c ON t.command_id = c.id
WHERE c.session_id = ?1
ORDER BY t.created_at DESC",
)?;
let tokens = stmt
.query_map(params![session_id], |row| {
Ok(Token {
id: row.get(0)?,
command_id: row.get(1)?,
token_type: row.get(2)?,
placeholder: row.get(3)?,
original_value: row.get(4)?,
created_at: row
.get::<_, String>(5)?
.parse()
.unwrap_or_else(|_| Utc::now()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(tokens)
}
pub fn get_tokens_by_directory(&self, directory: &str) -> Result<Vec<Token>> {
let mut stmt = self.conn.prepare(
"SELECT t.id, t.command_id, t.token_type, t.placeholder, t.original_value, t.created_at
FROM tokens t
JOIN commands c ON t.command_id = c.id
WHERE c.directory = ?1
ORDER BY t.created_at DESC",
)?;
let tokens = stmt
.query_map(params![directory], |row| {
Ok(Token {
id: row.get(0)?,
command_id: row.get(1)?,
token_type: row.get(2)?,
placeholder: row.get(3)?,
original_value: row.get(4)?,
created_at: row
.get::<_, String>(5)?
.parse()
.unwrap_or_else(|_| Utc::now()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(tokens)
}
#[must_use = "Search results should be used"]
pub fn search_commands(
&self,
query: &str,
directory_filter: Option<&str>,
host_filter: Option<&str>,
limit: Option<usize>,
) -> Result<Vec<CommandEntry>> {
let mut sql = String::from(
"SELECT c.id, c.session_id, c.command, c.timestamp, c.directory, c.redacted, c.exit_code
FROM commands c
JOIN sessions s ON c.session_id = s.id
JOIN hosts h ON s.host_id = h.id
WHERE c.command LIKE ?1",
);
let mut params: Vec<Box<dyn rusqlite::ToSql>> = vec![Box::new(format!("%{}%", query))];
if let Some(dir) = directory_filter {
sql.push_str(" AND c.directory LIKE ?");
params.push(Box::new(format!("%{}%", dir)));
}
if let Some(host) = host_filter {
sql.push_str(" AND h.hostname = ?");
params.push(Box::new(host.to_string()));
}
sql.push_str(" ORDER BY c.timestamp DESC");
if let Some(lim) = limit {
sql.push_str(" LIMIT ?");
params.push(Box::new(lim as i64));
}
let mut stmt = self.conn.prepare(&sql)?;
let param_refs: Vec<&dyn rusqlite::ToSql> = params.iter().map(|b| b.as_ref()).collect();
let commands = stmt
.query_map(param_refs.as_slice(), |row| {
Ok(CommandEntry {
id: row.get(0)?,
session_id: row.get(1)?,
command: row.get(2)?,
timestamp: row
.get::<_, String>(3)?
.parse()
.unwrap_or_else(|_| Utc::now()),
directory: row.get(4)?,
redacted: row.get::<_, i32>(5)? != 0,
exit_code: row.get(6)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(commands)
}
#[must_use = "Query results should be used"]
pub fn get_recent_commands(&self, limit: usize) -> Result<Vec<CommandEntry>> {
let mut stmt = self.conn.prepare(
"SELECT id, session_id, command, timestamp, directory, redacted, exit_code
FROM commands
ORDER BY timestamp DESC
LIMIT ?1",
)?;
let commands = stmt
.query_map(params![limit as i64], |row| {
Ok(CommandEntry {
id: row.get(0)?,
session_id: row.get(1)?,
command: row.get(2)?,
timestamp: row
.get::<_, String>(3)?
.parse()
.unwrap_or_else(|_| Utc::now()),
directory: row.get(4)?,
redacted: row.get::<_, i32>(5)? != 0,
exit_code: row.get(6)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(commands)
}
#[must_use = "Query results should be used"]
pub fn get_all_commands(&self) -> Result<Vec<CommandEntry>> {
let mut stmt = self.conn.prepare(
"SELECT id, session_id, command, timestamp, directory, redacted, exit_code
FROM commands
ORDER BY timestamp ASC",
)?;
let commands = stmt
.query_map([], |row| {
Ok(CommandEntry {
id: row.get(0)?,
session_id: row.get(1)?,
command: row.get(2)?,
timestamp: row
.get::<_, String>(3)?
.parse()
.unwrap_or_else(|_| Utc::now()),
directory: row.get(4)?,
redacted: row.get::<_, i32>(5)? != 0,
exit_code: row.get(6)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(commands)
}
pub fn get_commands_paginated(&self, offset: usize, limit: usize) -> Result<Vec<CommandEntry>> {
let mut stmt = self.conn.prepare(
"SELECT id, session_id, command, timestamp, directory, redacted, exit_code
FROM commands
WHERE directory != '<imported>'
ORDER BY timestamp DESC
LIMIT ?1 OFFSET ?2",
)?;
let commands = stmt
.query_map(rusqlite::params![limit as i64, offset as i64], |row| {
Ok(CommandEntry {
id: row.get(0)?,
session_id: row.get(1)?,
command: row.get(2)?,
timestamp: row
.get::<_, String>(3)?
.parse()
.unwrap_or_else(|_| Utc::now()),
directory: row.get(4)?,
redacted: row.get::<_, i32>(5)? != 0,
exit_code: row.get(6)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(commands)
}
pub fn get_unique_commands_paginated(
&self,
offset: usize,
limit: usize,
) -> Result<Vec<CommandEntry>> {
let mut stmt = self.conn.prepare(
"SELECT MAX(id), session_id, command, MAX(timestamp) as ts, directory, redacted, exit_code
FROM commands
WHERE directory != '<imported>'
GROUP BY command, directory
ORDER BY ts DESC
LIMIT ?1 OFFSET ?2",
)?;
let commands = stmt
.query_map(rusqlite::params![limit as i64, offset as i64], |row| {
Ok(CommandEntry {
id: row.get(0)?,
session_id: row.get(1)?,
command: row.get(2)?,
timestamp: row
.get::<_, String>(3)?
.parse()
.unwrap_or_else(|_| Utc::now()),
directory: row.get(4)?,
redacted: row.get::<_, i32>(5)? != 0,
exit_code: row.get(6)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(commands)
}
pub fn count_unique_commands(&self) -> Result<usize> {
let count: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM (
SELECT 1 FROM commands
WHERE directory != '<imported>'
GROUP BY command, directory
)",
[],
|row| row.get(0),
)?;
Ok(count as usize)
}
pub fn get_commands_for_directory(&self, directory: &str) -> Result<Vec<CommandEntry>> {
let mut stmt = self.conn.prepare(
"SELECT MAX(id), session_id, command, MAX(timestamp) as ts, directory, redacted, exit_code
FROM commands
WHERE directory = ?1
GROUP BY command
ORDER BY ts DESC",
)?;
let commands = stmt
.query_map(rusqlite::params![directory], |row| {
Ok(CommandEntry {
id: row.get(0)?,
session_id: row.get(1)?,
command: row.get(2)?,
timestamp: row
.get::<_, String>(3)?
.parse()
.unwrap_or_else(|_| Utc::now()),
directory: row.get(4)?,
redacted: row.get::<_, i32>(5)? != 0,
exit_code: row.get(6)?,
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(commands)
}
pub fn get_frequent_commands(&self, limit: usize) -> Result<Vec<(String, usize)>> {
let mut stmt = self.conn.prepare(
"SELECT command, COUNT(*) as cnt
FROM commands
WHERE directory != '<imported>'
GROUP BY command
ORDER BY cnt DESC
LIMIT ?1",
)?;
let results = stmt
.query_map(rusqlite::params![limit as i64], |row| {
Ok((row.get::<_, String>(0)?, row.get::<_, i64>(1)? as usize))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(results)
}
pub fn count_commands(&self) -> Result<usize> {
let count: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM commands WHERE directory != '<imported>'",
[],
|row| row.get(0),
)?;
Ok(count as usize)
}
pub fn get_stats(&self) -> Result<DatabaseStats> {
let total_commands: i64 =
self.conn
.query_row("SELECT COUNT(*) FROM commands", [], |row| row.get(0))?;
let total_sessions: i64 =
self.conn
.query_row("SELECT COUNT(*) FROM sessions", [], |row| row.get(0))?;
let total_hosts: i64 = self
.conn
.query_row("SELECT COUNT(*) FROM hosts", [], |row| row.get(0))?;
let redacted_commands: i64 = self.conn.query_row(
"SELECT COUNT(*) FROM commands WHERE redacted = 1",
[],
|row| row.get(0),
)?;
let stored_tokens: i64 = self
.conn
.query_row("SELECT COUNT(*) FROM tokens", [], |row| row.get(0))?;
let oldest_entry: Option<String> = self
.conn
.query_row(
"SELECT timestamp FROM commands ORDER BY timestamp ASC LIMIT 1",
[],
|row| row.get(0),
)
.optional()?;
let newest_entry: Option<String> = self
.conn
.query_row(
"SELECT timestamp FROM commands ORDER BY timestamp DESC LIMIT 1",
[],
|row| row.get(0),
)
.optional()?;
Ok(DatabaseStats {
total_commands: total_commands as usize,
total_sessions: total_sessions as usize,
total_hosts: total_hosts as usize,
redacted_commands: redacted_commands as usize,
stored_tokens: stored_tokens as usize,
oldest_entry: oldest_entry.and_then(|s| s.parse().ok()),
newest_entry: newest_entry.and_then(|s| s.parse().ok()),
})
}
pub fn get_hosts(&self) -> Result<Vec<Host>> {
let mut stmt = self
.conn
.prepare("SELECT id, hostname, created_at FROM hosts ORDER BY hostname")?;
let hosts = stmt
.query_map([], |row| {
Ok(Host {
id: row.get(0)?,
hostname: row.get(1)?,
created_at: row
.get::<_, String>(2)?
.parse()
.unwrap_or_else(|_| Utc::now()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(hosts)
}
pub fn get_sessions_for_host(&self, host_id: HostId) -> Result<Vec<Session>> {
let mut stmt = self.conn.prepare(
"SELECT id, host_id, started_at, ended_at
FROM sessions
WHERE host_id = ?1
ORDER BY started_at DESC",
)?;
let sessions = stmt
.query_map(params![host_id.as_i64()], |row| {
Ok(Session {
id: SessionId::new(row.get(0)?),
host_id: HostId::new(row.get(1)?),
started_at: row
.get::<_, String>(2)?
.parse()
.unwrap_or_else(|_| Utc::now()),
ended_at: row
.get::<_, Option<String>>(3)?
.and_then(|s| s.parse().ok()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(sessions)
}
pub fn import_from_bash_history(&mut self, bash_history_path: &Path) -> Result<usize> {
let content = std::fs::read_to_string(bash_history_path)?;
let mut imported_count = 0;
let now = Utc::now();
for line in content.lines() {
let line = line.trim();
if line.is_empty() || line.starts_with('#') {
continue;
}
self.add_command(line, "<imported>", now, false, None)?;
imported_count += 1;
}
Ok(imported_count)
}
pub fn import_from_zsh_history(&mut self, zsh_history_path: &Path) -> Result<usize> {
let content = std::fs::read_to_string(zsh_history_path)?;
let mut imported_count = 0;
let re = regex::Regex::new(r"^: (\d+):\d+;(.*)").unwrap();
for line in content.lines() {
if let Some(caps) = re.captures(line) {
let timestamp_str = caps.get(1).unwrap().as_str();
let command = caps.get(2).unwrap().as_str();
if let Ok(timestamp_secs) = timestamp_str.parse::<i64>()
&& let Some(datetime) = DateTime::from_timestamp(timestamp_secs, 0)
{
self.add_command(command, "<imported>", datetime, false, None)?;
imported_count += 1;
}
}
}
Ok(imported_count)
}
pub fn merge_from_database(&mut self, other_db_path: &Path) -> Result<usize> {
let other_conn = Connection::open(other_db_path)?;
let mut imported_count = 0;
let mut stmt = other_conn.prepare(
"SELECT c.command, c.timestamp, c.directory, c.redacted, c.exit_code,
s.started_at, h.hostname
FROM commands c
JOIN sessions s ON c.session_id = s.id
JOIN hosts h ON s.host_id = h.id
ORDER BY c.timestamp ASC",
)?;
let commands: Vec<_> = stmt
.query_map([], |row| {
Ok((
row.get::<_, String>(0)?,
row.get::<_, String>(1)?,
row.get::<_, String>(2)?,
row.get::<_, i32>(3)? != 0,
row.get::<_, Option<i32>>(4)?,
))
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
for (command, timestamp_str, directory, redacted, exit_code) in commands {
if let Ok(timestamp) = timestamp_str.parse() {
self.add_command(&command, &directory, timestamp, redacted, exit_code)?;
imported_count += 1;
}
}
Ok(imported_count)
}
pub fn add_alias(&self, alias: &str, command: &str, description: &str) -> Result<()> {
let now = Utc::now().to_rfc3339();
self.conn.execute(
"INSERT INTO aliases (alias, command, description, date_created, date_updated)
VALUES (?1, ?2, ?3, ?4, ?5)",
params![alias, command, description, now, now],
)?;
Ok(())
}
pub fn update_alias(
&self,
alias: &str,
command: &str,
description: Option<&str>,
) -> Result<()> {
let now = Utc::now().to_rfc3339();
if let Some(desc) = description {
self.conn.execute(
"UPDATE aliases SET command = ?1, description = ?2, date_updated = ?3
WHERE alias = ?4",
params![command, desc, now, alias],
)?;
} else {
self.conn.execute(
"UPDATE aliases SET command = ?1, date_updated = ?2 WHERE alias = ?3",
params![command, now, alias],
)?;
}
Ok(())
}
pub fn remove_alias(&self, alias: &str) -> Result<()> {
self.conn
.execute("DELETE FROM aliases WHERE alias = ?1", params![alias])?;
Ok(())
}
#[must_use = "Alias list should be used"]
pub fn list_aliases(&self) -> Result<Vec<Alias>> {
let mut stmt = self.conn.prepare(
"SELECT alias, command, description, date_created, date_updated
FROM aliases ORDER BY alias ASC",
)?;
let aliases = stmt
.query_map([], |row| {
Ok(Alias {
alias: row.get(0)?,
command: row.get(1)?,
description: row.get(2)?,
date_created: row
.get::<_, String>(3)?
.parse()
.unwrap_or_else(|_| Utc::now()),
date_updated: row
.get::<_, String>(4)?
.parse()
.unwrap_or_else(|_| Utc::now()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(aliases)
}
pub fn sync_aliases(&self, aliases: &[(String, String)]) -> Result<usize> {
let now = Utc::now().to_rfc3339();
let mut count = 0;
for (name, command) in aliases {
self.conn.execute(
"INSERT INTO aliases (alias, command, description, date_created, date_updated)
VALUES (?1, ?2, '', ?3, ?4)
ON CONFLICT(alias) DO UPDATE SET command = ?2, date_updated = ?4",
params![name, command, now, now],
)?;
count += 1;
}
Ok(count)
}
pub fn clear(&self) -> Result<()> {
self.conn.execute("DELETE FROM tokens", [])?;
self.conn.execute("DELETE FROM commands", [])?;
self.conn.execute("DELETE FROM sessions", [])?;
self.conn.execute("DELETE FROM hosts", [])?;
Ok(())
}
pub fn delete_command(&self, id: CommandId) -> Result<()> {
self.conn
.execute("DELETE FROM commands WHERE id = ?1", [id.0])?;
Ok(())
}
pub fn get_all_sessions(&self) -> Result<Vec<Session>> {
let mut stmt = self.conn.prepare(
"SELECT id, host_id, started_at, ended_at
FROM sessions
ORDER BY started_at DESC",
)?;
let sessions = stmt
.query_map([], |row| {
Ok(Session {
id: SessionId::new(row.get(0)?),
host_id: HostId::new(row.get(1)?),
started_at: row
.get::<_, String>(2)?
.parse()
.unwrap_or_else(|_| Utc::now()),
ended_at: row
.get::<_, Option<String>>(3)?
.and_then(|s| s.parse().ok()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(sessions)
}
pub fn get_all_tokens(&self) -> Result<Vec<Token>> {
let mut stmt = self.conn.prepare(
"SELECT id, command_id, token_type, placeholder, original_value, created_at
FROM tokens
ORDER BY created_at DESC",
)?;
let tokens = stmt
.query_map([], |row| {
Ok(Token {
id: row.get(0)?,
command_id: CommandId::new(row.get(1)?),
token_type: row.get(2)?,
placeholder: row.get(3)?,
original_value: row.get(4)?,
created_at: row
.get::<_, String>(5)?
.parse()
.unwrap_or_else(|_| Utc::now()),
})
})?
.collect::<rusqlite::Result<Vec<_>>>()?;
Ok(tokens)
}
pub fn delete_host(&self, id: HostId) -> Result<()> {
self.conn
.execute("DELETE FROM hosts WHERE id = ?1", [id.as_i64()])?;
Ok(())
}
pub fn delete_session(&self, id: &str) -> Result<()> {
self.conn
.execute("DELETE FROM sessions WHERE id = ?1", [id])?;
Ok(())
}
pub fn delete_token(&self, id: i64) -> Result<()> {
self.conn
.execute("DELETE FROM tokens WHERE id = ?1", [id])?;
Ok(())
}
}
#[cfg(test)]
mod tests {
use super::*;
use tempfile::NamedTempFile;
#[test]
fn test_database_creation() {
let temp_file = NamedTempFile::new().unwrap();
let db = Database::new(temp_file.path()).unwrap();
let stats = db.get_stats().unwrap();
assert_eq!(stats.total_commands, 0);
}
#[test]
fn test_add_command() {
let temp_file = NamedTempFile::new().unwrap();
let mut db = Database::new(temp_file.path()).unwrap();
let cmd_id = db
.add_command("ls -la", "/home/user", Utc::now(), false, Some(0))
.unwrap();
assert!(cmd_id > 0);
let stats = db.get_stats().unwrap();
assert_eq!(stats.total_commands, 1);
}
#[test]
fn test_token_storage() {
let temp_file = NamedTempFile::new().unwrap();
let mut db = Database::new(temp_file.path()).unwrap();
let cmd_id = db
.add_command("echo password123", "/home", Utc::now(), true, None)
.unwrap();
db.store_token(cmd_id, "password", "<redacted>", "password123")
.unwrap();
let tokens = db.get_tokens_for_command(CommandId::new(cmd_id)).unwrap();
assert_eq!(tokens.len(), 1);
assert_eq!(tokens[0].original_value, "password123");
}
#[test]
fn test_alias_crud() {
let temp_file = NamedTempFile::new().unwrap();
let db = Database::new(temp_file.path()).unwrap();
db.add_alias("ll", "ls -la", "long listing").unwrap();
let aliases = db.list_aliases().unwrap();
assert_eq!(aliases.len(), 1);
assert_eq!(aliases[0].alias, "ll");
assert_eq!(aliases[0].command, "ls -la");
assert_eq!(aliases[0].description, "long listing");
db.update_alias("ll", "ls -lah", Some("long listing with human sizes"))
.unwrap();
let aliases = db.list_aliases().unwrap();
assert_eq!(aliases[0].command, "ls -lah");
assert_eq!(aliases[0].description, "long listing with human sizes");
db.remove_alias("ll").unwrap();
let aliases = db.list_aliases().unwrap();
assert!(aliases.is_empty());
}
#[test]
fn test_alias_sync() {
let temp_file = NamedTempFile::new().unwrap();
let db = Database::new(temp_file.path()).unwrap();
let aliases = vec![
("ll".to_string(), "ls -la".to_string()),
("gs".to_string(), "git status".to_string()),
];
let count = db.sync_aliases(&aliases).unwrap();
assert_eq!(count, 2);
let updated = vec![("ll".to_string(), "ls -lah".to_string())];
db.sync_aliases(&updated).unwrap();
let all = db.list_aliases().unwrap();
assert_eq!(all.len(), 2);
let ll = all.iter().find(|a| a.alias == "ll").unwrap();
assert_eq!(ll.command, "ls -lah");
}
}