P-Aimon-Pen-ROFMAD/src/sql/subnets.rs

// src/sql/subnets.rs
use sqlx::Row;
use ipnetwork::IpNetwork;
use std::error::Error;
use crate::workflow::{dispatch_event, Hook};
use crate::utilities::{tag_info, tag_debug};
use crate::VERBOSE;
use super::context::ProjectContext;

impl ProjectContext {
    // Context: Subnet scanning phase completion tracking for workflow progression management
    // Operation: Updates phase completion status with dynamic schema migration and logging
    pub async fn mark_subnet_phase_complete(&self, subnet: &str, phase: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        let column = match phase {
            "subnet_phase1" => "phase1_complete",
            "subnet_phase2" => "phase2_complete", 
            _ => return Err(format!("Unknown subnet phase: {}", phase).into()),
        };
        
        // Add columns if they don't exist
        let info = sqlx::query("PRAGMA table_info(subnets);")
            .fetch_all(&self.pool)
            .await?;
        if !info.iter().any(|r| r.get::<String, _>("name") == column) {
            let alter_sql = format!("ALTER TABLE subnets ADD COLUMN {} BOOLEAN DEFAULT FALSE", column);
            sqlx::query(&alter_sql).execute(&self.pool).await?;
        }
        
        let update_sql = format!("UPDATE subnets SET {} = TRUE WHERE run_id = ? AND target = ?", column);
        sqlx::query(&update_sql)
            .bind(&self.run_id)
            .bind(subnet)
            .execute(&self.pool)
            .await?;
        
        if VERBOSE.load(std::sync::atomic::Ordering::SeqCst) {
            tag_debug(&format!("Marked {} complete for subnet {}", phase, subnet));
        }
        
        Ok(())
    }

    // Context: Subnet phase verification for workflow decision making and dependency checking
    // Operation: Queries database to check if specific subnet phase has been completed
    pub async fn is_subnet_phase_complete(&self, subnet: &str, phase: &str) -> Result<bool, Box<dyn Error + Send + Sync>> {
        let column = match phase {
            "subnet_phase1" => "phase1_complete",
            "subnet_phase2" => "phase2_complete",
            _ => return Ok(false),
        };
        
        // Check if column exists
        let info = sqlx::query("PRAGMA table_info(subnets);")
            .fetch_all(&self.pool)
            .await?;
        if !info.iter().any(|r| r.get::<String, _>("name") == column) {
            return Ok(false); // Column doesn't exist, so not complete
        }
        
        let query_sql = format!("SELECT {} FROM subnets WHERE run_id = ? AND target = ?", column);
        let complete: bool = sqlx::query_scalar(&query_sql)
            .bind(&self.run_id)
            .bind(subnet)
            .fetch_optional(&self.pool)
            .await?
            .unwrap_or(false);
        
        Ok(complete)
    }

    // Context: Masscan scanning progress tracking to prevent duplicate subnet enumeration
    // Operation: Marks subnet as having completed masscan discovery phase
    pub async fn mark_masscan_scanned(&self, subnet: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        sqlx::query(
            "UPDATE subnets SET masscan_scanned = 1
             WHERE run_id = ? AND target = ?",
        )
        .bind(&self.run_id)
        .bind(subnet)
        .execute(&self.pool)
        .await?;
        Ok(())
    }


