// Copyright (c) 2024 Jason Witty <jasonpwitty+socktop@proton.me>.
// All rights reserved.
//
// Security validation for command execution

use std::path::Path;

/// Error type for command validation failures
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct ValidationError {
    pub message: String,
}

impl ValidationError {
    fn new(message: impl Into<String>) -> Self {
        Self {
            message: message.into(),
        }
    }
}

impl std::fmt::Display for ValidationError {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "Command validation error: {}", self.message)
    }
}

impl std::error::Error for ValidationError {}

/// Whitelist of allowed shell commands
const ALLOWED_SHELLS: &[&str] = &[
    "/bin/sh",
    "/bin/bash",
    "/bin/zsh",
    "/bin/dash",
    "/usr/bin/bash",
    "/usr/bin/zsh",
    "/usr/bin/fish",
];

/// Maximum allowed command path length
const MAX_COMMAND_LENGTH: usize = 4096;

/// Validates a command path for security concerns
///
/// # Security Checks
/// - Must be an absolute path
/// - Must not contain path traversal sequences (..)
/// - Must not contain shell metacharacters
/// - Must not contain null bytes
/// - Must be ASCII only
/// - Must not be in user-writable directories
/// - Must be in the whitelist (if whitelist checking is enabled)
/// - Must not exceed maximum length
///
/// # Example
/// ```
/// use webterm::security::validate_command;
///
/// // Valid command
/// assert!(validate_command("/bin/sh", true).is_ok());
///
/// // Invalid: shell injection attempt
/// assert!(validate_command("/bin/sh; rm -rf /", true).is_err());
///
/// // Invalid: path traversal
/// assert!(validate_command("../../bin/sh", true).is_err());
/// ```
pub fn validate_command(command: &str, check_whitelist: bool) -> Result<(), ValidationError> {
    // Check for empty command
    if command.is_empty() {
        return Err(ValidationError::new("Command cannot be empty"));
    }

    // Check length
    if command.len() > MAX_COMMAND_LENGTH {
        return Err(ValidationError::new(format!(
            "Command path too long: {} bytes (max: {})",
            command.len(),
            MAX_COMMAND_LENGTH
        )));
    }

    // Must be absolute path
    if !command.starts_with('/') {
        return Err(ValidationError::new(
            "Command must be an absolute path starting with '/'",
        ));
    }

    // Check for path traversal
    if command.contains("..") {
        return Err(ValidationError::new(
            "Command path contains '..' (path traversal attempt)",
        ));
    }

    // Check for shell metacharacters
    let dangerous_chars = [';', '&', '|', '`', '$', '\n', '\r', '\0', '<', '>'];
    for ch in dangerous_chars {
        if command.contains(ch) {
            return Err(ValidationError::new(format!(
                "Command contains dangerous character: {:?}",
                ch
            )));
        }
    }

    // Check for null bytes
    if command.contains('\0') {
        return Err(ValidationError::new("Command contains null byte"));
    }

    // Must be ASCII only (avoid Unicode tricks)
    if !command.is_ascii() {
        return Err(ValidationError::new("Command must be ASCII only"));
    }

    // Check for control characters (except common ones that might be in paths)
    for ch in command.chars() {
        if ch.is_control() {
            return Err(ValidationError::new(format!(
                "Command contains control character: {:?}",
                ch
            )));
        }
    }

    // Must not be in user-writable directories
    let dangerous_prefixes = ["/tmp/", "/var/tmp/", "/home/", "/Users/", "/root/"];
    for prefix in dangerous_prefixes {
        if command.starts_with(prefix) {
            return Err(ValidationError::new(format!(
                "Command in user-writable directory: {}",
                prefix
            )));
        }
    }

    // Check against whitelist if enabled
    if check_whitelist && !ALLOWED_SHELLS.contains(&command) {
        return Err(ValidationError::new(format!(
            "Command '{}' not in whitelist. Allowed: {:?}",
            command, ALLOWED_SHELLS
        )));
    }

    // Verify the command exists (if not checking whitelist, we should at least verify it's a file)
    if !check_whitelist {
        let path = Path::new(command);
        if !path.exists() {
            return Err(ValidationError::new(format!(
                "Command path does not exist: {}",
                command
            )));
        }
        if !path.is_file() {
            return Err(ValidationError::new(format!(
                "Command path is not a file: {}",
                command
            )));
        }
    }

    Ok(())
}

/// Validates an environment variable value for security
///
/// # Security Checks
/// - Must not contain shell metacharacters
/// - Must not contain null bytes
/// - Must not contain newlines
/// - Must be reasonable length
///
/// # Example
/// ```
/// use webterm::security::validate_env_value;
///
/// assert!(validate_env_value("xterm").is_ok());
/// assert!(validate_env_value("xterm; rm -rf /").is_err());
/// ```
pub fn validate_env_value(value: &str) -> Result<(), ValidationError> {
    // Check length
    if value.len() > 4096 {
        return Err(ValidationError::new("Environment value too long"));
    }

    // Check for dangerous characters
    let dangerous_chars = [';', '&', '|', '`', '$', '\n', '\r', '\0'];
    for ch in dangerous_chars {
        if value.contains(ch) {
            return Err(ValidationError::new(format!(
                "Environment value contains dangerous character: {:?}",
                ch
            )));
        }
    }

    Ok(())
}

/// Returns the list of allowed shells
pub fn allowed_shells() -> &'static [&'static str] {
    ALLOWED_SHELLS
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_valid_command() {
        assert!(validate_command("/bin/sh", true).is_ok());
        assert!(validate_command("/bin/bash", true).is_ok());
    }

    #[test]
    fn test_reject_relative_path() {
        assert!(validate_command("bin/sh", true).is_err());
        assert!(validate_command("./bin/sh", true).is_err());
    }

    #[test]
    fn test_reject_path_traversal() {
        assert!(validate_command("/../bin/sh", true).is_err());
        assert!(validate_command("/bin/../sh", true).is_err());
    }

    #[test]
    fn test_reject_shell_metacharacters() {
        assert!(validate_command("/bin/sh;", true).is_err());
        assert!(validate_command("/bin/sh&", true).is_err());
        assert!(validate_command("/bin/sh|", true).is_err());
        assert!(validate_command("/bin/sh`", true).is_err());
        assert!(validate_command("/bin/sh$", true).is_err());
    }

    #[test]
    fn test_reject_null_bytes() {
        assert!(validate_command("/bin/sh\0", true).is_err());
    }

    #[test]
    fn test_reject_tmp_directory() {
        assert!(validate_command("/tmp/malicious.sh", true).is_err());
        assert!(validate_command("/var/tmp/evil.sh", true).is_err());
    }

    #[test]
    fn test_reject_non_ascii() {
        assert!(validate_command("/bin/sh™", true).is_err());
    }

    #[test]
    fn test_reject_empty_command() {
        assert!(validate_command("", true).is_err());
    }

    #[test]
    fn test_reject_too_long() {
        let long_command = format!("/{}", "a".repeat(5000));
        assert!(validate_command(&long_command, true).is_err());
    }

    #[test]
    fn test_whitelist_enforcement() {
        assert!(validate_command("/bin/sh", true).is_ok());
        assert!(validate_command("/usr/local/bin/custom", true).is_err());
    }

    #[test]
    fn test_valid_env_value() {
        assert!(validate_env_value("xterm").is_ok());
        assert!(validate_env_value("xterm-256color").is_ok());
    }

    #[test]
    fn test_reject_env_with_metacharacters() {
        assert!(validate_env_value("xterm; rm -rf /").is_err());
        assert!(validate_env_value("xterm && curl evil.com").is_err());
    }
}
