机器学习在安全领域的真实用法

import pyshark
import time
import json

class NetworkTrafficCollector:
    def __init__(self, interface, output_file):
        self.interface = interface
        self.output_file = output_file
        self.capture = None
    
    def start_capture(self):
        """开始捕获网络流量"""
        self.capture = pyshark.LiveCapture(interface=self.interface)
        self.capture.sniff(timeout=60)  # 捕获60秒
        
        # 处理捕获的数据包
        for packet in self.capture:
            traffic_data = self._extract_traffic_features(packet)
            self._save_traffic_data(traffic_data)
    
    def _extract_traffic_features(self, packet):
        """提取流量特征"""
        features = {}
        
        # 基本信息
        features['timestamp'] = time.time()
        features['protocol'] = packet.transport_layer if hasattr(packet, 'transport_layer') else 'unknown'
        
        # 网络层信息
        if hasattr(packet, 'ip'):
            features['src_ip'] = packet.ip.src
            features['dst_ip'] = packet.ip.dst
            features['ip_len'] = int(packet.ip.len)
        
        # 传输层信息
        if hasattr(packet, 'tcp'):
            features['src_port'] = int(packet.tcp.srcport)
            features['dst_port'] = int(packet.tcp.dstport)
            features['tcp_flags'] = packet.tcp.flags
            features['tcp_window_size'] = int(packet.tcp.window_size)
        elif hasattr(packet, 'udp'):
            features['src_port'] = int(packet.udp.srcport)
            features['dst_port'] = int(packet.udp.dstport)
            features['udp_length'] = int(packet.udp.length)
        
        # 应用层信息
        if hasattr(packet, 'http'):
            features['http_method'] = packet.http.request_method if hasattr(packet.http, 'request_method') else None
            features['http_uri'] = packet.http.request_uri if hasattr(packet.http, 'request_uri') else None
            features['http_status'] = packet.http.response_code if hasattr(packet.http, 'response_code') else None
        
        return features
    
    def _save_traffic_data(self, data):
        """保存流量数据到文件"""
        with open(self.output_file, 'a') as f:
            json.dump(data, f)
            f.write('\n')

# 使用示例
if __name__ == '__main__':
    collector = NetworkTrafficCollector(interface='eth0', output_file='network_traffic.json')
    collector.start_capture()

import pandas as pd
import numpy as np
from sklearn.preprocessing import StandardScaler, LabelEncoder

class FeatureExtractor:
    def __init__(self):
        self.scaler = StandardScaler()
        self.label_encoder = LabelEncoder()
    
    def load_data(self, file_path):
        """加载数据"""
        return pd.read_json(file_path, lines=True)
    
    def preprocess_data(self, df):
        """预处理数据"""
        # 处理缺失值
        df = df.dropna()
        
        # 时间特征工程
        df['hour'] = pd.to_datetime(df['timestamp'], unit='s').dt.hour
        df['day_of_week'] = pd.to_datetime(df['timestamp'], unit='s').dt.dayofweek
        
        # 分类特征编码
        categorical_features = ['protocol', 'tcp_flags', 'http_method']
        for feature in categorical_features:
            if feature in df.columns:
                df[feature] = self.label_encoder.fit_transform(df[feature].astype(str))
        
        # 数值特征标准化
        numerical_features = ['ip_len', 'tcp_window_size', 'udp_length']
        for feature in numerical_features:
            if feature in df.columns:
                df[feature] = self.scaler.fit_transform(df[[feature]])
        
        return df
    
    def extract_features(self, df):
        """提取特征"""
        # 选择特征列
        feature_columns = [col for col in df.columns if col != 'timestamp']
        
        # 提取特征
        X = df[feature_columns].values
        
        return X

# 使用示例
if __name__ == '__main__':
    extractor = FeatureExtractor()
    df = extractor.load_data('network_traffic.json')
    df_processed = extractor.preprocess_data(df)
    features = extractor.extract_features(df_processed)
    print(f"提取的特征数量: {features.shape[1]}")
    print(f"特征示例: {features[0]}")

from sklearn.ensemble import IsolationForest
from sklearn.model_selection import train_test_split
import numpy as np
import joblib

class ThreatDetector:
    def __init__(self):
        self.model = IsolationForest(n_estimators=100, contamination=0.1, random_state=42)
    
    def train(self, X):
        """训练模型"""
        self.model.fit(X)
    
    def predict(self, X):
        """预测异常"""
        # Isolation Forest返回-1表示异常，1表示正常
        predictions = self.model.predict(X)
        # 转换为0表示正常，1表示异常
        anomalies = np.where(predictions == -1, 1, 0)
        return anomalies
    
    def get_anomaly_scores(self, X):
        """获取异常分数"""
        return self.model.decision_function(X)
    
    def save_model(self, file_path):
        """保存模型"""
        joblib.dump(self.model, file_path)
    
    def load_model(self, file_path):
        """加载模型"""
        self.model = joblib.load(file_path)