    // Context: Host exploration phase tracking for comprehensive subnet enumeration
    // Operation: Marks subnet as having completed host exploration scanning phase
    pub async fn mark_explorer_scanned(&self, subnet: &str) -> Result<(), Box<dyn Error + Send + Sync>> {
        sqlx::query(
            "UPDATE subnets SET explorer_scanned = 1
             WHERE run_id = ? AND target = ?",
        )
        .bind(&self.run_id)
        .bind(subnet)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    // Context: Subnet insertion for masscan-managed paths that own their own Phase 1 → Phase 2 chain
    // Operation: Inserts subnet without firing Hook::Subnet — prevents Phase 1 queuing before masscan runs
    pub async fn insert_subnet_no_hook(
        &self,
        target: &str,
    ) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
        let new_net: IpNetwork = target.parse()?;
        if new_net.prefix() == 32 {
            return Ok(false);
        }
        let new_cidr = format!("{}/{}", new_net.network(), new_net.prefix());

        let (count,): (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM subnets WHERE run_id = ? AND target = ?",
        )
        .bind(&self.run_id)
        .bind(&new_cidr)
        .fetch_one(&self.pool)
        .await?;
        if count > 0 {
            return Ok(false);
        }

        sqlx::query("INSERT INTO subnets (run_id, target) VALUES (?1, ?2)")
            .bind(&self.run_id)
            .bind(&new_cidr)
            .execute(&self.pool)
            .await?;

        tag_info(&format!("New subnet (masscan): {}", new_cidr));
        Ok(true)
    }

    // Context: Forced subnet insertion for masscan-discovered networks without consolidation
    // Operation: Adds valid IPv4 networks to database, skipping duplicate checks and consolidation rules
    pub async fn insert_subnet_force(
        &self,
        target: &str,
    ) -> Result<bool, Box<dyn std::error::Error + Send + Sync>> {
        // Normalize & ensure it's a valid IPv4 network
        let new_net: IpNetwork = target.parse()?;
        // Don't treat a single host as a "subnet"
        if new_net.prefix() == 32 {
            return Ok(false);
        }
        let new_cidr = format!("{}/{}", new_net.network(), new_net.prefix());

        // Check for exact duplicate in the DB
        let (count,): (i64,) = sqlx::query_as(
            "SELECT COUNT(*) FROM subnets WHERE run_id = ? AND target = ?",
        )
        .bind(&self.run_id)
        .bind(&new_cidr)
        .fetch_one(&self.pool)
        .await?;
        if count > 0 {
            // Already present
            return Ok(false);
        }

        // Not found—insert it
        sqlx::query("INSERT INTO subnets (run_id, target) VALUES (?1, ?2)")
            .bind(&self.run_id)
            .bind(&new_cidr)
            .execute(&self.pool)
            .await?;

        tag_info(&format!("[MASSCAN] New Subnet with alive hosts discovered: {}", new_cidr));
        dispatch_event(Hook::Subnet, &new_cidr, self);
        Ok(true)
    }

