tetanus/src/main.rs

use clap::Parser;
use rand::RngExt;
use rand::distr::Alphanumeric;
use std::sync::{Arc, Mutex};
use std::thread::{sleep, spawn};
use std::time::Duration;
use std::{env, path::PathBuf, process::exit};
use sysinfo::*;
use tetanus::funcs::*;
use tetanus::server::start_server;
use tetanus::{AppState, server};

mod install;

#[derive(Parser, Debug)]
#[command(
    author,
    version,
    about = "The Tetanus Redteaming tool. This will start the client, server, or both"
)]
struct Args {
    #[arg(short, long, help = "start in client mode")]
    client: bool,

    #[arg(short, long, help = "start in server mode")]
    server: bool,

    #[arg(
        short,
        long,
        help = "The IP address for the server, in server mode it will listen on this IP, in client mode it will add this ip as its first server. Default is 127.0.0.1"
    )]
    ip: Option<String>,

    #[arg(
        short,
        long,
        help = "The port for the server, in server mode this will be the port the server listens on, in client mode this will be the default server port it connects to."
    )]
    port: Option<String>,

    #[arg(
        short = 'C',
        long,
        help = "Custom Config file. Give the path to a config file to use."
    )]
    config: Option<PathBuf>,

    #[arg(
        short = 'P',
        long,
        help = "the password for the server, in server mode it will set the password clients need to log in, in client mode it sets the password the client will use to login to the server. If blank it will generate a random password."
    )]
    password: Option<String>,

    #[arg(
        short,
        long,
        help = "The name to set this instance, on servers it will name the server so clients know who it is, on clients it names the clinet for server communications. This defaults to the hostname of the machine, or if the hostname can't be determined it falls back to the username of who runs it."
    )]
    name: Option<String>,
}

#[tokio::main]
async fn main() {
    println!("checking for server or client config files...");
    let args = Args::parse();
    let mut config_path = PathBuf::new();
    let mut config_path_opt = env::home_dir();
    if let Some(given_config_path) = args.config {
        config_path = given_config_path.clone();
    } else if let Some(home_config_path) = config_path_opt.as_mut() {
        println!("made it to home_config_path");
        home_config_path.push(".config/tetanus");
        if !home_config_path.exists() {
            println!("no config directory found in home directory...");
            config_path = PathBuf::from("/etc/tetanus");
            if !config_path.exists() {
                println!("no config directory found in /etc/tetanus... Installing!");
                let install_result = install::install();
                if install_result.is_ok() {
                    install_result.unwrap();
                    println!("install succeeded!");
                    exit(0);
                }
            }
        } else {
            println!("found home config_path.");
            config_path = home_config_path.clone();
        }
    }
    let mut client_config_path = config_path.clone();
    client_config_path.push("client.conf");
    let mut name = String::new();
    let mut server_hostname = String::new();
    if let Some(aname) = args.name {
        name = aname;
    } else {
        if let Some(hostname) = System::host_name() {
            name = hostname.clone();
            server_hostname = hostname;
        } else {
            for (key, value) in env::vars() {
                if key == "USERNAME" || key == "USER" {
                    name = value;
                }
            }
        }
    }
    if args.client {
        let banner = "
        XXXXXXXXXXXXXXXXXXXXXXXXXX
         XXXXXXXXXXXXXXXXXXXXXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                  XXXXX
                   XXX
                    X

              ++----------++
              ++----------++
              ||Welcome To||
              ||  TETANUS ||
              ++----------++
              ++----------++      ";
        if !client_config_path.exists() {
            eprintln!(
                "error: no client config path found at {}",
                client_config_path.display()
            );
            exit(1);
        }
        let mut appstate = AppState::new();
        let config_load = appstate.load_config(client_config_path.clone(), true);
        if config_load.is_err() {
            eprintln!("error loading config!");
            exit(1);
        }
        appstate.name = name;
        banner.lines().for_each(|line| {
            appstate.output.push(line.to_string());
        });
        let handles = Arc::new(Mutex::new(Vec::new()));
        let handles_clone = handles.clone();
        let primary_handle = spawn(move || {
            println!("entering startup function...");
            let _ = startup(appstate, false, handles_clone);
        });
        if let Ok(mut lock) = handles.lock() {
            lock.push(primary_handle);
        }
        loop {
            sleep(Duration::from_secs(3));
            if let Ok(lock) = handles.lock() {
                if lock.iter().all(|handle| handle.is_finished()) {
                    break;
                }
            }
        }
    } else if args.server {
        let mut address = "127.0.0.1:31337".to_string();
        let mut certificate_path = config_path.clone();
        certificate_path.push("server");
        let mut key_path = certificate_path.clone();
        certificate_path.push("cert.pem");
        key_path.push("key.pem");
        println!("key paht: {}", key_path.display());
        println!("cert path: {}", certificate_path.display());
        if let Some(ip) = args.ip {
            if let Some(port) = args.port {
                address = format!("{}:{}", ip.trim(), port.trim());
            } else {
                address = format!("{}:31337", ip.trim())
            }
        }
        let mut server_pass: String = rand::rng()
            .sample_iter(&Alphanumeric)
            .take(12)
            .map(char::from)
            .collect();
        if let Some(pass) = args.password {
            server_pass = pass;
        }
        let new_server = server::Server {
            address,
            clients: Vec::new(),
            certificate_path,
            key_path,
            password: server_pass,
            name,
            hostname: server_hostname,
        };
        let res = start_server(Arc::new(Mutex::new(new_server))).await;
        match res {
            Ok(_) => {}
            Err(e) => {
                println!("error running server {e}");
            }
        }
    }
}