# 使用示例
if __name__ == '__main__':
    # 假设features是从FeatureExtractor提取的特征
    # 这里使用随机数据模拟
    np.random.seed(42)
    normal_data = np.random.randn(1000, 10)
    anomalous_data = np.random.randn(100, 10) * 2 + 5  # 异常数据
    X = np.vstack([normal_data, anomalous_data])
    
    # 划分训练集和测试集
    X_train, X_test = train_test_split(X, test_size=0.2, random_state=42)
    
    # 训练模型
    detector = ThreatDetector()
    detector.train(X_train)
    
    # 预测测试集
    predictions = detector.predict(X_test)
    anomaly_scores = detector.get_anomaly_scores(X_test)
    
    # 保存模型
    detector.save_model('isolation_forest_model.joblib')
    
    print(f"测试集大小: {X_test.shape[0]}")
    print(f"检测到的异常数量: {np.sum(predictions)}")
    print(f"异常分数范围: [{np.min(anomaly_scores):.4f}, {np.max(anomaly_scores):.4f}]")

import numpy as np
from sklearn.feature_selection import mutual_info_classif
from sklearn.preprocessing import StandardScaler

class DynamicFeatureLearner:
    def __init__(self, window_size=1000, threshold=0.1):
        self.window_size = window_size
        self.threshold = threshold
        self.scaler = StandardScaler()
        self.features_history = []
    
    def update_features(self, new_features, labels=None):
        """更新特征"""
        # 标准化新特征
        new_features_scaled = self.scaler.fit_transform(new_features)
        
        # 添加到历史特征
        self.features_history.append(new_features_scaled)
        
        # 保持窗口大小
        if len(self.features_history) > self.window_size:
            self.features_history.pop(0)
        
        # 合并历史特征
        all_features = np.vstack(self.features_history)
        
        # 如果有标签，计算特征重要性
        if labels is not None and len(labels) > 0:
            feature_importance = self._calculate_feature_importance(all_features, labels)
            return self._select_important_features(all_features, feature_importance)
        
        return all_features
    
    def _calculate_feature_importance(self, X, y):
        """计算特征重要性"""
        return mutual_info_classif(X, y, random_state=42)
    
    def _select_important_features(self, X, feature_importance):
        """选择重要特征"""
        important_indices = np.where(feature_importance > self.threshold)[0]
        return X[:, important_indices]
    
    def get_current_features(self):
        """获取当前特征"""
        if len(self.features_history) == 0:
            return None
        return np.vstack(self.features_history)

# 使用示例
if __name__ == '__main__':
    # 生成模拟数据
    np.random.seed(42)
    X = np.random.randn(2000, 20)
    y = np.random.randint(0, 2, 2000)
    
    # 初始化动态特征学习器
    feature_learner = DynamicFeatureLearner(window_size=1000, threshold=0.05)
    
    # 分批次更新特征
    for i in range(0, 2000, 100):
        X_batch = X[i:i+100]
        y_batch = y[i:i+100]
        updated_features = feature_learner.update_features(X_batch, y_batch)
        print(f"批次 {i//100 + 1}: 特征数量从 {X_batch.shape[1]} 减少到 {updated_features.shape[1]}")
    
    current_features = feature_learner.get_current_features()
    print(f"当前特征总数: {current_features.shape[0]}")
    print(f"当前特征维度: {current_features.shape[1]}")

import numpy as np
from sklearn.metrics import precision_recall_curve, f1_score

class AdaptiveThresholdAdjuster:
    def __init__(self, initial_threshold=0.5, window_size=1000, alpha=0.1):
        self.current_threshold = initial_threshold
        self.window_size = window_size
        self.alpha = alpha  # 平滑系数
        self.predictions_history = []
        self.labels_history = []
    
    def update_threshold(self, predictions, labels):
        """更新阈值"""
        # 添加到历史记录
        self.predictions_history.extend(predictions)
        self.labels_history.extend(labels)
        
        # 保持窗口大小
        if len(self.predictions_history) > self.window_size:
            self.predictions_history = self.predictions_history[-self.window_size:]
            self.labels_history = self.labels_history[-self.window_size:]
        
        # 计算最佳阈值
        best_threshold = self._find_best_threshold(
            np.array(self.predictions_history), 
            np.array(self.labels_history)
        )
        
        # 平滑更新阈值
        self.current_threshold = (1 - self.alpha) * self.current_threshold + self.alpha * best_threshold
        
        return self.current_threshold
    
    def _find_best_threshold(self, scores, labels):
        """寻找最佳阈值"""
        if len(np.unique(labels)) < 2:
            return self.current_threshold
        
        # 计算精确率-召回率曲线
        precision, recall, thresholds = precision_recall_curve(labels, scores)
        
        # 计算F1分数
        f1_scores = 2 * (precision * recall) / (precision + recall)
        f1_scores = f1_scores[:-1]  # 移除最后一个无效值
        
        # 寻找最佳阈值
        best_idx = np.argmax(f1_scores)
        best_threshold = thresholds[best_idx]
        
        return best_threshold
    
    def get_current_threshold(self):
        """获取当前阈值"""
        return self.current_threshold

