This commit is contained in:
AB
2024-07-07 21:02:39 +03:00
commit dff21a5f9c
8 changed files with 2892 additions and 0 deletions

1
.gitignore vendored Normal file
View File

@ -0,0 +1 @@
/target

5
.idea/.gitignore generated vendored Normal file
View File

@ -0,0 +1,5 @@
# Default ignored files
/shelf/
/workspace.xml
# Editor-based HTTP Client requests
/httpRequests/

2435
Cargo.lock generated Normal file

File diff suppressed because it is too large Load Diff

17
Cargo.toml Normal file
View File

@ -0,0 +1,17 @@
[package]
name = "untitled"
version = "0.1.0"
edition = "2021"
[dependencies]
actix-web = "4"
serde = { version = "1.0", features = ["derive"] }
serde_json = "1.0"
env_logger = "0.11.3"
log = "0.4"
regex = "1.10.5"
tokio = { version = "1", features = ["full"] }
tokio-postgres = { version = "0.7", features = ["with-chrono-0_4"] }
clap = { version = "4", features = ["derive"] }
chrono = "0.4.38"
reqwest = { version = "0.12", features = ["json"] }

62
README.MD Normal file
View File

@ -0,0 +1,62 @@
# KHM - Known Hosts Manager
KHM allows you to synchronize the `known_hosts` file across multiple hosts. 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. In client mode, it sends keys to the server and can update the `known_hosts` file with keys from the server.
## Features
- Synchronize `known_hosts` file across multiple hosts.
- Manage SSH keys and flows in a PostgreSQL database.
- Operate in both server and client modes.
- Automatically update `known_hosts` file with keys from the server.
## Usage
### Server Mode
To run the application in server mode, use the following command:
```bash
khm --server --ip 127.0.0.1 --port 8080 --db-host 127.0.0.1 --db-name khm --db-user admin --db-password <SECRET> --flows work,home
```
### Client Mode
To run the application in client mode, use the following command:
```bash
khm --host http://khm.example.com:8080 --known-hosts ~/.ssh/known_hosts --in-place
```
### Arguments
- `--server`: Run in server mode (default: false).
- `--ip`: IP address to bind the server or client to (default: 127.0.0.1).
- `--port`: Port to bind the server or client to (default: 8080).
- `--db-host`: Hostname or IP address of the PostgreSQL database (default: 127.0.0.1).
- `--db-name`: Name of the PostgreSQL database (default: khm).
- `--db-user`: Username for the PostgreSQL database (required in server mode).
- `--db-password`: Password for the PostgreSQL database (required in server mode).
- `--host`: Host address of the server to connect to in client mode (required in client mode).
- `--known-hosts`: Path to the `known_hosts` file (default: ~/.ssh/known_hosts).
- `--in-place`: Update the `known_hosts` file with keys from the server after sending keys (default: false).
- `--flows`: Comma-separated list of flows to manage (default: default).
## Installation
1. Ensure you have Rust installed. If not, you can install it from [rustup.rs](https://rustup.rs/).
2. Clone the repository:
```bash
git clone https://github.com/house-of-vanity/khm.git
cd khm
```
3. Run the project:
```bash
cargo run --release -- --help
```
## Contributing
Contributions are welcome! Please open an issue or submit a pull request for any changes.
## License
This project is licensed under the WTFPL License.

91
src/client.rs Normal file
View File

@ -0,0 +1,91 @@
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 {
server: String,
public_key: String,
}
fn read_known_hosts(file_path: &str) -> io::Result<Vec<SshKey>> {
let path = Path::new(file_path);
let file = File::open(&path)?;
let reader = io::BufReader::new(file);
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 });
}
}
}
Ok(keys)
}
fn write_known_hosts(file_path: &str, keys: &[SshKey]) -> io::Result<()> {
let path = Path::new(file_path);
let mut file = File::create(&path)?;
for key in keys {
writeln!(file, "{} {}", key.server, key.public_key)?;
}
Ok(())
}
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?;
if response.status().is_success() {
println!("Keys successfully sent to server.");
} else {
eprintln!("Failed to send keys to server. Status: {}", response.status());
}
Ok(())
}
async fn get_keys_from_server(host: &str) -> Result<Vec<SshKey>, reqwest::Error> {
let client = Client::new();
let url = format!("{}/keys", host);
let response = client.get(&url).send().await?;
if response.status().is_success() {
let keys: Vec<SshKey> = response.json().await?;
Ok(keys)
} else {
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 host = args.host.expect("host is required in client mode");
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
.expect("Failed to get keys from server");
write_known_hosts(&args.known_hosts, &server_keys)
.expect("Failed to write known hosts file");
}
Ok(())
}

