P-Aimon-Pen-ROFMAD/src/web/handlers/hosts.rs

// src/web/handlers/hosts.rs - Host-related page handlers

use askama::Template;
use axum::{
    extract::{Path, Query, State, Form},
    response::{Html, IntoResponse},
};
use serde::{Deserialize, Serialize};
use std::collections::HashMap;
use sqlx::Row;
use crate::sql::ProjectContext;
use super::utils::{WebResult, WebError, parse_host_summary, get_bool_from_row, xml_ports_to_port_info, calculate_host_progress, parse_action_outputs, find_host_screenshots, get_combined_tags_with_source, TagWithSource, HostSummary, PortInfo, ActionOutput, ScreenshotInfo, HostProgress};


// Add to the existing HostDetail struct in handlers.rs
#[derive(Serialize)]
pub struct HostDetail {
    pub ip: String,
    pub hostname: String,
    pub short_hostname: String,
    pub tags: Vec<TagWithSource>, // Combined manual + refined with source info
    pub automated_tags: Vec<String>, // Raw automated tags for bottom section
    pub ad_domain: String,
    pub ports: Vec<PortInfo>,
    pub progress: HostProgress,
    
    // Raw nmap scan data - simplified to avoid filter issues
    pub nmap_scan_raw: String,
    
    // Action outputs split by whether they produced tags
    pub action_findings: Vec<ActionOutput>,
    pub action_other: Vec<ActionOutput>,
    
    // HTTP services and screenshots
    pub http_services: HashMap<String, crate::http::probe::HttpServiceInfo>,
    pub screenshots: Vec<ScreenshotInfo>,

    // Nuclei findings
    pub nuclei_results: Vec<crate::sql::nuclei::NucleiPortResult>,
    pub nuclei_findings_count: usize,
    pub nuclei_scanned_ports: Vec<String>,
}

#[derive(Deserialize)]
pub struct HostsListParams {
    pub sort:   Option<String>,
    pub dir:    Option<String>,
    pub filter: Option<String>,
}

#[derive(Template)]
#[template(path = "hosts.html")]
pub struct HostsListTemplate {
    pub hosts:   Vec<HostSummary>,
    pub sort:    String,
    pub dir:     String,
    pub filter:  String,
    // Precomputed filter URLs (preserve current sort)
    pub url_filter_clear:           String,
    pub url_filter_interesting:     String,
    pub url_filter_not_interesting: String,
    pub url_filter_analyzed_ready:  String,
    pub url_filter_analyzed_done:   String,
}

#[derive(Template)]
#[template(path = "hosts_rows.html")]
pub struct HostsRowsTemplate {
    pub hosts: Vec<HostSummary>,
}

fn make_filter_url(filter_val: &str, cur_sort: &str, cur_dir: &str) -> String {
    if filter_val.is_empty() {
        format!("/hosts?sort={}&dir={}", cur_sort, cur_dir)
    } else {
        format!("/hosts?sort={}&dir={}&filter={}", cur_sort, cur_dir, filter_val)
    }
}

#[derive(Template)]
#[template(path = "host_detail.html")]
pub struct HostDetailTemplate {
    pub host: HostDetail,
}

#[derive(Deserialize)]
pub struct InterestingForm {
    pub ip: String,
    pub interesting: bool,
}

// Add new form struct for analyzed status
#[derive(Deserialize)]
pub struct AnalyzedForm {
    pub ip: String,
    pub analyzed: bool,
}

#[derive(Deserialize)]
pub struct NucleiReviewedForm {
    pub ip: String,
    pub nuclei_reviewed: bool,
}

