Hexor/khm
1
0
Fork 0
mirror of https://github.com/house-of-vanity/khm.git synced 2025-08-21 22: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

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
}