自律型ペネトレーションテストと有向攻撃経路グラフ（DAG）探索エンジンのRust 2024実装

#![forbid(unsafe_code)]

use std::cmp::Ordering;
use std::collections::{BinaryHeap, BTreeMap, BTreeSet};

/// Unique identifier for a security state or privilege node in V.
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq, PartialOrd, Ord, Hash)]
pub struct NodeId(pub u32);

/// Atomic security capability flag bitmask (Preconditions & Postconditions).
#[repr(transparent)]
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct CapabilityMask(pub u64);

impl CapabilityMask {
    pub const NONE: Self = Self(0);
    pub const PUBLIC_NET: Self = Self(1 << 0);
    pub const SUBDOMAIN_COOKIE_SCOPE: Self = Self(1 << 1);
    pub const INTERNAL_HOSTNAME_KNOWN: Self = Self(1 << 2);
    pub const WEBHOOK_TOKEN_FORGED: Self = Self(1 << 3);
    pub const TENANT_ADMIN_EXEC: Self = Self(1 << 4);
    pub const CROWN_JEWEL_DB_READ: Self = Self(1 << 5);

    #[inline]
    pub const fn contains(self, required: Self) -> bool {
        (self.0 & required.0) == required.0
    }

    #[inline]
    pub const fn union(self, other: Self) -> Self {
        Self(self.0 | other.0)
    }
}

/// Non-destructive active probe observation confidence factor.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum ProbeObservation {
    /// Vulnerability actively confirmed via safe idempotent state reflection (1.00)
    ConfirmedReachable,
    /// Configuration indicator observed, partial compensating control present (0.60)
    PartialExposure,
    /// Theoretical surface only; unconfirmed by active non-destructive probe (0.25)
    TheoreticalOnly,
    /// Blocked at kernel perimeter (e.g., FreeBSD 14 pf synproxy / FORCE RLS) (0.05)
    MitigatedAtBoundary,
}

impl ProbeObservation {
    #[inline]
    pub const fn alpha(self) -> f64 {
        match self {
            Self::ConfirmedReachable => 1.00,
            Self::PartialExposure => 0.60,
            Self::TheoreticalOnly => 0.25,
            Self::MitigatedAtBoundary => 0.05,
        }
    }
}

/// CVSS v4.0 + EPSS Bayesian risk input for an attack transition edge e = (u, v).
#[derive(Clone, Debug)]
pub struct ExploitEdge {
    pub edge_id: &'static str,
    pub cwe_id: u32,
    pub mitre_technique: &'static str,
    pub from: NodeId,
    pub to: NodeId,
    pub precondition: CapabilityMask,
    pub postcondition: CapabilityMask,
    /// CVSS v4.0 Exploitability sub-score in [0.0, 10.0]
    pub cvss_v4_exploitability: f64,
    /// FIRST EPSS 30-day exploitation probability in (0.0, 1.0)
    pub epss_probability: f64,
    /// Compensating control attenuation factor delta in [0.0, 0.99]
    pub mitigation_delta: f64,
    /// Normalized attacker computational/time cost C(e) >= 0.0
    pub attacker_cost: f64,
    pub observation: ProbeObservation,
}

impl ExploitEdge {
    /// Computes Bayesian transition probability p(e | Probe_obs) in [1e-6, 0.9999].
    #[inline]
    pub fn transition_probability(&self) -> f64 {
        let cvss_norm = (self.cvss_v4_exploitability / 10.0).clamp(0.0, 1.0);
        let epss_norm = self.epss_probability.clamp(1e-6, 0.9999);
        let prior = (0.45 * cvss_norm) + (0.55 * epss_norm);
        let delta = self.mitigation_delta.clamp(0.0, 0.99);
        let posterior = prior * self.observation.alpha() * (1.0 - delta);
        posterior.clamp(1e-6, 0.9999)
    }

    /// Non-negative edge weight w(e) = -ln(p(e)) + lambda * C(e) >= 0.
    #[inline]
    pub fn log_weight(&self, cost_lambda: f64) -> f64 {
        let p = self.transition_probability();
        (-p.ln()) + (cost_lambda.max(0.0) * self.attacker_cost.max(0.0))
    }
}

