mirror of
https://github.com/house-of-vanity/khm.git
synced 2025-07-07 15:24:07 +00:00
Compare commits
2 Commits
| Author | SHA1 | Date | |
|---|---|---|---|
| 4b2b56bcd2 | |||
| 2cfc2c6c3a |
40
Cargo.lock
generated
40
Cargo.lock
generated
@ -836,6 +836,17 @@ dependencies = [
|
||||
"digest",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "hostname"
|
||||
version = "0.3.1"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "3c731c3e10504cc8ed35cfe2f1db4c9274c3d35fa486e3b31df46f068ef3e867"
|
||||
dependencies = [
|
||||
"libc",
|
||||
"match_cfg",
|
||||
"winapi",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "http"
|
||||
version = "0.2.12"
|
||||
@ -1060,6 +1071,7 @@ dependencies = [
|
||||
"chrono",
|
||||
"clap",
|
||||
"env_logger",
|
||||
"hostname",
|
||||
"log",
|
||||
"regex",
|
||||
"reqwest",
|
||||
@ -1120,6 +1132,12 @@ version = "0.4.22"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "a7a70ba024b9dc04c27ea2f0c0548feb474ec5c54bba33a7f72f873a39d07b24"
|
||||
|
||||
[[package]]
|
||||
name = "match_cfg"
|
||||
version = "0.1.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ffbee8634e0d45d258acb448e7eaab3fce7a0a467395d4d9f228e3c1f01fb2e4"
|
||||
|
||||
[[package]]
|
||||
name = "md-5"
|
||||
version = "0.10.6"
|
||||
@ -2223,6 +2241,28 @@ dependencies = [
|
||||
"web-sys",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winapi"
|
||||
version = "0.3.9"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "5c839a674fcd7a98952e593242ea400abe93992746761e38641405d28b00f419"
|
||||
dependencies = [
|
||||
"winapi-i686-pc-windows-gnu",
|
||||
"winapi-x86_64-pc-windows-gnu",
|
||||
]
|
||||
|
||||
[[package]]
|
||||
name = "winapi-i686-pc-windows-gnu"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "ac3b87c63620426dd9b991e5ce0329eff545bccbbb34f3be09ff6fb6ab51b7b6"
|
||||
|
||||
[[package]]
|
||||
name = "winapi-x86_64-pc-windows-gnu"
|
||||
version = "0.4.0"
|
||||
source = "registry+https://github.com/rust-lang/crates.io-index"
|
||||
checksum = "712e227841d057c1ee1cd2fb22fa7e5a5461ae8e48fa2ca79ec42cfc1931183f"
|
||||
|
||||
[[package]]
|
||||
name = "windows-core"
|
||||
version = "0.52.0"
|
||||
|
||||
@ -1,6 +1,6 @@
|
||||
[package]
|
||||
name = "khm"
|
||||
version = "0.4.0"
|
||||
version = "0.4.1"
|
||||
edition = "2021"
|
||||
authors = ["AB <ab@hexor.cy>"]
|
||||
|
||||
@ -17,3 +17,4 @@ 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"] }
|
||||
hostname = "0.3"
|
||||
|
||||
@ -50,6 +50,14 @@ fn write_known_hosts(file_path: &str, keys: &[SshKey]) -> io::Result<()> {
|
||||
Ok(())
|
||||
}
|
||||
|
||||
// Get local hostname for request headers
|
||||
fn get_hostname() -> String {
|
||||
match hostname::get() {
|
||||
Ok(name) => name.to_string_lossy().to_string(),
|
||||
Err(_) => "unknown-host".to_string(),
|
||||
}
|
||||
}
|
||||
|
||||
async fn send_keys_to_server(
|
||||
host: &str,
|
||||
keys: Vec<SshKey>,
|
||||
@ -61,6 +69,17 @@ async fn send_keys_to_server(
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
// Add hostname header
|
||||
let hostname = get_hostname();
|
||||
headers.insert(
|
||||
"X-Client-Hostname",
|
||||
HeaderValue::from_str(&hostname).unwrap_or_else(|_| {
|
||||
error!("Failed to create hostname header value");
|
||||
HeaderValue::from_static("unknown-host")
|
||||
}),
|
||||
);
|
||||
info!("Adding hostname header: {}", hostname);
|
||||
|
||||
if !auth_string.is_empty() {
|
||||
let parts: Vec<&str> = auth_string.splitn(2, ':').collect();
|
||||
if parts.len() == 2 {
|
||||
@ -105,6 +124,17 @@ async fn get_keys_from_server(
|
||||
|
||||
let mut headers = HeaderMap::new();
|
||||
|
||||
// Add hostname header
|
||||
let hostname = get_hostname();
|
||||
headers.insert(
|
||||
"X-Client-Hostname",
|
||||
HeaderValue::from_str(&hostname).unwrap_or_else(|_| {
|
||||
error!("Failed to create hostname header value");
|
||||
HeaderValue::from_static("unknown-host")
|
||||
}),
|
||||
);
|
||||
info!("Adding hostname header: {}", hostname);
|
||||
|
||||
if !auth_string.is_empty() {
|
||||
let parts: Vec<&str> = auth_string.splitn(2, ':').collect();
|
||||
if parts.len() == 2 {
|
||||
|
||||
35
src/db.rs
35
src/db.rs
@ -4,13 +4,12 @@ use std::collections::HashMap;
|
||||
use std::collections::HashSet;
|
||||
use tokio_postgres::Client;
|
||||
|
||||
// Структура для хранения статистики обработки ключей
|
||||
// Structure for storing key processing statistics
|
||||
pub struct KeyInsertStats {
|
||||
pub total: usize, // Общее количество полученных ключей
|
||||
pub inserted: usize, // Количество новых ключей
|
||||
pub updated: usize, // Количество обновленных ключей
|
||||
pub unchanged: usize, // Количество неизмененных ключей
|
||||
pub key_id_map: Vec<(SshKey, i32)>, // Связь ключей с их ID в базе
|
||||
pub total: usize, // Total number of received keys
|
||||
pub inserted: usize, // Number of new keys
|
||||
pub unchanged: usize, // Number of unchanged keys
|
||||
pub key_id_map: Vec<(SshKey, i32)>, // Mapping of keys to their IDs in the database
|
||||
}
|
||||
|
||||
pub async fn initialize_db_schema(client: &Client) -> Result<(), tokio_postgres::Error> {
|
||||
@ -93,13 +92,12 @@ pub async fn batch_insert_keys(
|
||||
return Ok(KeyInsertStats {
|
||||
total: 0,
|
||||
inserted: 0,
|
||||
updated: 0,
|
||||
unchanged: 0,
|
||||
key_id_map: Vec::new(),
|
||||
});
|
||||
}
|
||||
|
||||
// Подготавливаем массивы для пакетной вставки
|
||||
// Prepare arrays for batch insertion
|
||||
let mut host_values: Vec<&str> = Vec::with_capacity(keys.len());
|
||||
let mut key_values: Vec<&str> = Vec::with_capacity(keys.len());
|
||||
|
||||
@ -108,7 +106,7 @@ pub async fn batch_insert_keys(
|
||||
key_values.push(&key.public_key);
|
||||
}
|
||||
|
||||
// Сначала проверяем, какие ключи уже существуют в базе
|
||||
// First, check which keys already exist in the database
|
||||
let mut existing_keys = HashMap::new();
|
||||
let mut key_query = String::from("SELECT host, key, key_id FROM public.keys WHERE ");
|
||||
|
||||
@ -135,7 +133,7 @@ pub async fn batch_insert_keys(
|
||||
existing_keys.insert((host, key), key_id);
|
||||
}
|
||||
|
||||
// Определяем, какие ключи нужно вставить, а какие уже существуют
|
||||
// Determine which keys need to be inserted and which already exist
|
||||
let mut keys_to_insert = Vec::new();
|
||||
let mut unchanged_keys = Vec::new();
|
||||
|
||||
@ -150,7 +148,7 @@ pub async fn batch_insert_keys(
|
||||
|
||||
let mut inserted_keys = Vec::new();
|
||||
|
||||
// Если есть ключи для вставки, выполняем вставку
|
||||
// If there are keys to insert, perform the insertion
|
||||
if !keys_to_insert.is_empty() {
|
||||
let mut insert_sql = String::from("INSERT INTO public.keys (host, key, updated) VALUES ");
|
||||
|
||||
@ -185,11 +183,11 @@ pub async fn batch_insert_keys(
|
||||
}
|
||||
}
|
||||
|
||||
// Сохраняем количество элементов перед объединением
|
||||
// Save the number of elements before combining
|
||||
let inserted_count = inserted_keys.len();
|
||||
let unchanged_count = unchanged_keys.len();
|
||||
|
||||
// Комбинируем результаты и формируем статистику
|
||||
// Combine results and generate statistics
|
||||
let mut key_id_map = Vec::with_capacity(unchanged_count + inserted_count);
|
||||
key_id_map.extend(unchanged_keys);
|
||||
key_id_map.extend(inserted_keys);
|
||||
@ -197,7 +195,6 @@ pub async fn batch_insert_keys(
|
||||
let stats = KeyInsertStats {
|
||||
total: keys.len(),
|
||||
inserted: inserted_count,
|
||||
updated: 0, // В этой версии мы не обновляем существующие ключи
|
||||
unchanged: unchanged_count,
|
||||
key_id_map,
|
||||
};
|
||||
@ -220,7 +217,7 @@ pub async fn batch_insert_flow_keys(
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Сначала проверим, какие связи уже существуют
|
||||
// First, check which associations already exist
|
||||
let mut existing_query =
|
||||
String::from("SELECT key_id FROM public.flows WHERE name = $1 AND key_id IN (");
|
||||
|
||||
@ -247,7 +244,7 @@ pub async fn batch_insert_flow_keys(
|
||||
existing_associations.insert(key_id);
|
||||
}
|
||||
|
||||
// Фильтруем только те ключи, которые еще не связаны с потоком
|
||||
// Filter only keys that are not yet associated with the flow
|
||||
let new_key_ids: Vec<&i32> = key_ids
|
||||
.iter()
|
||||
.filter(|&id| !existing_associations.contains(id))
|
||||
@ -262,7 +259,7 @@ pub async fn batch_insert_flow_keys(
|
||||
return Ok(0);
|
||||
}
|
||||
|
||||
// Строим SQL запрос с множественными значениями только для новых связей
|
||||
// Build SQL query with multiple values only for new associations
|
||||
let mut sql = String::from("INSERT INTO public.flows (name, key_id) VALUES ");
|
||||
|
||||
for i in 0..new_key_ids.len() {
|
||||
@ -274,7 +271,7 @@ pub async fn batch_insert_flow_keys(
|
||||
|
||||
sql.push_str(" ON CONFLICT (name, key_id) DO NOTHING");
|
||||
|
||||
// Подготавливаем параметры для запроса
|
||||
// Prepare parameters for the query
|
||||
let mut insert_params: Vec<&(dyn tokio_postgres::types::ToSql + Sync)> =
|
||||
Vec::with_capacity(new_key_ids.len() + 1);
|
||||
insert_params.push(&flow_name);
|
||||
@ -282,7 +279,7 @@ pub async fn batch_insert_flow_keys(
|
||||
insert_params.push(*key_id);
|
||||
}
|
||||
|
||||
// Выполняем запрос
|
||||
// Execute query
|
||||
let affected = client.execute(&sql, &insert_params[..]).await?;
|
||||
|
||||
let affected_usize = affected as usize;
|
||||
|
||||
150
src/server.rs
150
src/server.rs
@ -1,4 +1,4 @@
|
||||
use actix_web::{web, App, HttpResponse, HttpServer, Responder};
|
||||
use actix_web::{web, App, HttpRequest, HttpResponse, HttpServer, Responder};
|
||||
use log::{error, info};
|
||||
use regex::Regex;
|
||||
use serde::{Deserialize, Serialize};
|
||||
@ -35,50 +35,7 @@ pub fn is_valid_ssh_key(key: &str) -> bool {
|
||||
|| 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?;
|
||||
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))
|
||||
}
|
||||
}
|
||||
|
||||
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?;
|
||||
info!("Inserted key_id {} into flow: {}", key_id, flow_name);
|
||||
Ok(())
|
||||
}
|
||||
// Note: Removed unused functions insert_key_if_not_exists and insert_flow_key
|
||||
|
||||
pub async fn get_keys_from_db(client: &Client) -> Result<Vec<Flow>, tokio_postgres::Error> {
|
||||
let rows = client.query(
|
||||
@ -115,24 +72,53 @@ pub async fn get_keys_from_db(client: &Client) -> Result<Vec<Flow>, tokio_postgr
|
||||
Ok(flows_map.into_values().collect())
|
||||
}
|
||||
|
||||
// Extract client hostname from request headers
|
||||
fn get_client_hostname(req: &HttpRequest) -> String {
|
||||
if let Some(hostname) = req.headers().get("X-Client-Hostname") {
|
||||
if let Ok(hostname_str) = hostname.to_str() {
|
||||
return hostname_str.to_string();
|
||||
}
|
||||
}
|
||||
"unknown-client".to_string()
|
||||
}
|
||||
|
||||
pub async fn get_keys(
|
||||
flows: web::Data<Flows>,
|
||||
flow_id: web::Path<String>,
|
||||
allowed_flows: web::Data<Vec<String>>,
|
||||
req: HttpRequest,
|
||||
) -> impl Responder {
|
||||
let client_hostname = get_client_hostname(&req);
|
||||
let flow_id_str = flow_id.into_inner();
|
||||
|
||||
info!(
|
||||
"Received keys request from client '{}' for flow '{}'",
|
||||
client_hostname, flow_id_str
|
||||
);
|
||||
|
||||
if !allowed_flows.contains(&flow_id_str) {
|
||||
error!("Flow ID not allowed: {}", flow_id_str);
|
||||
error!(
|
||||
"Flow ID not allowed for client '{}': {}",
|
||||
client_hostname, 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);
|
||||
info!(
|
||||
"Returning {} keys for flow '{}' to client '{}'",
|
||||
servers.len(),
|
||||
flow_id_str,
|
||||
client_hostname
|
||||
);
|
||||
HttpResponse::Ok().json(servers)
|
||||
} else {
|
||||
error!("Flow ID not found: {}", flow_id_str);
|
||||
error!(
|
||||
"Flow ID not found for client '{}': {}",
|
||||
client_hostname, flow_id_str
|
||||
);
|
||||
HttpResponse::NotFound().body("Flow ID not found")
|
||||
}
|
||||
}
|
||||
@ -143,18 +129,34 @@ pub async fn add_keys(
|
||||
new_keys: web::Json<Vec<SshKey>>,
|
||||
db_client: web::Data<Arc<Client>>,
|
||||
allowed_flows: web::Data<Vec<String>>,
|
||||
req: HttpRequest,
|
||||
) -> impl Responder {
|
||||
let client_hostname = get_client_hostname(&req);
|
||||
let flow_id_str = flow_id.into_inner();
|
||||
|
||||
info!(
|
||||
"Received {} keys from client '{}' for flow '{}'",
|
||||
new_keys.len(),
|
||||
client_hostname,
|
||||
flow_id_str
|
||||
);
|
||||
|
||||
if !allowed_flows.contains(&flow_id_str) {
|
||||
error!("Flow ID not allowed: {}", flow_id_str);
|
||||
error!(
|
||||
"Flow ID not allowed for client '{}': {}",
|
||||
client_hostname, flow_id_str
|
||||
);
|
||||
return HttpResponse::Forbidden().body("Flow ID not allowed");
|
||||
}
|
||||
|
||||
// Проверяем формат SSH ключей
|
||||
// Check SSH key format
|
||||
let mut valid_keys = Vec::new();
|
||||
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);
|
||||
error!(
|
||||
"Invalid SSH key format from client '{}' for server: {}",
|
||||
client_hostname, new_key.server
|
||||
);
|
||||
return HttpResponse::BadRequest().body(format!(
|
||||
"Invalid SSH key format for server: {}",
|
||||
new_key.server
|
||||
@ -164,48 +166,62 @@ pub async fn add_keys(
|
||||
}
|
||||
|
||||
info!(
|
||||
"Processing batch of {} keys for flow: {}",
|
||||
"Processing batch of {} keys from client '{}' for flow: {}",
|
||||
valid_keys.len(),
|
||||
client_hostname,
|
||||
flow_id_str
|
||||
);
|
||||
|
||||
// Батчевая вставка ключей с получением статистики
|
||||
// Batch insert keys with statistics
|
||||
let key_stats = match crate::db::batch_insert_keys(&db_client, &valid_keys).await {
|
||||
Ok(stats) => stats,
|
||||
Err(e) => {
|
||||
error!("Failed to batch insert keys into database: {}", e);
|
||||
error!(
|
||||
"Failed to batch insert keys from client '{}' into database: {}",
|
||||
client_hostname, e
|
||||
);
|
||||
return HttpResponse::InternalServerError()
|
||||
.body("Failed to batch insert keys into database");
|
||||
}
|
||||
};
|
||||
|
||||
// Если нет новых ключей, нет необходимости обновлять связи с flow
|
||||
// If there are no new keys, no need to update flow associations
|
||||
if key_stats.inserted > 0 {
|
||||
// Извлекаем только ID ключей из статистики
|
||||
// Extract only key IDs from statistics
|
||||
let key_ids: Vec<i32> = key_stats.key_id_map.iter().map(|(_, id)| *id).collect();
|
||||
|
||||
// Батчевая вставка связей ключей с flow
|
||||
// Batch insert key-flow associations
|
||||
if let Err(e) = crate::db::batch_insert_flow_keys(&db_client, &flow_id_str, &key_ids).await
|
||||
{
|
||||
error!("Failed to batch insert flow keys into database: {}", e);
|
||||
error!(
|
||||
"Failed to batch insert flow keys from client '{}' into database: {}",
|
||||
client_hostname, e
|
||||
);
|
||||
return HttpResponse::InternalServerError()
|
||||
.body("Failed to batch insert flow keys into database");
|
||||
}
|
||||
|
||||
info!(
|
||||
"Added flow associations for {} keys in flow '{}'",
|
||||
"Added flow associations for {} keys from client '{}' in flow '{}'",
|
||||
key_ids.len(),
|
||||
client_hostname,
|
||||
flow_id_str
|
||||
);
|
||||
} else {
|
||||
info!("No new keys to associate with flow '{}'", flow_id_str);
|
||||
info!(
|
||||
"No new keys to associate from client '{}' with flow '{}'",
|
||||
client_hostname, flow_id_str
|
||||
);
|
||||
}
|
||||
|
||||
// Получаем обновленные данные
|
||||
// Get updated data
|
||||
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);
|
||||
error!(
|
||||
"Failed to get updated flows from database after client '{}' request: {}",
|
||||
client_hostname, e
|
||||
);
|
||||
return HttpResponse::InternalServerError()
|
||||
.body("Failed to refresh flows from database");
|
||||
}
|
||||
@ -218,7 +234,8 @@ pub async fn add_keys(
|
||||
if let Some(flow) = updated_flow {
|
||||
let servers: Vec<&SshKey> = flow.servers.iter().collect();
|
||||
info!(
|
||||
"Keys summary for flow '{}': total received={}, new={}, unchanged={}, total in flow={}",
|
||||
"Keys summary for client '{}', flow '{}': total received={}, new={}, unchanged={}, total in flow={}",
|
||||
client_hostname,
|
||||
flow_id_str,
|
||||
key_stats.total,
|
||||
key_stats.inserted,
|
||||
@ -226,7 +243,7 @@ pub async fn add_keys(
|
||||
servers.len()
|
||||
);
|
||||
|
||||
// Добавляем статистику в HTTP заголовки ответа
|
||||
// Add statistics to HTTP response headers
|
||||
let mut response = HttpResponse::Ok();
|
||||
response.append_header(("X-Keys-Total", key_stats.total.to_string()));
|
||||
response.append_header(("X-Keys-New", key_stats.inserted.to_string()));
|
||||
@ -234,7 +251,10 @@ pub async fn add_keys(
|
||||
|
||||
response.json(servers)
|
||||
} else {
|
||||
error!("Flow ID not found after update: {}", flow_id_str);
|
||||
error!(
|
||||
"Flow ID not found after update from client '{}': {}",
|
||||
client_hostname, flow_id_str
|
||||
);
|
||||
HttpResponse::NotFound().body("Flow ID not found")
|
||||
}
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user