// Shared helper: whitelists params, builds SQL, returns sorted host list
async fn fetch_hosts_sorted(
    ctx: &ProjectContext,
    params: &HostsListParams,
) -> Result<(Vec<HostSummary>, &'static str, &'static str, &'static str), WebError> {
    let sort: &'static str = match params.sort.as_deref() {
        Some("hostname")        => "hostname",
        Some("ports")           => "ports",
        Some("http")            => "http",
        Some("analyzed")        => "analyzed",
        Some("interesting")     => "interesting",
        Some("nuclei_reviewed") => "nuclei_reviewed",
        _                       => "ip",
    };
    let dir: &'static str = match params.dir.as_deref() {
        Some("desc") => "desc",
        _            => "asc",
    };
    let filter: &'static str = match params.filter.as_deref() {
        Some("interesting")     => "interesting",
        Some("not_interesting") => "not_interesting",
        Some("analyzed_ready")  => "analyzed_ready",
        Some("analyzed_done")   => "analyzed_done",
        _                       => "",
    };

    let filter_clause = match filter {
        "interesting"     => " AND interesting = 1",
        "not_interesting" => " AND (interesting = 0 OR interesting IS NULL)",
        "analyzed_ready"  => " AND full_scan_complete = 1 AND service_scan_complete = 1 \
                               AND http_discovery_complete = 1 AND screenshots_complete = 1 \
                               AND COALESCE(user_workflow_complete, 0) = 1 \
                               AND (analyzed = 0 OR analyzed IS NULL)",
        "analyzed_done"   => " AND analyzed = 1",
        _                 => "",
    };

    // All values in order_clause are static strings — no user input reaches the query
    let order_clause: &str = match (sort, dir) {
        ("hostname", "asc")    => "CASE WHEN hostname IS NULL OR hostname = '' THEN 1 ELSE 0 END ASC, LOWER(COALESCE(hostname,'')) ASC",
        ("hostname", _)        => "CASE WHEN hostname IS NULL OR hostname = '' THEN 1 ELSE 0 END ASC, LOWER(COALESCE(hostname,'')) DESC",
        ("ports", "asc")       => "CASE WHEN ports IS NULL OR ports = '' THEN 0 ELSE LENGTH(ports) - LENGTH(REPLACE(ports, ',', '')) + 1 END ASC",
        ("ports", _)           => "CASE WHEN ports IS NULL OR ports = '' THEN 0 ELSE LENGTH(ports) - LENGTH(REPLACE(ports, ',', '')) + 1 END DESC",
        ("http", "asc")        => "(SELECT COUNT(*) FROM http_services WHERE run_id = targets.run_id AND ip = targets.ip) ASC",
        ("http", _)            => "(SELECT COUNT(*) FROM http_services WHERE run_id = targets.run_id AND ip = targets.ip) DESC",
        ("analyzed", "asc")        => "COALESCE(analyzed, 0) ASC, ip ASC",
        ("analyzed", _)            => "COALESCE(analyzed, 0) DESC, ip ASC",
        ("interesting", "asc")     => "COALESCE(interesting, 0) ASC, ip ASC",
        ("interesting", _)         => "COALESCE(interesting, 0) DESC, ip ASC",
        ("nuclei_reviewed", "asc") => "COALESCE(nuclei_reviewed, 0) ASC, ip ASC",
        ("nuclei_reviewed", _)     => "COALESCE(nuclei_reviewed, 0) DESC, ip ASC",
        (_, "desc")                => "ip DESC",
        _                          => "ip ASC",
    };

    let query_str = format!(
        "SELECT ip, \
         COALESCE(hostname,'') as hostname, \
         COALESCE(ports,'') as ports, \
         COALESCE(ad_domain,'') as ad_domain, \
         COALESCE(full_scan_complete, 0) as full_scan_complete, \
         COALESCE(service_scan_complete, 0) as service_scan_complete, \
         COALESCE(http_discovery_complete, 0) as http_discovery_complete, \
         COALESCE(screenshots_complete, 0) as screenshots_complete, \
         (SELECT COUNT(*) FROM action_outputs WHERE run_id = targets.run_id AND ip = targets.ip) as action_count, \
         COALESCE(analyzed, 0) as analyzed, \
         COALESCE(interesting, 0) as interesting, \
         COALESCE(nuclei_reviewed, 0) as nuclei_reviewed, \
         CASE WHEN EXISTS (SELECT 1 FROM action_outputs WHERE run_id = targets.run_id AND ip = targets.ip AND action_name LIKE 'nuclei_%') THEN 1 ELSE 0 END as has_nuclei_findings \
         FROM targets WHERE run_id = ?{} ORDER BY {}",
        filter_clause, order_clause
    );

    let rows = sqlx::query(&query_str)
        .bind(&ctx.run_id)
        .fetch_all(&ctx.pool)
        .await
        .map_err(|e| WebError::DatabaseError(e.to_string()))?;

    let mut hosts = Vec::new();
    for row in rows {
        hosts.push(parse_host_summary(row, ctx).await);
    }

    Ok((hosts, sort, dir, filter))
}