#[derive(Clone, Debug)]
struct SearchState {
    neg_log_cost: f64,
    current_node: NodeId,
    capabilities: CapabilityMask,
    path_edges: Vec<&'static str>,
}

impl PartialEq for SearchState {
    fn eq(&self, other: &Self) -> bool {
        self.neg_log_cost.to_bits() == other.neg_log_cost.to_bits()
            && self.current_node == other.current_node
    }
}

impl Eq for SearchState {}

impl PartialOrd for SearchState {
    fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
        Some(self.cmp(other))
    }
}

impl Ord for SearchState {
    fn cmp(&self, other: &Self) -> Ordering {
        // Min-heap ordering on non-negative log weight
        other
            .neg_log_cost
            .partial_cmp(&self.neg_log_cost)
            .unwrap_or(Ordering::Equal)
    }
}

/// Result of maximum-likelihood exploit path resolution.
#[derive(Clone, Debug)]
pub struct CriticalExploitPath {
    pub ordered_edge_ids: Vec<&'static str>,
    pub total_log_weight: f64,
    pub cumulative_probability: f64,
    pub acquired_capabilities: CapabilityMask,
    pub recommended_cut_edge: &'static str,
}

pub struct AttackSurfaceGraph {
    adjacency: BTreeMap<NodeId, Vec<ExploitEdge>>,
    cost_lambda: f64,
}

impl AttackSurfaceGraph {
    pub fn new(cost_lambda: f64) -> Self {
        Self {
            adjacency: BTreeMap::new(),
            cost_lambda,
        }
    }

    pub fn add_edge(&mut self, edge: ExploitEdge) {
        self.adjacency.entry(edge.from).or_default().push(edge);
    }

    /// Solves for the highest-probability multi-hop exploit path from `source` to `target`
    /// while respecting capability preconditions and avoiding cycles.
    pub fn find_critical_path(
        &self,
        source: NodeId,
        target: NodeId,
        initial_capabilities: CapabilityMask,
        excluded_edges: &BTreeSet<&'static str>,
    ) -> Option<CriticalExploitPath> {
        let mut dist: BTreeMap<NodeId, f64> = BTreeMap::new();
        let mut heap = BinaryHeap::new();

        dist.insert(source, 0.0);
        heap.push(SearchState {
            neg_log_cost: 0.0,
            current_node: source,
            capabilities: initial_capabilities,
            path_edges: Vec::new(),
        });

        while let Some(state) = heap.pop() {
            if state.current_node == target {
                let cumulative_probability = (-state.neg_log_cost).exp().clamp(0.0, 1.0);
                let recommended_cut_edge = state.path_edges.first().copied().unwrap_or("NONE");
                return Some(CriticalExploitPath {
                    ordered_edge_ids: state.path_edges,
                    total_log_weight: state.neg_log_cost,
                    cumulative_probability,
                    acquired_capabilities: state.capabilities,
                    recommended_cut_edge,
                });
            }

            if let Some(&best) = dist.get(&state.current_node) {
                if state.neg_log_cost > best + 1e-12 {
                    continue;
                }
            }

            if let Some(edges) = self.adjacency.get(&state.current_node) {
                for edge in edges {
                    if excluded_edges.contains(edge.edge_id) {
                        continue;
                    }
                    if !state.capabilities.contains(edge.precondition) {
                        continue;
                    }
                    let next_cost = state.neg_log_cost + edge.log_weight(self.cost_lambda);
                    let is_better = match dist.get(&edge.to) {
                        Some(&prev) => next_cost < prev - 1e-12,
                        None => true,
                    };
                    if is_better {
                        dist.insert(edge.to, next_cost);
                        let mut next_path = state.path_edges.clone();
                        next_path.push(edge.edge_id);
                        heap.push(SearchState {
                            neg_log_cost: next_cost,
                            current_node: edge.to,
                            capabilities: state.capabilities.union(edge.postcondition),
                            path_edges: next_path,
                        });
                    }
                }
            }
        }
        None
    }
}