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

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",
]
[[package]]
name = "khm"
version = "0.1.0"
dependencies = [
"actix-web",
"chrono",
"clap",
"env_logger",
"log",
"regex",
"reqwest",
"serde",
"serde_json",
"tokio",
"tokio-postgres",
]
[[package]]
name = "language-tags"
version = "0.3.2"
@ -2062,23 +2079,6 @@ version = "0.1.1"
source = "registry+https://github.com/rust-lang/crates.io-index"
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]]
name = "untrusted"
version = "0.9.0"

3
Cargo.toml
View File

@ -1,7 +1,8 @@
[package]
name = "untitled"
name = "khm"
version = "0.1.0"
edition = "2021"
authors = ["AB <ab@hexor.cy>"]
[dependencies]
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::io::{self, BufRead, Write};
use std::path::Path;
use reqwest::Client;
use serde::{Deserialize, Serialize};
#[derive(Serialize, Deserialize, Clone, Debug)]
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> {
let client = Client::new();
let url = format!("{}/keys", host);
let response = client.post(&url)
.json(&keys)
.send()
.await?;
let response = client.post(&url).json(&keys).send().await?;
if response.status().is_success() {
println!("Keys successfully sent to server.");
} else {
eprintln!("Failed to send keys to server. Status: {}", response.status());
eprintln!(
"Failed to send keys to server. Status: {}",
response.status()
);
}
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?;
Ok(keys)
} else {
eprintln!("Failed to get keys from server. Status: {}", response.status());
eprintln!(
"Failed to get keys from server. Status: {}",
response.status()
);
Ok(vec![])
}
}
pub async fn run_client(args: crate::Args) -> std::io::Result<()> {
let keys = read_known_hosts(&args.known_hosts)
.expect("Failed to read 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");
send_keys_to_server(&host, keys).await
send_keys_to_server(&host, keys)
.await
.expect("Failed to send keys to server");
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");
write_known_hosts(&args.known_hosts, &server_keys)

65
src/main.rs
View File

@ -1,5 +1,5 @@
mod server;
mod client;
mod server;
use clap::Parser;
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.
#[derive(Parser, Debug)]
#[command(
author = "Your Name",
version = "1.0",
author = env!("CARGO_PKG_AUTHORS"),
version = env!("CARGO_PKG_VERSION"),
about = "SSH Key Manager",
long_about = None,
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\
\n\
"
)]struct Args {
)]
struct Args {
/// 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,
/// 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,
/// 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,
/// 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,
/// 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>,
/// 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>,
/// 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>,
/// Run in server mode (default: false)
@ -56,11 +87,18 @@ use env_logger;
server: bool,
/// 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,
/// 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,
/// Comma-separated list of flows to manage (default: default)
@ -68,7 +106,6 @@ use env_logger;
flows: Vec<String>,
}
#[actix_web::main]
async fn main() -> std::io::Result<()> {
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::sync::{Arc, Mutex};
use actix_web::{web, App, HttpServer, Responder, HttpResponse};
use serde::{Deserialize, Serialize};
use regex::Regex;
use tokio_postgres::{NoTls, Client};
use log;
use tokio_postgres::{Client, NoTls};
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct SshKey {
@ -23,23 +23,34 @@ 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) || 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> {
let row = client.query_opt(
"SELECT key_id FROM public.keys WHERE host = $1 AND key = $2",
&[&key.server, &key.public_key]
).await?;
pub async fn insert_key_if_not_exists(
client: &Client,
key: &SshKey,
) -> Result<i32, tokio_postgres::Error> {
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 {
client.execute(
"UPDATE public.keys SET updated = NOW() WHERE key_id = $1",
&[&row.get::<_, i32>(0)]
).await?;
client
.execute(
"UPDATE public.keys SET updated = NOW() WHERE key_id = $1",
&[&row.get::<_, i32>(0)],
)
.await?;
Ok(row.get(0))
} else {
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> {
client.execute(
"INSERT INTO public.flows (name, key_id) VALUES ($1, $2) ON CONFLICT DO NOTHING",
&[&flow_name, &key_id]
).await?;
pub async fn insert_flow_key(
client: &Client,
flow_name: &str,
key_id: i32,
) -> 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(())
}
@ -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 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) {
flow_entry.servers.push(ssh_key);
} else {
flows_map.insert(flow.clone(), Flow {
name: flow,
servers: vec![ssh_key],
});
flows_map.insert(
flow.clone(),
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(
flows: web::Data<Flows>,
flow_id: web::Path<String>,
allowed_flows: web::Data<Vec<String>>
allowed_flows: web::Data<Vec<String>>,
) -> impl Responder {
let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&flow_id_str) {
@ -110,7 +133,7 @@ pub async fn add_keys(
flow_id: web::Path<String>,
new_keys: web::Json<Vec<SshKey>>,
db_client: web::Data<Arc<Client>>,
allowed_flows: web::Data<Vec<String>>
allowed_flows: web::Data<Vec<String>>,
) -> impl Responder {
let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&flow_id_str) {
@ -119,25 +142,32 @@ pub async fn add_keys(
for new_key in new_keys.iter() {
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 {
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);
return HttpResponse::InternalServerError().body("Failed to insert flow key into database");
return HttpResponse::InternalServerError()
.body("Failed to insert flow key into database");
}
},
}
Err(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
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();
*flows_guard = updated_flows;
@ -152,7 +182,9 @@ pub async fn add_keys(
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_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!(
"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
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::post().to(add_keys))
})
.bind((args.ip.as_str(), args.port))?
.run()
.await
.bind((args.ip.as_str(), args.port))?
.run()
.await
}