switch to async

This commit is contained in:
2022-08-28 23:07:52 +02:00
parent 29c101fb1c
commit 96f5d2d781
2 changed files with 70 additions and 32 deletions

View File

@@ -6,7 +6,6 @@ edition = "2021"
# See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html # See more keys and their definitions at https://doc.rust-lang.org/cargo/reference/manifest.html
[dependencies] [dependencies]
dirs = "4.0.0"
serde = { version = "1.0", features = ["derive"] } serde = { version = "1.0", features = ["derive"] }
serde_yaml = "0.9.9" serde_yaml = "0.9.9"
ssh2 = "0.9" ssh2 = "0.9"

View File

@@ -1,30 +1,61 @@
mod config; mod config;
use std::{net::{TcpStream, ToSocketAddrs, SocketAddr}, path::{Path, PathBuf}, io::Write, time::Duration}; use std::io::Write;
use std::net::{TcpStream, SocketAddr, ToSocketAddrs};
use std::path::{PathBuf, Path};
use std::sync::{Arc, Mutex};
use std::thread;
use std::time::Duration;
use ssh2::Session; use ssh2::Session;
fn upload_file(host: &str, port: u16, username: &str, private_key: &Path, passphrase: Option<&str>, content: &str) -> Result<(), Box<dyn std::error::Error>> { #[derive(Clone)]
let addr: SocketAddr = format!("{}:{}", host, port).to_socket_addrs().unwrap().nth(0).expect(format!("Invalid host/port in `{}:{}`", host, port).as_str()); struct Host {
let tcp: TcpStream = TcpStream::connect_timeout(&addr, Duration::from_secs(1))?; host: Arc<Mutex<String>>,
let mut sess: Session = Session::new()?; port: u16,
username: Arc<Mutex<String>>,
content: Arc<Mutex<String>>
}
fn upload_config(h: &Host) -> Result<String, String> {
let host = h.host.lock().unwrap().to_owned();
let port = h.port;
let username = h.username.lock().unwrap().to_owned();
let private_key = PathBuf::from("/Users/toby/.ssh/id_ed25519");
let content = h.content.lock().unwrap();
let connection_string = format!("{}@{}", username, host);
let addr: SocketAddr = format!("{}:{}", host, port)
.to_socket_addrs()
.unwrap().nth(0)
.expect(format!("Invalid host/port in `{}:{}`", host, port).as_str());
let tcp: TcpStream = match TcpStream::connect_timeout(&addr, Duration::from_secs(1)) {
Ok(s) => s,
Err(_) => return Err(connection_string)
};
let mut sess = Session::new().unwrap();
sess.set_tcp_stream(tcp); sess.set_tcp_stream(tcp);
sess.handshake()?; sess.handshake().unwrap();
let public_key: PathBuf = private_key.with_extension("pub"); sess.userauth_pubkey_file(&username, None, &private_key, None).unwrap();
sess.userauth_pubkey_file(&username, Some(&public_key), &private_key, passphrase)?; let mut remote_file = match sess.scp_send(Path::new(".ssh").join("authorized_keys2").as_path(), 0o644, content.len() as u64, None) {
Ok(rf) => rf,
Err(_) => return Err(connection_string)
};
let mut remote_file = sess.scp_send(Path::new(".ssh").join("authorized_keys2").as_path(), 0o644, content.len() as u64, None)?; match remote_file.write(content.as_bytes()) {
Ok(rf) => rf,
Err(_) => return Err(connection_string)
};
remote_file.write(content.as_bytes())?; remote_file.send_eof().unwrap();
remote_file.wait_eof().unwrap();
remote_file.close().unwrap();
remote_file.wait_close().unwrap();
remote_file.send_eof()?; Ok(connection_string)
remote_file.wait_eof()?;
remote_file.close()?;
remote_file.wait_close()?;
Ok(())
} }
fn generate_authorized_keys(host_keys: Vec<String>) -> String { fn generate_authorized_keys(host_keys: Vec<String>) -> String {
@@ -32,9 +63,7 @@ fn generate_authorized_keys(host_keys: Vec<String>) -> String {
} }
fn main() { fn main() {
let ssh_dir: PathBuf = dirs::home_dir().unwrap().join(".ssh"); let mut hosts: Vec<Host> = vec![];
let private_key: PathBuf = ssh_dir.join("id_ed25519");
let passphrase: Option<&str> = None;
let config = config::read(); let config = config::read();
for host in &config.hosts { for host in &config.hosts {
@@ -73,19 +102,29 @@ fn main() {
continue; continue;
} }
let mut port: u16 = 22; let content = generate_authorized_keys(host_keys);
if host.port.is_some() { hosts.push(Host {
port = host.port.unwrap(); host: Arc::new(Mutex::new(host.host.to_string())),
port: host.port.unwrap_or(22),
username: Arc::new(Mutex::new(user_name.to_string())),
content: Arc::new(Mutex::new(content))
});
}
} }
let content = &generate_authorized_keys(host_keys); let mut handles = vec![];
// TODO: Multithreadding for host in hosts {
match upload_file(&host.host, port, user_name, &private_key.as_path(), passphrase, content) { handles.push(thread::spawn(move || {
Ok(_) => println!("{}@{}", user_name, host.host), match upload_config(&host) {
Err(_) => println!(" {}@{}", user_name, host.host) Ok(s) => println!(" {}", s),
} Err(s) => println!("{}", s)
} }
}));
}
for handler in handles {
handler.join().unwrap();
} }
} }