    // Context: Intelligent subnet discovery system with consolidation rules and deduplication
    // Operation: Inserts subnet with /24 boundary rules, consolidation logic, and workflow hooks
    pub async fn insert_subnet(
    &self,
    target: &str,
) -> Result<bool, Box<dyn Error + Send + Sync>> {
    let new_net: IpNetwork = target.parse()?;
    if new_net.prefix() == 32 {
        return Ok(false);
    }
    let new_cidr = format!("{}/{}", new_net.network(), new_net.prefix());

    // Check for exact duplicate
    let (count,): (i64,) = sqlx::query_as(
        "SELECT COUNT(*) FROM subnets WHERE run_id = ? AND target = ?",
    )
    .bind(&self.run_id)
    .bind(&new_cidr)
    .fetch_one(&self.pool)
    .await?;
    if count > 0 {
        return Ok(false);
    }

    // RULE 2: /24 always bypasses consolidation
    if new_net.prefix() == 24 {
        sqlx::query("INSERT INTO subnets (run_id, target) VALUES (?1, ?2)")
            .bind(&self.run_id)
            .bind(&new_cidr)
            .execute(&self.pool)
            .await?;
        
        tag_info(&format!("New subnet: {}", new_cidr));
        dispatch_event(Hook::Subnet, &new_cidr, self);
        return Ok(true);
    }

    // For non-/24 subnets, apply consolidation rules
    let rows = sqlx::query("SELECT target FROM subnets WHERE run_id = ?")
        .bind(&self.run_id)
        .fetch_all(&self.pool)
        .await?;

    let mut to_delete = Vec::new();
    for row in rows {
        let existing_str: String = row.get("target");
        let existing_net: IpNetwork = existing_str.parse()?;
        
        // RULE 1: Only consolidate smaller subnets UP TO /24
        if new_net.prefix() < 24 && existing_net.prefix() < 24 {
            // Both broader than /24 - apply normal consolidation
            if existing_net.prefix() <= new_net.prefix()
                && existing_net.contains(new_net.network())
            {
                return Ok(false); // Already covered
            }
            if new_net.prefix() <= existing_net.prefix()
                && new_net.contains(existing_net.network())
            {
                to_delete.push(existing_str); // Replace with broader
            }
        } else if new_net.prefix() >= 24 && existing_net.prefix() >= 24 {
            // Both /24 or smaller - consolidate UP TO /24 only
            if existing_net.prefix() <= new_net.prefix()
                && existing_net.contains(new_net.network())
            {
                return Ok(false); // Already covered
            }
            if new_net.prefix() <= existing_net.prefix()
                && new_net.contains(existing_net.network())
                && new_net.prefix() >= 24  // Don't consolidate beyond /24
            {
                to_delete.push(existing_str); // Replace with consolidated /24
            }
        }
        // RULE 3: Don't consolidate across /24 boundary
        // If one is >/24 and other is </24, no consolidation
    }

    for old in to_delete {
        sqlx::query("DELETE FROM subnets WHERE run_id = ? AND target = ?")
            .bind(&self.run_id)
            .bind(&old)
            .execute(&self.pool)
            .await?;
    }

    sqlx::query("INSERT INTO subnets (run_id, target) VALUES (?1, ?2)")
        .bind(&self.run_id)
        .bind(&new_cidr)
        .execute(&self.pool)
        .await?;

    tag_info(&format!("New subnet: {}", new_cidr));
    dispatch_event(Hook::Subnet, &new_cidr, self);
    Ok(true)
}
    

    // Context: Phase-1 subnet data retrieval for host enumeration and phase progression
    // Operation: Returns subnet list with associated discovered hosts from phase-1 scanning
    pub async fn fetch_subnets_phase1(
        &self,
    ) -> Result<Vec<(String, Vec<String>)>, Box<dyn Error + Send + Sync>> {
        let rows = sqlx::query("SELECT target, phase1_hosts FROM subnets WHERE run_id = ?")
            .bind(&self.run_id)
            .fetch_all(&self.pool)
            .await?;

        let mut out = Vec::new();
        for row in rows {
            let target: String = row.get("target");
            let raw: Option<String> = row.get("phase1_hosts");
            let list = raw
                .unwrap_or_default()
                .split(',')
                .map(str::trim)
                .filter(|s| !s.is_empty())
                .map(String::from)
                .collect();
            out.push((target, list));
        }
        Ok(out)
    }

    // Context: Phase-1 host discovery result storage for tracking subnet enumeration progress
    // Operation: Updates comma-separated host list for specific subnet in database
    pub async fn update_phase1_hosts(
        &self,
        subnet: &str,
        hosts: &[String],
    ) -> Result<(), Box<dyn Error + Send + Sync>> {
        let joined = hosts.join(",");
        sqlx::query(
            "UPDATE subnets SET phase1_hosts = ? WHERE run_id = ? AND target = ?",
        )
        .bind(&joined)
        .bind(&self.run_id)
        .bind(subnet)
        .execute(&self.pool)
        .await?;
        Ok(())
    }

    // Context: Subnet rescan flag verification for handling incomplete or failed scans
    // Operation: Queries database to check if subnet has been marked for rescanning
    pub async fn is_subnet_marked_rescan(&self, subnet: &str) -> Result<bool, Box<dyn Error + Send + Sync>> {
        let (flag,): (i64,) = sqlx::query_as(
            "SELECT rescan FROM subnets WHERE run_id = ? AND target = ?",
        )
        .bind(&self.run_id)
        .bind(subnet)
        .fetch_one(&self.pool)
        .await?;
        Ok(flag != 0)
    }
}