Hexor/khm
1
0
Fork 0
mirror of https://github.com/house-of-vanity/khm.git synced 2025-08-21 14:27:14 +00:00
Code Issues Actions Packages Projects Releases Wiki Activity

Fix metadata

Browse Source
This commit is contained in:
AB
2024-07-07 21:13:19 +03:00
parent 34dccc5772
commit d5fa1e0d9a
5 changed files with 158 additions and 82 deletions
Download Patch File Download Diff File Expand all files Collapse all files

34
Cargo.lock generated
View File

@@ -1051,6 +1051,23 @@ dependencies = [
"wasm-bindgen", "wasm-bindgen",
] ]
[[package]]
name = "khm"
version = "0.1.0"
dependencies = [
"actix-web",
"chrono",
"clap",
"env_logger",
"log",
"regex",
"reqwest",
"serde",
"serde_json",
"tokio",
"tokio-postgres",
]
[[package]] [[package]]
name = "language-tags" name = "language-tags"
version = "0.3.2" version = "0.3.2"
@@ -2062,23 +2079,6 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index" source = "registry+https://github.com/rust-lang/crates.io-index"
checksum = "e4259d9d4425d9f0661581b804cb85fe66a4c631cadd8f490d1c13a35d5d9291" checksum = "e4259d9d4425d9f0661581b804cb85fe66a4c631cadd8f490d1c13a35d5d9291"
[[package]]
name = "untitled"
version = "0.1.0"
dependencies = [
"actix-web",
"chrono",
"clap",
"env_logger",
"log",
"regex",
"reqwest",
"serde",
"serde_json",
"tokio",
"tokio-postgres",
]
[[package]] [[package]]
name = "untrusted" name = "untrusted"
version = "0.9.0" version = "0.9.0"

3
Cargo.toml
View File

@@ -1,7 +1,8 @@
[package] [package]
name = "untitled" name = "khm"
version = "0.1.0" version = "0.1.0"
edition = "2021" edition = "2021"
authors = ["AB <ab@hexor.cy>"]
[dependencies] [dependencies]
actix-web = "4" actix-web = "4"

28
src/client.rs
View File

