Files
khm/src/client.rs

157 lines
4.8 KiB
Rust
Raw Normal View History

2025-03-14 02:14:15 +02:00
use base64::{engine::general_purpose, Engine as _};
use log::{error, info};
2025-03-14 02:14:15 +02:00
use reqwest::header::{HeaderMap, HeaderValue, AUTHORIZATION};
2024-07-07 21:13:19 +03:00
use reqwest::Client;
use serde::{Deserialize, Serialize};
2024-07-07 21:02:39 +03:00
use std::fs::File;
use std::io::{self, BufRead, Write};
use std::path::Path;
#[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() {
2024-07-09 02:28:44 +03:00
match line {
Ok(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 });
}
}
2024-07-09 02:28:44 +03:00
Err(e) => {
error!("Error reading line from known_hosts file: {}", e);
2024-07-07 21:02:39 +03:00
}
}
}
2024-07-09 02:28:44 +03:00
info!("Read {} keys from known_hosts file", keys.len());
2024-07-07 21:02:39 +03:00
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)?;
}
2024-07-09 02:28:44 +03:00
info!("Wrote {} keys to known_hosts file", keys.len());
2024-07-07 21:02:39 +03:00
Ok(())
}
2025-03-14 02:14:15 +02:00
async fn send_keys_to_server(
host: &str,
keys: Vec<SshKey>,
auth_string: &str,
) -> Result<(), reqwest::Error> {
2024-07-07 21:02:39 +03:00
let client = Client::new();
let url = format!("{}/keys", host);
2025-03-14 02:14:15 +02:00
info!("URL: {} ", url);
let mut headers = HeaderMap::new();
if !auth_string.is_empty() {
let parts: Vec<&str> = auth_string.splitn(2, ':').collect();
if parts.len() == 2 {
let username = parts[0];
let password = parts[1];
let auth_header_value = format!("{}:{}", username, password);
let encoded_auth = general_purpose::STANDARD.encode(auth_header_value);
let auth_header = format!("Basic {}", encoded_auth);
headers.insert(AUTHORIZATION, HeaderValue::from_str(&auth_header).unwrap());
} else {
error!("Invalid auth string format. Expected 'username:password'");
}
}
let response = client
.post(&url)
.headers(headers)
.json(&keys)
.send()
.await?;
2024-07-07 21:02:39 +03:00
if response.status().is_success() {
2024-07-09 02:28:44 +03:00
info!("Keys successfully sent to server.");
2024-07-07 21:02:39 +03:00
} else {
2024-07-09 02:28:44 +03:00
error!(
2024-07-07 21:13:19 +03:00
"Failed to send keys to server. Status: {}",
response.status()
);
2024-07-07 21:02:39 +03:00
}
Ok(())
}
2025-03-14 02:14:15 +02:00
async fn get_keys_from_server(
host: &str,
auth_string: &str,
) -> Result<Vec<SshKey>, reqwest::Error> {
2024-07-07 21:02:39 +03:00
let client = Client::new();
let url = format!("{}/keys", host);
2025-03-14 02:14:15 +02:00
let mut headers = HeaderMap::new();
if !auth_string.is_empty() {
let parts: Vec<&str> = auth_string.splitn(2, ':').collect();
if parts.len() == 2 {
let username = parts[0];
let password = parts[1];
let auth_header_value = format!("{}:{}", username, password);
let encoded_auth = general_purpose::STANDARD.encode(auth_header_value);
let auth_header = format!("Basic {}", encoded_auth);
headers.insert(AUTHORIZATION, HeaderValue::from_str(&auth_header).unwrap());
} else {
error!("Invalid auth string format. Expected 'username:password'");
}
2024-07-07 21:02:39 +03:00
}
2025-03-14 02:14:15 +02:00
let response = client.get(&url).headers(headers).send().await?;
let response = response.error_for_status()?;
let keys: Vec<SshKey> = response.json().await?;
info!("Received {} keys from server", keys.len());
Ok(keys)
2024-07-07 21:02:39 +03:00
}
pub async fn run_client(args: crate::Args) -> std::io::Result<()> {
2024-07-09 02:28:44 +03:00
info!("Client mode: Reading known_hosts file");
2024-07-07 21:13:19 +03:00
let keys = read_known_hosts(&args.known_hosts).expect("Failed to read known hosts file");
2024-07-07 21:02:39 +03:00
let host = args.host.expect("host is required in client mode");
2024-07-09 02:28:44 +03:00
info!("Client mode: Sending keys to server at {}", host);
2025-03-14 02:14:15 +02:00
send_keys_to_server(&host, keys, &args.basic_auth)
2024-07-07 21:13:19 +03:00
.await
2024-07-07 21:02:39 +03:00
.expect("Failed to send keys to server");
if args.in_place {
2024-07-09 02:28:44 +03:00
info!("Client mode: In-place update is enabled. Fetching keys from server.");
2025-03-14 02:14:15 +02:00
let server_keys = get_keys_from_server(&host, &args.basic_auth)
2024-07-07 21:13:19 +03:00
.await
2024-07-07 21:02:39 +03:00
.expect("Failed to get keys from server");
2024-07-09 02:28:44 +03:00
info!("Client mode: Writing updated known_hosts file");
2024-07-07 21:02:39 +03:00
write_known_hosts(&args.known_hosts, &server_keys)
.expect("Failed to write known hosts file");
}
2024-07-09 02:28:44 +03:00
info!("Client mode: Finished operations");
2024-07-07 21:02:39 +03:00
Ok(())
}