# 使用示例
if __name__ == '__main__':
    # 生成模拟数据
    np.random.seed(42)
    normal_scores = np.random.normal(0, 1, 500)
    anomalous_scores = np.random.normal(3, 1, 500)
    scores = np.concatenate([normal_scores, anomalous_scores])
    labels = np.concatenate([np.zeros(500), np.ones(500)])
    
    # 初始化自适应阈值调整器
    threshold_adjuster = AdaptiveThresholdAdjuster(initial_threshold=0.5, window_size=500, alpha=0.1)
    
    # 分批次更新阈值
    for i in range(0, 1000, 100):
        scores_batch = scores[i:i+100]
        labels_batch = labels[i:i+100]
        new_threshold = threshold_adjuster.update_threshold(scores_batch, labels_batch)
        print(f"批次 {i//100 + 1}: 阈值调整为 {new_threshold:.4f}")
    
    current_threshold = threshold_adjuster.get_current_threshold()
    print(f"最终阈值: {current_threshold:.4f}")
    
    # 使用最终阈值进行预测
    final_predictions = (scores > current_threshold).astype(int)
    final_f1 = f1_score(labels, final_predictions)
    print(f"使用最终阈值的F1分数: {final_f1:.4f}")

import numpy as np
import tensorflow as tf
from tensorflow.keras.layers import Input, Dense, Attention, concatenate, Softmax
from tensorflow.keras.models import Model

class MultimodalFusion:
    def __init__(self, input_dims, hidden_dims):
        self.input_dims = input_dims
        self.hidden_dims = hidden_dims
        self.model = self._build_model()
    
    def _build_model(self):
        """构建多模态融合模型"""
        # 输入层
        inputs = [Input(shape=(dim,)) for dim in self.input_dims]
        
        # 模态嵌入层
        embeddings = []
        for i, input_ in enumerate(inputs):
            embedding = Dense(self.hidden_dims, activation='relu', name=f'modal_embedding_{i}')(input_)
            embeddings.append(embedding)
        
        # 注意力机制
        attention_input = concatenate(embeddings, axis=1)
        attention_scores = Dense(len(self.input_dims), activation='softmax', name='attention_scores')(attention_input)
        
        # 加权融合
        weighted_embeddings = []
        for i, embedding in enumerate(embeddings):
            weighted = tf.multiply(embedding, attention_scores[:, i:i+1])
            weighted_embeddings.append(weighted)
        
        fusion_output = concatenate(weighted_embeddings, axis=1)
        
        # 输出层
        output = Dense(1, activation='sigmoid', name='fusion_output')(fusion_output)
        
        # 构建模型
        model = Model(inputs=inputs, outputs=[output, attention_scores])
        
        # 编译模型
        model.compile(optimizer='adam', loss='binary_crossentropy', metrics=['accuracy'])
        
        return model
    
    def train(self, X_list, y, epochs=10, batch_size=32, validation_split=0.2):
        """训练模型"""
        history = self.model.fit(X_list, y, epochs=epochs, batch_size=batch_size, validation_split=validation_split)
        return history
    
    def predict(self, X_list):
        """预测"""
        return self.model.predict(X_list)
    
    def save_model(self, file_path):
        """保存模型"""
        self.model.save(file_path)
    
    def load_model(self, file_path):
        """加载模型"""
        self.model = tf.keras.models.load_model(file_path)

# 使用示例
if __name__ == '__main__':
    # 生成模拟数据
    np.random.seed(42)
    
    # 三种模态数据
    X1 = np.random.randn(1000, 10)  # 网络流量特征
    X2 = np.random.randn(1000, 8)   # 日志特征
    X3 = np.random.randn(1000, 12)  # 终端特征
    
    # 标签
    y = np.random.randint(0, 2, 1000)
    
    # 初始化多模态融合模型
    fusion_model = MultimodalFusion(input_dims=[10, 8, 12], hidden_dims=20)
    
    # 训练模型
    history = fusion_model.train([X1, X2, X3], y, epochs=5, batch_size=32)
    
    # 预测
    predictions, attention_scores = fusion_model.predict([X1, X2, X3])
    
    # 保存模型
    fusion_model.save_model('multimodal_fusion_model.h5')
    
    print(f"预测结果形状: {predictions.shape}")
    print(f"注意力分数形状: {attention_scores.shape}")
    print(f"样本注意力分数示例: {attention_scores[0]}")
    print(f"样本预测结果示例: {predictions[0][0]:.4f}")