// Context: Provides list view of all discovered hosts with server-side sort and filter
// Operation: Builds SQL from whitelisted query params, renders hosts listing template
pub async fn hosts_list(
    Query(params): Query<HostsListParams>,
    State(ctx): State<ProjectContext>,
) -> WebResult<impl IntoResponse> {
    let (hosts, sort, dir, filter) = fetch_hosts_sorted(&ctx, &params).await?;

    let template = HostsListTemplate {
        hosts,
        sort:   sort.to_string(),
        dir:    dir.to_string(),
        filter: filter.to_string(),
        url_filter_clear:           make_filter_url("",                sort, dir),
        url_filter_interesting:     make_filter_url("interesting",     sort, dir),
        url_filter_not_interesting: make_filter_url("not_interesting", sort, dir),
        url_filter_analyzed_ready:  make_filter_url("analyzed_ready",  sort, dir),
        url_filter_analyzed_done:   make_filter_url("analyzed_done",   sort, dir),
    };
    Ok(Html(template.render().map_err(|e| WebError::ParseError(e.to_string()))?))
}

// Context: Returns only the <tbody> row HTML for AJAX table refresh (no full page reload)
// Operation: Runs same SQL as hosts_list, renders fragment template for client-side tbody replacement
pub async fn hosts_rows(
    Query(params): Query<HostsListParams>,
    State(ctx): State<ProjectContext>,
) -> WebResult<impl IntoResponse> {
    let (hosts, _sort, _dir, _filter) = fetch_hosts_sorted(&ctx, &params).await?;
    let template = HostsRowsTemplate { hosts };
    Ok(Html(template.render().map_err(|e| WebError::ParseError(e.to_string()))?))
}

// Context: Builds port list for host detail using XML as sole structured source
// Operation: Parses nmap_xml JSON map, extracts service_scan XML, converts to PortInfo.
fn parse_ports_from_xml(nmap_xml: &str) -> Vec<PortInfo> {
    if !nmap_xml.is_empty() && nmap_xml != "{}" {
        if let Ok(map) = serde_json::from_str::<HashMap<String, String>>(nmap_xml) {
            if let Some(xml) = map.get("service_scan") {
                if let Some(result) = crate::tag::parse_nmap_xml(xml) {
                    let mut ports = xml_ports_to_port_info(result.ports);
                    ports.sort_by_key(|p| p.port.parse::<u32>().unwrap_or(0));
                    return ports;
                }
            }
        }
    }
    Vec::new()
}

