diff --git a/Cargo.toml b/Cargo.toml index 450a5f4..4671774 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -31,7 +31,7 @@ rand = "0.8.5" chrono = "0.4.23" users = "0.11.0" serde = { version = "1.0.152", features = ["derive"] } -serde_json = "1.0.92" +serde_json = { version = "1.0.92", features = ["arbitrary_precision"] } regex = "1.7.1" hostname = "0.3.1" rusqlite = { version = "0.28.0", features = ["bundled", "functions"] } diff --git a/src/lib.rs b/src/lib.rs index 6db635f..4b7fc90 100644 --- a/src/lib.rs +++ b/src/lib.rs @@ -19,7 +19,7 @@ use bstr::{BString, ByteSlice}; use chrono::prelude::{Local, TimeZone}; use itertools::Itertools; use regex::bytes::Regex; -use rusqlite::{functions::FunctionFlags, Connection, Error, Result, Transaction}; +use rusqlite::{functions::FunctionFlags, Connection, Error, Result, Row, Transaction}; use serde::{Deserialize, Serialize}; type BoxError = Box; @@ -226,9 +226,9 @@ fn dedup_invocations(invocations: Vec) -> Vec { } } -pub struct InvocationExport { +pub struct InvocationDatabaseRow { pub session_id: i64, - pub full_command: Vec, + pub command: Vec, pub shellname: String, pub working_directory: Option>, pub hostname: Option>, @@ -238,21 +238,89 @@ pub struct InvocationExport { pub end_unix_timestamp: Option, } -pub fn json_export(rows: &[InvocationExport]) -> Result<(), Box> { - let invocations: Vec = rows - .iter() - .map(|row| Invocation { - command: BString::from(row.full_command.as_slice()), +impl InvocationDatabaseRow { + pub fn from_row(row: &Row) -> Result { + Ok(InvocationDatabaseRow { + session_id: row.get("session_id")?, + command: row.get("full_command")?, + shellname: row.get("shellname")?, + working_directory: row.get("working_directory")?, + hostname: row.get("hostname")?, + username: row.get("username")?, + exit_status: row.get("exit_status")?, + start_unix_timestamp: row.get("start_unix_timestamp")?, + end_unix_timestamp: row.get("end_unix_timestamp")?, + }) + } +} + +// Create a pretty export string that gets serialized as an array of +// bytes only if it isn't valid UTF-8; this makes the json export +// prettier. +#[derive(Debug, Eq, PartialEq, Serialize, Deserialize, Clone)] +#[serde(untagged)] +enum PrettyExportString { + Readable(String), + Encoded(Vec), +} + +impl From<&[u8]> for PrettyExportString { + fn from(bytes: &[u8]) -> Self { + match str::from_utf8(bytes) { + Ok(v) => Self::Readable(v.to_string()), + _ => Self::Encoded(bytes.to_vec()), + } + } +} + +impl From>> for PrettyExportString { + fn from(bytes: Option<&Vec>) -> Self { + // TODO: make this pretty + if let Some(v) = bytes { + match str::from_utf8(v.as_slice()) { + Ok(v) => Self::Readable(v.to_string()), + _ => Self::Encoded(v.to_vec()), + } + } else { + PrettyExportString::Readable("".to_string()) + } + } +} + +// A copy of Invocation byt witgh PrettyExportString instead of +// BString; TODO: reconcile into one class perhaps? +#[derive(Debug, Eq, PartialEq, Serialize, Deserialize, Clone)] +struct InvocationJsonExporter { + session_id: i64, + command: PrettyExportString, + shellname: String, + working_directory: PrettyExportString, + hostname: PrettyExportString, + username: PrettyExportString, + exit_status: Option, + start_unix_timestamp: Option, + end_unix_timestamp: Option, +} + +impl InvocationJsonExporter { + fn new(row: &InvocationDatabaseRow) -> Self { + Self { + session_id: row.session_id, + command: row.command.as_slice().into(), shellname: row.shellname.clone(), - hostname: row.hostname.clone().map(|v| BString::from(v.as_slice())), - username: row.username.clone().map(|v| BString::from(v.as_slice())), - working_directory: row.working_directory.clone().map(|v| BString::from(v.as_slice())), + working_directory: row.working_directory.as_ref().into(), + hostname: row.hostname.as_ref().into(), + username: row.username.as_ref().into(), exit_status: row.exit_status, start_unix_timestamp: row.start_unix_timestamp, end_unix_timestamp: row.end_unix_timestamp, - session_id: row.session_id, - }) - .collect(); + } + } +} + +pub fn json_export(rows: &[InvocationDatabaseRow]) -> Result<(), Box> { + let invocations: Vec = + rows.iter().map(InvocationJsonExporter::new).collect(); serde_json::to_writer(io::stdout(), &invocations)?; Ok(()) } @@ -261,7 +329,7 @@ pub fn json_export(rows: &[InvocationExport]) -> Result<(), Box String>, + displayer: Box String>, } fn time_display_helper(t: Option) -> String { @@ -283,7 +351,7 @@ fn displayers() -> HashMap<&'static str, QueryResultColumnDisplayer> { "command", QueryResultColumnDisplayer { header: "Command", - displayer: Box::new(|row| binary_display_helper(&row.full_command)), + displayer: Box::new(|row| binary_display_helper(&row.command)), }, ); ret.insert( @@ -360,7 +428,7 @@ fn displayers() -> HashMap<&'static str, QueryResultColumnDisplayer> { pub fn present_results_human_readable( fields: &[&str], - rows: &[InvocationExport], + rows: &[InvocationDatabaseRow], suppress_headers: bool, ) -> Result<(), Box> { let displayers = displayers(); diff --git a/src/main.rs b/src/main.rs index 92400a1..00071ce 100644 --- a/src/main.rs +++ b/src/main.rs @@ -251,21 +251,8 @@ SELECT session_id, full_command, shellname, hostname, username, working_director FROM command_history h ORDER BY id"#, )?; - let rows: Result, _> = stmt - .query_map([], |row| { - Ok(pxh::InvocationExport { - session_id: row.get("session_id")?, - full_command: row.get("full_command")?, - shellname: row.get("shellname")?, - hostname: row.get("hostname")?, - username: row.get("username")?, - working_directory: row.get("working_directory")?, - exit_status: row.get("exit_status")?, - start_unix_timestamp: row.get("start_unix_timestamp")?, - end_unix_timestamp: row.get("end_unix_timestamp")?, - }) - })? - .collect(); + let rows: Result, _> = + stmt.query_map([], pxh::InvocationDatabaseRow::from_row)?.collect(); let rows = rows?; pxh::json_export(&rows)?; Ok(()) @@ -426,21 +413,8 @@ SELECT session_id, full_command, shellname, working_directory, hostname, usernam ORDER BY ch_start_unix_timestamp DESC, ch_id DESC "#)?; - let rows: Result, _> = stmt - .query_map((), |row| { - Ok(pxh::InvocationExport { - session_id: row.get("session_id")?, - full_command: row.get("full_command")?, - shellname: row.get("shellname")?, - working_directory: row.get("working_directory")?, - hostname: row.get("hostname")?, - username: row.get("username")?, - exit_status: row.get("exit_status")?, - start_unix_timestamp: row.get("start_unix_timestamp")?, - end_unix_timestamp: row.get("end_unix_timestamp")?, - }) - })? - .collect(); + let rows: Result, _> = + stmt.query_map([], pxh::InvocationDatabaseRow::from_row)?.collect(); let mut rows = rows?; rows.reverse(); if self.verbose {