Eris/src/utils.rs

use std::collections::HashMap;
use std::hash::Hasher;

// Find end of HTTP headers
pub fn find_header_end(data: &[u8]) -> Option<usize> {
    data.windows(4)
        .position(|window| window == b"\r\n\r\n")
        .map(|pos| pos + 4)
}

// Extract path from raw request data
pub fn extract_path_from_request(data: &[u8]) -> Option<&str> {
    // Get first line from request
    let first_line = data
        .split(|&b| b == b'\r' || b == b'\n')
        .next()
        .filter(|line| !line.is_empty())?;

    // Split by spaces and ensure we have at least 3 parts (METHOD PATH VERSION)
    let parts: Vec<&[u8]> = first_line.split(|&b| b == b' ').collect();
    if parts.len() < 3 || !parts[2].starts_with(b"HTTP/") {
        return None;
    }

    // Return the path (second element)
    std::str::from_utf8(parts[1]).ok()
}

// Extract header value from raw request data
pub fn extract_header_value(data: &[u8], header_name: &str) -> Option<String> {
    let data_str = std::str::from_utf8(data).ok()?;
    let header_prefix = format!("{header_name}: ").to_lowercase();

    for line in data_str.lines() {
        let line_lower = line.to_lowercase();
        if line_lower.starts_with(&header_prefix) {
            return Some(line[header_prefix.len()..].trim().to_string());
        }
    }
    None
}

// Extract all headers from request data
pub fn extract_all_headers(data: &[u8]) -> HashMap<String, String> {
    let mut headers = HashMap::new();

    if let Ok(data_str) = std::str::from_utf8(data) {
        let mut lines = data_str.lines();

        // Skip the request line
        let _ = lines.next();

        // Parse headers until empty line
        for line in lines {
            if line.is_empty() {
                break;
            }

            if let Some(colon_pos) = line.find(':') {
                let key = line[..colon_pos].trim().to_lowercase();
                let value = line[colon_pos + 1..].trim().to_string();
                headers.insert(key, value);
            }
        }
    }

    headers
}

// Determine response type based on request path
pub fn choose_response_type(path: &str) -> &'static str {
    if path.contains("phpunit") || path.contains("eval") {
        "php_exploit"
    } else if path.contains("wp-") {
        "wordpress"
    } else if path.contains("api") {
        "api"
    } else {
        "generic"
    }
}

// Get current timestamp in seconds
pub fn get_timestamp() -> u64 {
    std::time::SystemTime::now()
        .duration_since(std::time::UNIX_EPOCH)
        .unwrap_or_default()
        .as_secs()
}

// Create a unique session ID for tracking a connection
pub fn generate_session_id(ip: &str, user_agent: &str) -> String {
    let timestamp = get_timestamp();
    let random = rand::random::<u32>();

    // XXX: Is this fast enough for our case? I don't think hashing is a huge
    // bottleneck, but it's worth revisiting in the future to see if there is
    // an objectively faster algorithm that we can try.
    let mut hasher = std::collections::hash_map::DefaultHasher::new();
    std::hash::Hash::hash(&format!("{ip}_{user_agent}_{timestamp}"), &mut hasher);
    let hash = hasher.finish();

    format!("SID_{hash:x}_{random:x}")
}

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

    #[test]
    fn test_find_header_end() {
        let data = b"GET / HTTP/1.1\r\nHost: example.com\r\nUser-Agent: test\r\n\r\nBody content";
        assert_eq!(find_header_end(data), Some(55));

        let incomplete = b"GET / HTTP/1.1\r\nHost: example.com\r\n";
        assert_eq!(find_header_end(incomplete), None);
    }

    #[test]
    fn test_extract_path_from_request() {
        let data = b"GET /index.html HTTP/1.1\r\nHost: example.com\r\n\r\n";
        assert_eq!(extract_path_from_request(data), Some("/index.html"));

        let bad_data = b"INVALID DATA";
        assert_eq!(extract_path_from_request(bad_data), None);
    }

    #[test]
    fn test_extract_header_value() {
        let data = b"GET / HTTP/1.1\r\nHost: example.com\r\nUser-Agent: TestBot/1.0\r\n\r\n";
        assert_eq!(
            extract_header_value(data, "user-agent"),
            Some("TestBot/1.0".to_string())
        );
        assert_eq!(
            extract_header_value(data, "Host"),
            Some("example.com".to_string())
        );
        assert_eq!(extract_header_value(data, "nonexistent"), None);
    }

    #[test]
    fn test_extract_all_headers() {
        let data = b"GET / HTTP/1.1\r\nHost: example.com\r\nUser-Agent: TestBot/1.0\r\nAccept: */*\r\n\r\n";
        let headers = extract_all_headers(data);

        assert_eq!(headers.len(), 3);
        assert_eq!(headers.get("host").unwrap(), "example.com");
        assert_eq!(headers.get("user-agent").unwrap(), "TestBot/1.0");
        assert_eq!(headers.get("accept").unwrap(), "*/*");
    }

    #[test]
    fn test_choose_response_type() {
        assert_eq!(
            choose_response_type("/vendor/phpunit/whatever"),
            "php_exploit"
        );
        assert_eq!(
            choose_response_type("/path/to/eval-stdin.php"),
            "php_exploit"
        );
        assert_eq!(choose_response_type("/wp-admin/login.php"), "wordpress");
        assert_eq!(choose_response_type("/wp-login.php"), "wordpress");
        assert_eq!(choose_response_type("/api/v1/users"), "api");
        assert_eq!(choose_response_type("/index.html"), "generic");
    }
}