pub mod crypto;
pub mod db;
pub mod http;
pub mod utils;
#[cfg(feature = "vt")]
pub mod vt;

use crate::crypto::FileEncryption;
use crate::db::MDBConfig;
use malwaredb_api::ServerInfo;
//use utils::HashPath;

use std::collections::HashMap;
use std::io::{Cursor, Read, Write};
use std::net::{IpAddr, SocketAddr};
use std::path::PathBuf;
use std::str::FromStr;
use std::sync::Arc;
use std::time::{Duration, SystemTime};

use anyhow::{bail, Result};
use chrono::Local;
use chrono_humanize::{Accuracy, HumanTime, Tense};
use flate2::read::GzDecoder;
use flate2::write::GzEncoder;
use flate2::Compression;
use sha2::{Digest, Sha256};
use tokio::net::TcpListener;
#[cfg(feature = "vt")]
use zeroize::Zeroizing;

/// MDB version
pub const MDB_VERSION: &str = env!("CARGO_PKG_VERSION");

/// GZip magic number to see if a file is compressed
pub const GZIP_MAGIC: [u8; 2] = [0x1fu8, 0x8bu8];

/// State & configuration of the running server instance
pub struct State {
    /// The port which will be used to listen for connections.
    pub port: u16,

    /// The directory to store malware samples, if we're keeping them.
    pub directory: Option<PathBuf>,

    /// Maximum upload size
    pub max_upload: usize,

    /// The IP to use for listening for connections
    pub ip: IpAddr,

    /// Handle to the database connection
    pub db_type: db::DatabaseType,

    /// Start time of the server
    pub started: SystemTime,

    /// Configuration which is stored in the database
    pub db_config: MDBConfig,

    /// File encryption keys, may be empty
    pub(crate) keys: HashMap<u32, FileEncryption>,

    /// VirusTotal API key
    #[cfg(feature = "vt")]
    vt_api_key: Option<Zeroizing<String>>,
}

impl State {
    pub async fn new(
        port: u16,
        directory: Option<PathBuf>,
        max_upload: usize,
        ip: IpAddr,
        db_string: &str,
        #[cfg(feature = "vt")] vt_api_key: Option<Zeroizing<String>>,
    ) -> Result<Self> {
        if let Some(dir) = &directory {
            if !dir.exists() {
                bail!("data directory {dir:?} does not exist!");
            }
        }

        let db_type = db::DatabaseType::from_string(db_string).await?;
        // TODO: allow config and keys to be refreshed so changes don't require a restart
        let db_config = db_type.get_config().await?;
        let keys = db_type.get_encryption_keys().await?;

        Ok(Self {
            port,
            directory,
            max_upload,
            ip,
            db_type,
            db_config,
            keys,
            #[cfg(feature = "vt")]
            vt_api_key,
            started: SystemTime::now(),
        })
    }

    #[allow(clippy::too_many_arguments)]
    pub async fn new_first_run(
        port: u16,
        directory: Option<PathBuf>,
        max_upload: usize,
        ip: IpAddr,
        db_string: &str,
        #[cfg(feature = "vt")] vt_api_key: Option<Zeroizing<String>>,
        compress: bool,
        #[cfg(feature = "vt")] send_samples_to_vt: bool,
    ) -> Result<Self> {
        if let Some(dir) = &directory {
            if !dir.exists() {
                bail!("data directory {dir:?} does not exist!");
            }
        }

        let db_type = db::DatabaseType::from_string(db_string).await?;

        if db_type.first_run() {
            println!("Welcome to MalwareDB {}!", MDB_VERSION);
            if compress {
                db_type.enable_compression().await?;
            }
            #[cfg(feature = "vt")]
            if send_samples_to_vt {
                db_type.enable_vt_upload().await?;
            }
        } else {
            bail!("Attempting to set first-run options when not first-run");
        }

        let db_config = db_type.get_config().await?;
        let keys = db_type.get_encryption_keys().await?;

        Ok(Self {
            port,
            directory,
            max_upload,
            ip,
            db_type,
            db_config,
            keys,
            #[cfg(feature = "vt")]
            vt_api_key,
            started: SystemTime::now(),
        })
    }

