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