83
src/main.rs Normal file
View File

@ -0,0 +1,83 @@
mod server;
mod client;
use clap::Parser;
use env_logger;
/// 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.
/// 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",
about = "SSH Key Manager",
long_about = None,
after_help = "Examples:\n\
\n\
Running in server mode:\n\
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\
\n\
"
)]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")]
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")]
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")]
db_host: String,
/// Name of the PostgreSQL database (default: khm)
#[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")]
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")]
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")]
host: Option<String>,
/// Run in server mode (default: false)
#[arg(long, help = "Run in server mode")]
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")]
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")]
in_place: bool,
/// Comma-separated list of flows to manage (default: default)
#[arg(long, default_value = "default", value_parser, num_args = 1.., value_delimiter = ',', help = "Comma-separated list of flows to manage")]
flows: Vec<String>,
}
#[actix_web::main]
async fn main() -> std::io::Result<()> {
env_logger::init();
let args = Args::parse();
if args.server {
server::run_server(args).await
} else {
client::run_client(args).await
}
}

198
src/server.rs Normal file
View File

@ -0,0 +1,198 @@
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;
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct SshKey {
pub server: String,
pub public_key: String,
}
#[derive(Serialize, Deserialize, Clone, Debug)]
pub struct Flow {
pub name: String,
pub servers: Vec<SshKey>,
}
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 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)
}
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?;
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?;
Ok(row.get(0))
}
}
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(())
}
pub async fn get_keys_from_db(client: &Client) -> Result<Vec<Flow>, tokio_postgres::Error> {
let rows = client.query(
"SELECT k.host, k.key, f.name FROM public.keys k INNER JOIN public.flows f ON k.key_id = f.key_id",
&[]
).await?;
let mut flows_map: HashMap<String, Flow> = HashMap::new();
for row in rows {
let host: String = row.get(0);
let key: String = row.get(1);
let flow: String = row.get(2);
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],
});
}
}
Ok(flows_map.into_values().collect())
}
pub async fn get_keys(
flows: web::Data<Flows>,
flow_id: web::Path<String>,
allowed_flows: web::Data<Vec<String>>
) -> impl Responder {
let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&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();
HttpResponse::Ok().json(servers)
} else {
HttpResponse::NotFound().body("Flow ID not found")
}
}
pub async fn add_keys(
flows: web::Data<Flows>,
flow_id: web::Path<String>,
new_keys: web::Json<Vec<SshKey>>,
db_client: web::Data<Arc<Client>>,
allowed_flows: web::Data<Vec<String>>
) -> impl Responder {
let flow_id_str = flow_id.into_inner();
if !allowed_flows.contains(&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) {
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");
}
},
Err(e) => {
log::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 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();
HttpResponse::Ok().json(servers)
} else {
HttpResponse::NotFound().body("Flow ID not found")
}
}
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_conn_str = format!(
"host={} user={} password={} dbname={}",
args.db_host, db_user, db_password, args.db_name
);
let (db_client, connection) = tokio_postgres::connect(&db_conn_str, NoTls).await.unwrap();
let db_client = Arc::new(db_client);
// Spawn a new thread to run the database connection
tokio::spawn(async move {
if let Err(e) = connection.await {
eprintln!("Connection error: {}", e);
}
});
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 {
if !initial_flows.iter().any(|flow| &flow.name == allowed_flow) {
initial_flows.push(Flow {
name: allowed_flow.clone(),
servers: vec![],
});
}
}
let flows: Flows = Arc::new(Mutex::new(initial_flows));
let allowed_flows = web::Data::new(args.flows);
HttpServer::new(move || {
App::new()
.app_data(web::Data::new(flows.clone()))
.app_data(web::Data::new(db_client.clone()))
.app_data(allowed_flows.clone())
.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
}