Hexor/khm
1
0
Fork 0
mirror of https://github.com/house-of-vanity/khm.git synced 2025-07-06 23:04:07 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Improve logging and error handling

Browse Source
This commit is contained in:
AB
2024-07-09 02:28:44 +03:00
parent cd1ba739ab
commit 8991a0579b
3 changed files with 74 additions and 28 deletions
Download Patch File Download Diff File Expand all files Collapse all files

32
src/client.rs
View File

@ -3,6 +3,7 @@ use serde::{Deserialize, Serialize};
use std::fs::File;
use std::io::{self, BufRead, Write};
use std::path::Path;
use log::{info, error};
#[derive(Serialize, Deserialize, Clone, Debug)]
struct SshKey {
@ -17,15 +18,21 @@ fn read_known_hosts(file_path: &str) -> io::Result<Vec<SshKey>> {
let mut keys = Vec::new();
for line in reader.lines() {
if let Ok(line) = line {
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.len() >= 2 {
let server = parts[0].to_string();
let public_key = parts[1..].join(" ");
keys.push(SshKey { server, public_key });
match line {
Ok(line) => {
let parts: Vec<&str> = line.split_whitespace().collect();
if parts.len() >= 2 {
let server = parts[0].to_string();
let public_key = parts[1..].join(" ");
keys.push(SshKey { server, public_key });
}
},
Err(e) => {
error!("Error reading line from known_hosts file: {}", e);
}
}
}
info!("Read {} keys from known_hosts file", keys.len());
Ok(keys)
}
@ -36,6 +43,7 @@ fn write_known_hosts(file_path: &str, keys: &[SshKey]) -> io::Result<()> {
for key in keys {
writeln!(file, "{} {}", key.server, key.public_key)?;
}
info!("Wrote {} keys to known_hosts file", keys.len());
Ok(())
}
@ -46,9 +54,9 @@ async fn send_keys_to_server(host: &str, keys: Vec<SshKey>) -> Result<(), reqwes
let response = client.post(&url).json(&keys).send().await?;
if response.status().is_success() {
println!("Keys successfully sent to server.");
info!("Keys successfully sent to server.");
} else {
eprintln!(
error!(
"Failed to send keys to server. Status: {}",
response.status()
);
@ -64,9 +72,10 @@ async fn get_keys_from_server(host: &str) -> Result<Vec<SshKey>, reqwest::Error>
if response.status().is_success() {
let keys: Vec<SshKey> = response.json().await?;
info!("Received {} keys from server", keys.len());
Ok(keys)
} else {
eprintln!(
error!(
"Failed to get keys from server. Status: {}",
response.status()
);
@ -75,21 +84,26 @@ async fn get_keys_from_server(host: &str) -> Result<Vec<SshKey>, reqwest::Error>
}
pub async fn run_client(args: crate::Args) -> std::io::Result<()> {
info!("Client mode: Reading known_hosts file");
let keys = read_known_hosts(&args.known_hosts).expect("Failed to read known hosts file");
let host = args.host.expect("host is required in client mode");
info!("Client mode: Sending keys to server at {}", host);
send_keys_to_server(&host, keys)
.await
.expect("Failed to send keys to server");
if args.in_place {
info!("Client mode: In-place update is enabled. Fetching keys from server.");
let server_keys = get_keys_from_server(&host)
.await
.expect("Failed to get keys from server");
info!("Client mode: Writing updated known_hosts file");
write_known_hosts(&args.known_hosts, &server_keys)
.expect("Failed to write known hosts file");
}
info!("Client mode: Finished operations");
Ok(())
}

18
src/main.rs
View File

@ -3,6 +3,7 @@ mod server;
use clap::Parser;
use env_logger;
use log::{info, error};
/// This application manages SSH keys and flows, either as a server or client.
/// In server mode, it stores keys and flows in a PostgreSQL database.
@ -19,7 +20,7 @@ use env_logger;
khm --server --ip 0.0.0.0 --port 1337 --db-host psql.psql.svc --db-name khm --db-user admin --db-password <SECRET> --flows work,home\n\
\n\
Running in client mode to send diff and sync ~/.ssh/known_hosts with remote flow in place:\n\
khm --host http://kh.example.com:8080 --known_hosts ~/.ssh/known_hosts --in-place\n\
khm --host http://kh.example.com:8080 --known-hosts ~/.ssh/known_hosts --in-place\n\
\n\
"
)]
@ -109,12 +110,23 @@ struct Args {
#[actix_web::main]
async fn main() -> std::io::Result<()> {
env_logger::init();
info!("Starting SSH Key Manager");
let args = Args::parse();
if args.server {
server::run_server(args).await
info!("Running in server mode");
if let Err(e) = server::run_server(args).await {
error!("Failed to run server: {}", e);
}
} else {
client::run_client(args).await
info!("Running in client mode");
if let Err(e) = client::run_client(args).await {
error!("Failed to run client: {}", e);
}
}
info!("Application has exited");
Ok(())
}

52
src/server.rs
View File

@ -1,5 +1,5 @@
use actix_web::{web, App, HttpResponse, HttpServer, Responder};
use log;
use log::{info, error};
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
@ -23,8 +23,7 @@ pub type Flows = Arc<Mutex<Vec<Flow>>>;
pub fn is_valid_ssh_key(key: &str) -> bool {
let rsa_re = Regex::new(r"^ssh-rsa AAAA[0-9A-Za-z+/]+[=]{0,3}( .+)?$").unwrap();
let dsa_re = Regex::new(r"^ssh-dss AAAA[0-9A-Za-z+/]+[=]{0,3}( .+)?$").unwrap();
let ecdsa_re =
Regex::new(r"^ecdsa-sha2-nistp(256|384|521) AAAA[0-9A-Za-z+/]+[=]{0,3}( .+)?$").unwrap();
let ecdsa_re = Regex::new(r"^ecdsa-sha2-nistp(256|384|521) AAAA[0-9A-Za-z+/]+[=]{0,3}( .+)?$").unwrap();
let ed25519_re = Regex::new(r"^ssh-ed25519 AAAA[0-9A-Za-z+/]+[=]{0,3}( .+)?$").unwrap();
rsa_re.is_match(key)
@ -51,12 +50,14 @@ pub async fn insert_key_if_not_exists(
&[&row.get::<_, i32>(0)],
)
.await?;
info!("Updated existing key for server: {}", key.server);
Ok(row.get(0))
} else {
let row = client.query_one(
"INSERT INTO public.keys (host, key, updated) VALUES ($1, $2, NOW()) RETURNING key_id",
&[&key.server, &key.public_key]
).await?;
info!("Inserted new key for server: {}", key.server);
Ok(row.get(0))
}
}
@ -72,6 +73,7 @@ pub async fn insert_flow_key(
&[&flow_name, &key_id],
)
.await?;
info!("Inserted key_id {} into flow: {}", key_id, flow_name);
Ok(())
}
@ -106,6 +108,7 @@ pub async fn get_keys_from_db(client: &Client) -> Result<Vec<Flow>, tokio_postgr
}
}
info!("Retrieved {} flows from database", flows_map.len());
Ok(flows_map.into_values().collect())
}
@ -116,14 +119,17 @@ pub async fn get_keys(
) -> impl Responder {
let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&flow_id_str) {
error!("Flow ID not allowed: {}", flow_id_str);
return HttpResponse::Forbidden().body("Flow ID not allowed");
}
let flows = flows.lock().unwrap();
if let Some(flow) = flows.iter().find(|flow| flow.name == flow_id_str) {
let servers: Vec<&SshKey> = flow.servers.iter().collect();
info!("Returning {} keys for flow: {}", servers.len(), flow_id_str);
HttpResponse::Ok().json(servers)
} else {
error!("Flow ID not found: {}", flow_id_str);
HttpResponse::NotFound().body("Flow ID not found")
}
}
@ -137,11 +143,13 @@ pub async fn add_keys(
) -> impl Responder {
let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&flow_id_str) {
error!("Flow ID not allowed: {}", flow_id_str);
return HttpResponse::Forbidden().body("Flow ID not allowed");
}
for new_key in new_keys.iter() {
if !is_valid_ssh_key(&new_key.public_key) {
error!("Invalid SSH key format for server: {}", new_key.server);
return HttpResponse::BadRequest().body(format!(
"Invalid SSH key format for server: {}",
new_key.server
@ -151,31 +159,38 @@ pub async fn add_keys(
match insert_key_if_not_exists(&db_client, new_key).await {
Ok(key_id) => {
if let Err(e) = insert_flow_key(&db_client, &flow_id_str, key_id).await {
log::error!("Failed to insert flow key into database: {}", e);
error!("Failed to insert flow key into database: {}", e);
return HttpResponse::InternalServerError()
.body("Failed to insert flow key into database");
}
}
Err(e) => {
log::error!("Failed to insert key into database: {}", e);
error!("Failed to insert key into database: {}", e);
return HttpResponse::InternalServerError()
.body("Failed to insert key into database");
}
}
}
// Refresh the flows data from the database
let updated_flows = get_keys_from_db(&db_client)
.await
.unwrap_or_else(|_| Vec::new());
let updated_flows = match get_keys_from_db(&db_client).await {
Ok(flows) => flows,
Err(e) => {
error!("Failed to get updated flows from database: {}", e);
return HttpResponse::InternalServerError()
.body("Failed to refresh flows from database");
}
};
let mut flows_guard = flows.lock().unwrap();
*flows_guard = updated_flows;
let updated_flow = flows_guard.iter().find(|flow| flow.name == flow_id_str);
if let Some(flow) = updated_flow {
let servers: Vec<&SshKey> = flow.servers.iter().collect();
info!("Updated flow: {} with {} keys", flow_id_str, servers.len());
HttpResponse::Ok().json(servers)
} else {
error!("Flow ID not found after update: {}", flow_id_str);
HttpResponse::NotFound().body("Flow ID not found")
}
}
@ -197,13 +212,17 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> {
// Spawn a new thread to run the database connection
tokio::spawn(async move {
if let Err(e) = connection.await {
eprintln!("Connection error: {}", e);
error!("Connection error: {}", e);
}
});
let mut initial_flows = get_keys_from_db(&db_client)
.await
.unwrap_or_else(|_| Vec::new());
let mut initial_flows = match get_keys_from_db(&db_client).await {
Ok(flows) => flows,
Err(e) => {
error!("Failed to get initial flows from database: {}", e);
Vec::new()
}
};
// Ensure all allowed flows are initialized
for allowed_flow in &args.flows {
@ -218,6 +237,7 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> {
let flows: Flows = Arc::new(Mutex::new(initial_flows));
let allowed_flows = web::Data::new(args.flows);
info!("Starting HTTP server on {}:{}", args.ip, args.port);
HttpServer::new(move || {
App::new()
.app_data(web::Data::new(flows.clone()))
@ -226,7 +246,7 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> {
.route("/{flow_id}/keys", web::get().to(get_keys))
.route("/{flow_id}/keys", web::post().to(add_keys))
})
.bind((args.ip.as_str(), args.port))?
.run()
.await
.bind((args.ip.as_str(), args.port))?
.run()
.await
}