    /// Store the sample with a depth of three based on the sample's SHA-256 hash, even if compressed
    pub async fn store_bytes(&self, data: &[u8]) -> Result<bool> {
        if let Some(dest_path) = &self.directory {
            let mut hasher = Sha256::new();
            hasher.update(data);
            let sha256 = hex::encode(hasher.finalize());

            // Trait `HashPath` needs to be re-worked so it can work with Strings.
            // This code below ends up making the String into ASCII representations of the hash
            // See: https://github.com/malwaredb/malwaredb-rs/issues/60
            let hashed_path = format!(
                "{}/{}/{}/{}",
                &sha256[0..2],
                &sha256[2..4],
                &sha256[4..6],
                sha256
            );

            // The path which has the file name included, with the storage directory prepended.
            //let hashed_path = result.hashed_path(3);
            let mut dest_path = dest_path.clone();
            dest_path.push(hashed_path);

            // Remove the file name so we can just have the directory path.
            let mut just_the_dir = dest_path.clone();
            just_the_dir.pop();
            std::fs::create_dir_all(just_the_dir)?;

            let data = if self.db_config.compression {
                let mut compressor = GzEncoder::new(Vec::new(), Compression::default());
                compressor.write_all(data)?;
                compressor.finish()?
            } else {
                data.to_vec()
            };

            let data = if let Some(key_id) = self.db_config.default_key {
                if let Some(key) = self.keys.get(&key_id) {
                    let nonce = key.nonce();
                    self.db_type.set_file_nonce(&sha256, &nonce).await?;
                    key.encrypt(&data, nonce)?
                } else {
                    bail!("Key not available!")
                }
            } else {
                data
            };

            std::fs::write(dest_path, data)?;

            Ok(true)
        } else {
            Ok(false)
        }
    }

    /// Retrieve a sample given the SHA-256 hash
    /// Assumes that MalwareDB permissions have already been checked to ensure this is permitted.
    pub async fn retrieve_bytes(&self, sha256: &String) -> Result<Vec<u8>> {
        if let Some(dest_path) = &self.directory {
            let path = format!(
                "{}/{}/{}/{}",
                &sha256[0..2],
                &sha256[2..4],
                &sha256[4..6],
                sha256
            );
            // Trait `HashPath` needs to be re-worked so it can work with Strings.
            // This code below ends up making the String into ASCII representations of the hash
            // See: https://github.com/malwaredb/malwaredb-rs/issues/60
            //let path = sha256.as_bytes().iter().hashed_path(3);
            let contents = std::fs::read(dest_path.join(path))?;

            let contents = if !self.keys.is_empty() {
                let (key_id, nonce) = self.db_type.get_file_encryption_key_id(sha256).await?;
                if let Some(key_id) = key_id {
                    if let Some(key) = self.keys.get(&key_id) {
                        key.decrypt(&contents, nonce)?
                    } else {
                        bail!("File was encrypted but we don't have tke key!")
                    }
                } else {
                    // File was not encrypted
                    contents
                }
            } else {
                // We don't have encryption keys configured
                contents
            };

            if contents.starts_with(&GZIP_MAGIC) {
                let buff = Cursor::new(contents);
                let mut decompressor = GzDecoder::new(buff);
                let mut decompressed: Vec<u8> = vec![];
                decompressor.read_to_end(&mut decompressed)?;
                Ok(decompressed)
            } else {
                Ok(contents)
            }
        } else {
            bail!("files are not saved")
        }
    }

    pub fn since(&self) -> Duration {
        let now = SystemTime::now();
        now.duration_since(self.started).unwrap()
    }

    pub async fn get_info(&self) -> Result<ServerInfo> {
        let db_info = self.db_type.db_info().await?;

        let os_name = if cfg!(target_os = "linux") {
            "Linux"
        } else if cfg!(target_os = "macos") {
            "macOS"
        } else if cfg!(target_os = "windows") {
            "Windows"
        } else if cfg!(target_os = "freebsd") {
            "FreeBSD"
        } else if cfg!(target_os = "openbsd") {
            "OpenBSD"
        } else if cfg!(target_os = "netbsd") {
            "NetBSD"
        } else if cfg!(target_os = "wasi") {
            "WebAssembly WASI"
        } else {
            "Unknown"
        };

        let mem_size = if cfg!(target_family = "unix") {
            if let Ok(statm) = std::fs::read_to_string("/proc/self/statm") {
                let mut parts = statm.split(' ');
                if let Some(total_memory) = parts.next() {
                    u64::from_str(total_memory)
                        .map(|mem| {
                            humansize::SizeFormatter::new(mem, humansize::BINARY).to_string()
                        })
                        .unwrap_or_default()
                } else {
                    "".into()
                }
            } else {
                "".into()
            }
        } else {
            "".into()
        };

        let uptime = Local::now() - self.since();

        Ok(ServerInfo {
            os_name: os_name.into(),
            memory_used: mem_size,
            num_samples: db_info.num_files,
            num_users: db_info.num_users,
            uptime: HumanTime::from(uptime).to_text_en(Accuracy::Rough, Tense::Present),
            mdb_version: MDB_VERSION.into(),
            db_version: db_info.version,
            db_size: db_info.size,
        })
    }

    pub async fn serve(self) -> Result<()> {
        let socket = SocketAddr::new(self.ip, self.port);
        let listener = TcpListener::bind(socket).await.unwrap();
        println!("Listening on {socket:?}");
        axum::serve(listener, http::app(Arc::new(self)).into_make_service()).await?;

        Ok(())
    }
}