// Context: Displays comprehensive information for a specific host including scan results
// Operation: Aggregates host data, ports, services, screenshots and renders detail template
pub async fn host_detail(Path(ip): Path<String>, State(ctx): State<ProjectContext>) -> WebResult<impl IntoResponse> {
    let row = sqlx::query("SELECT * FROM targets WHERE run_id = ? AND ip = ?")
        .bind(&ctx.run_id)
        .bind(&ip)
        .fetch_optional(&ctx.pool)
        .await
        .map_err(|e| WebError::DatabaseError(e.to_string()))?
        .ok_or(WebError::NotFound(format!("Host {} not found", ip)))?;
    
    // Extract basic data
    let hostname: Option<String> = row.get("hostname");
    
    // Get combined tags with source info
    let tags = get_combined_tags_with_source(&ctx, &ip).await.unwrap_or_default();
    
    // Get automated tags separately
    let automated_tags_set = ctx.get_automated_tags(&ip).await.unwrap_or_default();
    let mut automated_tags: Vec<String> = automated_tags_set.into_iter().collect();
    automated_tags.sort();
    
    let ad_domain: Option<String> = row.get("ad_domain");

    // Raw nmap text — kept solely for the human-readable display block
    let nmap_scan_str: String = row.try_get("nmap_scan").unwrap_or_else(|_| "{}".to_string());
    let nmap_display_text = if let Ok(nmap_json) = serde_json::from_str::<HashMap<String, String>>(&nmap_scan_str) {
        if let Some(service_scan) = nmap_json.get("service_scan") {
            service_scan.clone()
        } else {
            let mut parts: Vec<_> = nmap_json.into_values().collect();
            parts.sort();
            parts.join("\n\n")
        }
    } else {
        nmap_scan_str.clone()
    };

    // XML is the sole structured source for port/service data
    let nmap_xml_str: String = row.try_get("nmap_xml").unwrap_or_else(|_| "{}".to_string());
    let ports = parse_ports_from_xml(&nmap_xml_str);

    // Fetch action outputs from normalized table
    let all_action_outputs = ctx.get_action_outputs_for_host(&ip).await
        .map(parse_action_outputs)
        .unwrap_or_default();
    let action_count = all_action_outputs.len();
    let (action_findings, action_other): (Vec<_>, Vec<_>) = all_action_outputs.into_iter().partition(|a| a.tagged);

    // Calculate progress with corrected logic
    let progress = calculate_host_progress(
        get_bool_from_row(&row, "full_scan_complete"),
        get_bool_from_row(&row, "service_scan_complete"),
        get_bool_from_row(&row, "http_discovery_complete"),
        get_bool_from_row(&row, "screenshots_complete"),
        action_count,
    );

    // Use the dedicated table
    let http_services_raw = ctx.get_http_services(&ip).await.unwrap_or_default();
    let mut http_services = HashMap::new();
    for (port, service_json) in http_services_raw {
        if let Ok(service_info) = serde_json::from_str::<crate::http::probe::HttpServiceInfo>(&service_json) {
            http_services.insert(port, service_info);
        }
    }
    
    // Find screenshots
    let screenshots_dir = crate::config::paths::get_screenshots_directory(&ctx.db_path, &ctx.run_id);
    let screenshots = find_host_screenshots(&ip, &http_services, &screenshots_dir);

    // Nuclei findings
    let mut nuclei_results = ctx.get_nuclei_results_for_host(&ip).await.unwrap_or_default();
    for result in &mut nuclei_results {
        if let Some(svc) = http_services.get(&result.port) {
            result.url = format!("{}://{}:{}", svc.protocol, ip, result.port);
        }
    }
    let nuclei_findings_count: usize = nuclei_results.iter().map(|r| r.findings.len()).sum();
    let nuclei_scanned_ports: Vec<String> = nuclei_results.iter().map(|r| r.port.clone()).collect();

    let hostname_str = hostname.unwrap_or_default();
    let ad_domain_str = ad_domain.unwrap_or_default();
    let short_hostname = crate::web::handlers::utils::strip_domain_suffix(&hostname_str, &ad_domain_str);

    let host_detail = HostDetail {
        ip: ip.clone(),
        hostname: hostname_str,
        short_hostname,
        tags,
        automated_tags,
        ad_domain: ad_domain_str,
        ports,
        progress,
        nmap_scan_raw: nmap_display_text,
        action_findings,
        action_other,
        http_services,
        screenshots,
        nuclei_results,
        nuclei_findings_count,
        nuclei_scanned_ports,
    };
    
    let template = HostDetailTemplate { host: host_detail };
    Ok(Html(template.render().map_err(|e| WebError::ParseError(e.to_string()))?))
}

// Handler for marking hosts as analyzed
// Context: Allows marking hosts as analyzed for workflow progress tracking
// Operation: Updates host analysis status with retry logic for database consistency
pub async fn mark_host_analyzed(
    State(ctx): State<ProjectContext>,
    Form(form): Form<AnalyzedForm>
) -> WebResult<impl IntoResponse> {
    // Add timeout and retry logic
    let mut attempts = 0;
    let max_attempts = 3;
    
    while attempts < max_attempts {
        match ctx.mark_host_analyzed(&form.ip, form.analyzed).await {
            Ok(_) => {
                return Ok(axum::Json(serde_json::json!({
                    "success": true,
                    "ip": form.ip,
                    "analyzed": form.analyzed,
                    "attempts": attempts + 1
                })));
            }
            Err(e) if attempts < max_attempts - 1 => {
                attempts += 1;
                eprintln!("Attempt {} failed for host {}: {}", attempts, form.ip, e);
                tokio::time::sleep(tokio::time::Duration::from_millis(100)).await;
            }
            Err(e) => {
                eprintln!("Final attempt failed for host {}: {}", form.ip, e);
                return Err(WebError::DatabaseError(format!("Failed to update analysis status after {} attempts: {}", max_attempts, e)));
            }
        }
    }
    
    unreachable!()
}

// Context: Provides current analysis status for host progress tracking
// Operation: Queries database and returns JSON with host analysis boolean status
pub async fn get_host_analysis_status(
    Path(ip): Path<String>, 
    State(ctx): State<ProjectContext>
) -> WebResult<impl IntoResponse> {
    let analyzed = ctx.is_host_analyzed(&ip).await
        .map_err(|e| WebError::DatabaseError(e.to_string()))?;
    
    Ok(axum::Json(serde_json::json!({
        "ip": ip,
        "analyzed": analyzed
    })))
}

