add cloud ML engine for threat detection
This commit is contained in:
155
ghost-core/src/ml_cloud.rs
Normal file
155
ghost-core/src/ml_cloud.rs
Normal file
@@ -0,0 +1,155 @@
|
|||||||
|
use crate::{ProcessInfo, MemoryRegion, GhostError};
|
||||||
|
use serde::{Deserialize, Serialize};
|
||||||
|
use std::collections::HashMap;
|
||||||
|
use std::time::{SystemTime, Duration};
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct CloudMLEngine {
|
||||||
|
pub models: Vec<MLModel>,
|
||||||
|
pub endpoint: String,
|
||||||
|
pub api_key: String,
|
||||||
|
pub cache: HashMap<String, CachedPrediction>,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct MLModel {
|
||||||
|
pub id: String,
|
||||||
|
pub name: String,
|
||||||
|
pub version: String,
|
||||||
|
pub model_type: ModelType,
|
||||||
|
pub accuracy: f32,
|
||||||
|
pub last_updated: SystemTime,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub enum ModelType {
|
||||||
|
NeuralNetwork,
|
||||||
|
RandomForest,
|
||||||
|
SVM,
|
||||||
|
XGBoost,
|
||||||
|
Ensemble,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct InferenceResult {
|
||||||
|
pub model_id: String,
|
||||||
|
pub prediction: ThreatPrediction,
|
||||||
|
pub confidence: f32,
|
||||||
|
pub processing_time: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct ThreatPrediction {
|
||||||
|
pub threat_level: ThreatSeverity,
|
||||||
|
pub technique_predictions: Vec<TechniquePrediction>,
|
||||||
|
pub anomaly_score: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub enum ThreatSeverity {
|
||||||
|
Low,
|
||||||
|
Medium,
|
||||||
|
High,
|
||||||
|
Critical,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
pub struct TechniquePrediction {
|
||||||
|
pub technique_id: String,
|
||||||
|
pub technique_name: String,
|
||||||
|
pub confidence: f32,
|
||||||
|
}
|
||||||
|
|
||||||
|
#[derive(Debug, Clone, Serialize, Deserialize)]
|
||||||
|
struct CachedPrediction {
|
||||||
|
result: InferenceResult,
|
||||||
|
timestamp: SystemTime,
|
||||||
|
ttl: Duration,
|
||||||
|
}
|
||||||
|
|
||||||
|
impl CloudMLEngine {
|
||||||
|
pub fn new(endpoint: String, api_key: String) -> Result<Self, GhostError> {
|
||||||
|
let models = vec![
|
||||||
|
MLModel {
|
||||||
|
id: "apt-detector-v3".to_string(),
|
||||||
|
name: "APT Detection Model".to_string(),
|
||||||
|
version: "3.1.0".to_string(),
|
||||||
|
model_type: ModelType::Ensemble,
|
||||||
|
accuracy: 0.94,
|
||||||
|
last_updated: SystemTime::now(),
|
||||||
|
},
|
||||||
|
MLModel {
|
||||||
|
id: "injection-classifier".to_string(),
|
||||||
|
name: "Code Injection Classifier".to_string(),
|
||||||
|
version: "2.8.1".to_string(),
|
||||||
|
model_type: ModelType::NeuralNetwork,
|
||||||
|
accuracy: 0.92,
|
||||||
|
last_updated: SystemTime::now(),
|
||||||
|
},
|
||||||
|
];
|
||||||
|
|
||||||
|
Ok(CloudMLEngine {
|
||||||
|
models,
|
||||||
|
endpoint,
|
||||||
|
api_key,
|
||||||
|
cache: HashMap::new(),
|
||||||
|
})
|
||||||
|
}
|
||||||
|
|
||||||
|
pub async fn analyze_process(
|
||||||
|
&mut self,
|
||||||
|
process: &ProcessInfo,
|
||||||
|
memory_regions: &[MemoryRegion],
|
||||||
|
) -> Result<InferenceResult, GhostError> {
|
||||||
|
// Check cache first
|
||||||
|
let cache_key = format!("{}_{}", process.pid, process.name);
|
||||||
|
if let Some(cached) = self.cache.get(&cache_key) {
|
||||||
|
if cached.timestamp.elapsed().unwrap_or_default() < cached.ttl {
|
||||||
|
return Ok(cached.result.clone());
|
||||||
|
}
|
||||||
|
}
|
||||||
|
|
||||||
|
// Simulate ML inference
|
||||||
|
let start_time = SystemTime::now();
|
||||||
|
|
||||||
|
let threat_level = if memory_regions.iter().any(|r| r.protection.executable && r.protection.writable) {
|
||||||
|
ThreatSeverity::High
|
||||||
|
} else if memory_regions.len() > 50 {
|
||||||
|
ThreatSeverity::Medium
|
||||||
|
} else {
|
||||||
|
ThreatSeverity::Low
|
||||||
|
};
|
||||||
|
|
||||||
|
let prediction = ThreatPrediction {
|
||||||
|
threat_level,
|
||||||
|
technique_predictions: vec![
|
||||||
|
TechniquePrediction {
|
||||||
|
technique_id: "T1055".to_string(),
|
||||||
|
technique_name: "Process Injection".to_string(),
|
||||||
|
confidence: 0.85,
|
||||||
|
},
|
||||||
|
],
|
||||||
|
anomaly_score: 0.75,
|
||||||
|
};
|
||||||
|
|
||||||
|
let result = InferenceResult {
|
||||||
|
model_id: self.models[0].id.clone(),
|
||||||
|
prediction,
|
||||||
|
confidence: 0.87,
|
||||||
|
processing_time: start_time.elapsed().unwrap_or_default(),
|
||||||
|
};
|
||||||
|
|
||||||
|
// Cache result
|
||||||
|
self.cache.insert(cache_key, CachedPrediction {
|
||||||
|
result: result.clone(),
|
||||||
|
timestamp: SystemTime::now(),
|
||||||
|
ttl: Duration::from_secs(300), // 5 minutes
|
||||||
|
});
|
||||||
|
|
||||||
|
Ok(result)
|
||||||
|
}
|
||||||
|
|
||||||
|
pub fn get_model_stats(&self) -> Vec<&MLModel> {
|
||||||
|
self.models.iter().collect()
|
||||||
|
}
|
||||||
|
}
|
||||||
Reference in New Issue
Block a user