@@ -1,8 +1,8 @@
use reqwest::Client;
use serde::{Deserialize, Serialize};
use std::fs::File; use std::fs::File;
use std::io::{self, BufRead, Write}; use std::io::{self, BufRead, Write};
use std::path::Path; use std::path::Path;
use reqwest::Client;
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Clone, Debug)] #[derive(Serialize, Deserialize, Clone, Debug)]
struct SshKey { struct SshKey {
@@ -43,15 +43,15 @@ fn write_known_hosts(file_path: &str, keys: &[SshKey]) -> io::Result<()> {
async fn send_keys_to_server(host: &str, keys: Vec<SshKey>) -> Result<(), reqwest::Error> { async fn send_keys_to_server(host: &str, keys: Vec<SshKey>) -> Result<(), reqwest::Error> {
let client = Client::new(); let client = Client::new();
let url = format!("{}/keys", host); let url = format!("{}/keys", host);
let response = client.post(&url) let response = client.post(&url).json(&keys).send().await?;
.json(&keys)
.send()
.await?;
if response.status().is_success() { if response.status().is_success() {
println!("Keys successfully sent to server."); println!("Keys successfully sent to server.");
} else { } else {
eprintln!("Failed to send keys to server. Status: {}", response.status()); eprintln!(
"Failed to send keys to server. Status: {}",
response.status()
);
} }
Ok(()) Ok(())
@@ -66,21 +66,25 @@ async fn get_keys_from_server(host: &str) -> Result<Vec<SshKey>, reqwest::Error>
let keys: Vec<SshKey> = response.json().await?; let keys: Vec<SshKey> = response.json().await?;
Ok(keys) Ok(keys)
} else { } else {
eprintln!("Failed to get keys from server. Status: {}", response.status()); eprintln!(
"Failed to get keys from server. Status: {}",
response.status()
);
Ok(vec![]) Ok(vec![])
} }
} }
pub async fn run_client(args: crate::Args) -> std::io::Result<()> { pub async fn run_client(args: crate::Args) -> std::io::Result<()> {
let keys = read_known_hosts(&args.known_hosts) let keys = read_known_hosts(&args.known_hosts).expect("Failed to read known hosts file");
.expect("Failed to read known hosts file");
let host = args.host.expect("host is required in client mode"); let host = args.host.expect("host is required in client mode");
send_keys_to_server(&host, keys).await send_keys_to_server(&host, keys)
.await
.expect("Failed to send keys to server"); .expect("Failed to send keys to server");
if args.in_place { if args.in_place {
let server_keys = get_keys_from_server(&host).await let server_keys = get_keys_from_server(&host)
.await
.expect("Failed to get keys from server"); .expect("Failed to get keys from server");
write_known_hosts(&args.known_hosts, &server_keys) write_known_hosts(&args.known_hosts, &server_keys)

65
src/main.rs
View File

@@ -1,5 +1,5 @@
mod server;
mod client; mod client;
mod server;
use clap::Parser; use clap::Parser;
use env_logger; use env_logger;
@@ -9,8 +9,8 @@ use env_logger;
/// In client mode, it sends keys to the server and can update the known_hosts file with keys from the server. /// In client mode, it sends keys to the server and can update the known_hosts file with keys from the server.
#[derive(Parser, Debug)] #[derive(Parser, Debug)]
#[command( #[command(
author = "Your Name", author = env!("CARGO_PKG_AUTHORS"),
version = "1.0", version = env!("CARGO_PKG_VERSION"),
about = "SSH Key Manager", about = "SSH Key Manager",
long_about = None, long_about = None,
after_help = "Examples:\n\ after_help = "Examples:\n\
@@ -22,33 +22,64 @@ use env_logger;
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\ \n\
" "
)]struct Args { )]
struct Args {
/// IP address to bind the server or client to (default: 127.0.0.1) /// IP address to bind the server or client to (default: 127.0.0.1)
#[arg(short, long, default_value = "127.0.0.1", help = "Server mode: IP address to bind the server to")] #[arg(
short,
long,
default_value = "127.0.0.1",
help = "Server mode: IP address to bind the server to"
)]
ip: String, ip: String,
/// Port to bind the server or client to (default: 8080) /// Port to bind the server or client to (default: 8080)
#[arg(short, long, default_value = "8080", help = "Server mode: Port to bind the server to")] #[arg(
short,
long,
default_value = "8080",
help = "Server mode: Port to bind the server to"
)]
port: u16, port: u16,
/// Hostname or IP address of the PostgreSQL database (default: 127.0.0.1) /// Hostname or IP address of the PostgreSQL database (default: 127.0.0.1)
#[arg(long, default_value = "127.0.0.1", help = "Server mode: Hostname or IP address of the PostgreSQL database")] #[arg(
long,
default_value = "127.0.0.1",
help = "Server mode: Hostname or IP address of the PostgreSQL database"
)]
db_host: String, db_host: String,
/// Name of the PostgreSQL database (default: khm) /// Name of the PostgreSQL database (default: khm)
#[arg(long, default_value = "khm", help = "Server mode: Name of the PostgreSQL database")] #[arg(
long,
default_value = "khm",
help = "Server mode: Name of the PostgreSQL database"
)]
db_name: String, db_name: String,
/// Username for the PostgreSQL database (required in server mode) /// Username for the PostgreSQL database (required in server mode)
#[arg(long, required_if_eq("server", "true"), help = "Server mode: Username for the PostgreSQL database")] #[arg(
long,
required_if_eq("server", "true"),
help = "Server mode: Username for the PostgreSQL database"
)]
db_user: Option<String>, db_user: Option<String>,
/// Password for the PostgreSQL database (required in server mode) /// Password for the PostgreSQL database (required in server mode)
#[arg(long, required_if_eq("server", "true"), help = "Server mode: Password for the PostgreSQL database")] #[arg(
long,
required_if_eq("server", "true"),
help = "Server mode: Password for the PostgreSQL database"
)]
db_password: Option<String>, db_password: Option<String>,
/// Host address of the server to connect to in client mode (required in client mode) /// Host address of the server to connect to in client mode (required in client mode)
#[arg(long, required_if_eq("server", "false"), help = "Client mode: Host address of the server to connect to")] #[arg(
long,
required_if_eq("server", "false"),
help = "Client mode: Host address of the server to connect to"
)]
host: Option<String>, host: Option<String>,
/// Run in server mode (default: false) /// Run in server mode (default: false)
@@ -56,11 +87,18 @@ use env_logger;
server: bool, server: bool,
/// Path to the known_hosts file (default: ~/.ssh/known_hosts) /// Path to the known_hosts file (default: ~/.ssh/known_hosts)
#[arg(long, default_value = "~/.ssh/known_hosts", help = "Client mode: Path to the known_hosts file")] #[arg(
long,
default_value = "~/.ssh/known_hosts",
help = "Client mode: Path to the known_hosts file"
)]
known_hosts: String, known_hosts: String,
/// Update the known_hosts file with keys from the server after sending keys (default: false) /// Update the known_hosts file with keys from the server after sending keys (default: false)
#[arg(long, help = "Server mode: Sync the known_hosts file with keys from the server")] #[arg(
long,
help = "Server mode: Sync the known_hosts file with keys from the server"
)]
in_place: bool, in_place: bool,
/// Comma-separated list of flows to manage (default: default) /// Comma-separated list of flows to manage (default: default)
@@ -68,7 +106,6 @@ use env_logger;
flows: Vec<String>, flows: Vec<String>,
} }
#[actix_web::main] #[actix_web::main]
async fn main() -> std::io::Result<()> { async fn main() -> std::io::Result<()> {
env_logger::init(); env_logger::init();

110
src/server.rs
View File

@@ -1,10 +1,10 @@
use actix_web::{web, App, HttpResponse, HttpServer, Responder};
use log;
use regex::Regex;
use serde::{Deserialize, Serialize};
use std::collections::HashMap; use std::collections::HashMap;
use std::sync::{Arc, Mutex}; use std::sync::{Arc, Mutex};
use actix_web::{web, App, HttpServer, Responder, HttpResponse}; use tokio_postgres::{Client, NoTls};
use serde::{Deserialize, Serialize};
use regex::Regex;
use tokio_postgres::{NoTls, Client};
use log;
#[derive(Serialize, Deserialize, Clone, Debug)] #[derive(Serialize, Deserialize, Clone, Debug)]
pub struct SshKey { pub struct SshKey {
@@ -23,23 +23,34 @@ pub type Flows = Arc<Mutex<Vec<Flow>>>;
pub fn is_valid_ssh_key(key: &str) -> bool { 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 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 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(); let ed25519_re = Regex::new(r"^ssh-ed25519 AAAA[0-9A-Za-z+/]+[=]{0,3}( .+)?$").unwrap();
rsa_re.is_match(key) || dsa_re.is_match(key) || ecdsa_re.is_match(key) || ed25519_re.is_match(key) rsa_re.is_match(key)
|| dsa_re.is_match(key)
|| ecdsa_re.is_match(key)
|| ed25519_re.is_match(key)
} }
pub async fn insert_key_if_not_exists(client: &Client, key: &SshKey) -> Result<i32, tokio_postgres::Error> { pub async fn insert_key_if_not_exists(
let row = client.query_opt( client: &Client,
"SELECT key_id FROM public.keys WHERE host = $1 AND key = $2", key: &SshKey,
&[&key.server, &key.public_key] ) -> Result<i32, tokio_postgres::Error> {
).await?; let row = client
.query_opt(
"SELECT key_id FROM public.keys WHERE host = $1 AND key = $2",
&[&key.server, &key.public_key],
)
.await?;
if let Some(row) = row { if let Some(row) = row {
client.execute( client
"UPDATE public.keys SET updated = NOW() WHERE key_id = $1", .execute(
&[&row.get::<_, i32>(0)] "UPDATE public.keys SET updated = NOW() WHERE key_id = $1",
).await?; &[&row.get::<_, i32>(0)],
)
.await?;
Ok(row.get(0)) Ok(row.get(0))
} else { } else {
let row = client.query_one( let row = client.query_one(
@@ -50,11 +61,17 @@ pub async fn insert_key_if_not_exists(client: &Client, key: &SshKey) -> Result<i
} }
} }
pub async fn insert_flow_key(client: &Client, flow_name: &str, key_id: i32) -> Result<(), tokio_postgres::Error> { pub async fn insert_flow_key(
client.execute( client: &Client,
"INSERT INTO public.flows (name, key_id) VALUES ($1, $2) ON CONFLICT DO NOTHING", flow_name: &str,
&[&flow_name, &key_id] key_id: i32,
).await?; ) -> Result<(), tokio_postgres::Error> {
client
.execute(
"INSERT INTO public.flows (name, key_id) VALUES ($1, $2) ON CONFLICT DO NOTHING",
&[&flow_name, &key_id],
)
.await?;
Ok(()) Ok(())
} }
@@ -71,15 +88,21 @@ pub async fn get_keys_from_db(client: &Client) -> Result<Vec<Flow>, tokio_postgr
let key: String = row.get(1); let key: String = row.get(1);
let flow: String = row.get(2); let flow: String = row.get(2);
let ssh_key = SshKey { server: host, public_key: key }; let ssh_key = SshKey {
server: host,
public_key: key,
};
if let Some(flow_entry) = flows_map.get_mut(&flow) { if let Some(flow_entry) = flows_map.get_mut(&flow) {
flow_entry.servers.push(ssh_key); flow_entry.servers.push(ssh_key);
} else { } else {
flows_map.insert(flow.clone(), Flow { flows_map.insert(
name: flow, flow.clone(),
servers: vec![ssh_key], Flow {
}); name: flow,
servers: vec![ssh_key],
},
);
} }
} }
@@ -89,7 +112,7 @@ pub async fn get_keys_from_db(client: &Client) -> Result<Vec<Flow>, tokio_postgr
pub async fn get_keys( pub async fn get_keys(
flows: web::Data<Flows>, flows: web::Data<Flows>,
flow_id: web::Path<String>, flow_id: web::Path<String>,
allowed_flows: web::Data<Vec<String>> allowed_flows: web::Data<Vec<String>>,
) -> impl Responder { ) -> impl Responder {
let flow_id_str = flow_id.into_inner(); let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&flow_id_str) { if !allowed_flows.contains(&flow_id_str) {
@@ -110,7 +133,7 @@ pub async fn add_keys(
flow_id: web::Path<String>, flow_id: web::Path<String>,
new_keys: web::Json<Vec<SshKey>>, new_keys: web::Json<Vec<SshKey>>,
db_client: web::Data<Arc<Client>>, db_client: web::Data<Arc<Client>>,
allowed_flows: web::Data<Vec<String>> allowed_flows: web::Data<Vec<String>>,
) -> impl Responder { ) -> impl Responder {
let flow_id_str = flow_id.into_inner(); let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&flow_id_str) { if !allowed_flows.contains(&flow_id_str) {
@@ -119,25 +142,32 @@ pub async fn add_keys(
for new_key in new_keys.iter() { for new_key in new_keys.iter() {
if !is_valid_ssh_key(&new_key.public_key) { if !is_valid_ssh_key(&new_key.public_key) {
return HttpResponse::BadRequest().body(format!("Invalid SSH key format for server: {}", new_key.server)); return HttpResponse::BadRequest().body(format!(
"Invalid SSH key format for server: {}",
new_key.server
));
} }
match insert_key_if_not_exists(&db_client, new_key).await { match insert_key_if_not_exists(&db_client, new_key).await {
Ok(key_id) => { Ok(key_id) => {
if let Err(e) = insert_flow_key(&db_client, &flow_id_str, key_id).await { 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); log::error!("Failed to insert flow key into database: {}", e);
return HttpResponse::InternalServerError().body("Failed to insert flow key into database"); return HttpResponse::InternalServerError()
.body("Failed to insert flow key into database");
} }
}, }
Err(e) => { Err(e) => {
log::error!("Failed to insert key into database: {}", e); log::error!("Failed to insert key into database: {}", e);
return HttpResponse::InternalServerError().body("Failed to insert key into database"); return HttpResponse::InternalServerError()
.body("Failed to insert key into database");
} }
} }
} }
// Refresh the flows data from the 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 = get_keys_from_db(&db_client)
.await
.unwrap_or_else(|_| Vec::new());
let mut flows_guard = flows.lock().unwrap(); let mut flows_guard = flows.lock().unwrap();
*flows_guard = updated_flows; *flows_guard = updated_flows;
@@ -152,7 +182,9 @@ pub async fn add_keys(
pub async fn run_server(args: crate::Args) -> std::io::Result<()> { pub async fn run_server(args: crate::Args) -> std::io::Result<()> {
let db_user = args.db_user.expect("db_user is required in server mode"); let db_user = args.db_user.expect("db_user is required in server mode");
let db_password = args.db_password.expect("db_password is required in server mode"); let db_password = args
.db_password
.expect("db_password is required in server mode");
let db_conn_str = format!( let db_conn_str = format!(
"host={} user={} password={} dbname={}", "host={} user={} password={} dbname={}",
@@ -169,7 +201,9 @@ pub async fn run_server(args: crate::Args) -> std::io::Result<()> {
} }
}); });
let mut initial_flows = get_keys_from_db(&db_client).await.unwrap_or_else(|_| Vec::new()); let mut initial_flows = get_keys_from_db(&db_client)
.await
.unwrap_or_else(|_| Vec::new());
// Ensure all allowed flows are initialized // Ensure all allowed flows are initialized
for allowed_flow in &args.flows { for allowed_flow in &args.flows {
@@ -192,7 +226,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::get().to(get_keys))
.route("/{flow_id}/keys", web::post().to(add_keys)) .route("/{flow_id}/keys", web::post().to(add_keys))
}) })
.bind((args.ip.as_str(), args.port))? .bind((args.ip.as_str(), args.port))?
.run() .run()
.await .await
} }