// Context: Enables marking hosts as interesting for prioritized analysis tracking
// Operation: Updates host interesting flag in database and returns confirmation
pub async fn mark_host_interesting(
    State(ctx): State<ProjectContext>,
    Form(form): Form<InterestingForm>
) -> WebResult<impl IntoResponse> {
    ctx.mark_host_interesting(&form.ip, form.interesting).await
        .map_err(|e| WebError::DatabaseError(e.to_string()))?;
    
    Ok(axum::Json(serde_json::json!({
        "success": true,
        "ip": form.ip,
        "interesting": form.interesting
    })))
}

// Context: Retrieves current interesting status for host prioritization display
// Operation: Queries database and returns JSON with host interesting boolean flag
pub async fn get_host_interesting_status(
    Path(ip): Path<String>, 
    State(ctx): State<ProjectContext>
) -> WebResult<impl IntoResponse> {
    let interesting = ctx.is_host_interesting(&ip).await
        .map_err(|e| WebError::DatabaseError(e.to_string()))?;
    
    Ok(axum::Json(serde_json::json!({
        "ip": ip,
        "interesting": interesting
    })))
}

pub async fn mark_host_nuclei_reviewed(
    State(ctx): State<ProjectContext>,
    Form(form): Form<NucleiReviewedForm>,
) -> WebResult<impl IntoResponse> {
    ctx.mark_host_nuclei_reviewed(&form.ip, form.nuclei_reviewed).await
        .map_err(|e| WebError::DatabaseError(e.to_string()))?;
    Ok(axum::Json(serde_json::json!({
        "success": true, "ip": form.ip, "nuclei_reviewed": form.nuclei_reviewed
    })))
}

pub async fn get_host_nuclei_reviewed_status(
    Path(ip): Path<String>,
    State(ctx): State<ProjectContext>,
) -> WebResult<impl IntoResponse> {
    let reviewed = ctx.is_host_nuclei_reviewed(&ip).await
        .map_err(|e| WebError::DatabaseError(e.to_string()))?;
    Ok(axum::Json(serde_json::json!({ "ip": ip, "nuclei_reviewed": reviewed })))
}

// Context: Serves screenshot images with security validation for web interface display
// Operation: Validates filename format and serves PNG files from project-specific screenshots directory
pub async fn screenshot(Path(filename): Path<String>, State(ctx): State<ProjectContext>) -> WebResult<impl IntoResponse> {
    // Enhanced filename validation for security
    if filename.contains("..") || filename.contains("/") || filename.contains("\\") {
        return Err(WebError::InvalidData("Invalid filename".to_string()));
    }
    
    // Must be PNG with expected format
    if !filename.ends_with(".png") {
        return Err(WebError::InvalidData("Invalid file type".to_string()));
    }
    
    // Validate filename format: IP_PORT_PROTOCOL_TIMESTAMP.png
    let name_without_ext = filename.strip_suffix(".png").unwrap();
    let parts: Vec<&str> = name_without_ext.split('_').collect();
    
    if parts.len() < 7 {
        return Err(WebError::InvalidData("Invalid filename format".to_string()));
    }
    
    // Validate IP format (4 numeric parts)
    for i in 0..4 {
        if parts[i].parse::<u8>().is_err() {
            return Err(WebError::InvalidData("Invalid IP format in filename".to_string()));
        }
    }
    
    // Validate port
    if parts[4].parse::<u16>().is_err() {
        return Err(WebError::InvalidData("Invalid port in filename".to_string()));
    }
    
    // Validate protocol
    if !["http", "https"].contains(&parts[5]) {
        return Err(WebError::InvalidData("Invalid protocol in filename".to_string()));
    }
    
    // Validate access within current project and get file path
    if !crate::config::paths::validate_screenshot_access(&ctx.db_path, &ctx.run_id, &filename) {
        return Err(WebError::NotFound("Screenshot not found in current project".to_string()));
    }
    
    let file_path = crate::config::paths::get_screenshot_file_path(&ctx.db_path, &ctx.run_id, &filename);
    
    if std::path::Path::new(&file_path).exists() {
        let contents = tokio::fs::read(&file_path)
            .await
            .map_err(|e| WebError::NotFound(format!("File read error: {}", e)))?;
        
        let content_type = "image/png";
        
        Ok(([(axum::http::header::CONTENT_TYPE, content_type)], contents))
    } else {
        Err(WebError::NotFound("Screenshot not found".to_string()))
    }
}