Update README.MD

Cargo.lock

@@ -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"
@@ -1053,13 +1064,14 @@ dependencies = [
 
 [[package]]
 name = "khm"
-version = "0.4.0"
+version = "0.4.1"
 dependencies = [
  "actix-web",
  "base64 0.21.7",
  "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"

Cargo.toml

@@ -1,6 +1,6 @@
 [package]
 name = "khm"
-version = "0.4.0"
+version = "0.4.2"
 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"

README.MD

@@ -38,6 +38,7 @@ Options:
 - `--db-name <DB_NAME>` Server mode: Name of the PostgreSQL database [default: khm]
 - `--db-user <DB_USER>` Server mode: Username for the PostgreSQL database
 - `--db-password <DB_PASSWORD>` Server mode: Password for the PostgreSQL database
+- `--basic-auth <BASIC_AUTH>` Client mode: Basic Auth credentials [default: ""]
 - `--host <HOST>` Client mode: Full host address of the server to connect to. Like `https://khm.example.com/<FLOW_NAME>`
 - `--known-hosts <KNOWN_HOSTS>` Client mode: Path to the known_hosts file [default: ~/.ssh/known_hosts]
 
@@ -61,4 +62,4 @@ Contributions are welcome! Please open an issue or submit a pull request for any
 
 ## License
 
-This project is licensed under the WTFPL License.
+This project is licensed under the WTFPL License.

src/client.rs

@@ -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 {

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 total: usize,                   // Total number of received keys
-    pub inserted: usize,                // Количество новых ключей
+    pub inserted: usize,                // Number of new keys
-    pub updated: usize,                 // Количество обновленных ключей
+    pub unchanged: usize,               // Number of unchanged keys
-    pub unchanged: usize,               // Количество неизмененных ключей
+    pub key_id_map: Vec<(SshKey, i32)>, // Mapping of keys to their IDs in the database
-    pub key_id_map: Vec<(SshKey, i32)>, // Связь ключей с их ID в базе
 }
 
 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;

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,51 +35,6 @@ 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(())
-}
-
 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",
@@ -115,24 +70,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 +127,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 +164,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
+    // Always try to associate all keys with the flow, regardless of whether they're new or existing
-    if key_stats.inserted > 0 {
+    if !key_stats.key_id_map.is_empty() {
-        // Извлекаем только ID ключей из статистики
+        // Extract all key IDs from statistics, both new and existing
         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 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 +232,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 +241,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 +249,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")
     }
 }