from flask import Flask, request, jsonify from flask_cors import CORS import pandas as pd import numpy as np import torch import torch.nn as nn import joblib from collections import Counter import os from pathlib import Path from huggingface_hub import hf_hub_download app = Flask(__name__) CORS(app, origins="*") @app.after_request def after_request(response): response.headers['Access-Control-Allow-Origin'] = '*' response.headers['Access-Control-Allow-Headers'] = 'Content-Type' response.headers['Access-Control-Allow-Methods'] = 'GET, POST, OPTIONS' return response # ══════════════════════════════════════════════════════════════════ # CONFIGURATION # ══════════════════════════════════════════════════════════════════ # Ton repo Hugging Face HF_REPO_ID = "saidimn/ids-cnn-cicids2017" # Dossier local pour stocker les modèles téléchargés CACHE_DIR = Path(__file__).parent / "model_cache" CACHE_DIR.mkdir(exist_ok=True) # ══════════════════════════════════════════════════════════════════ # ARCHITECTURES CNN-1D # ══════════════════════════════════════════════════════════════════ class CNN1D_Binary(nn.Module): def __init__(self, num_features): super().__init__() self.features = nn.Sequential( nn.Conv1d(1, 64, kernel_size=3, padding=1), nn.BatchNorm1d(64), nn.ReLU(), nn.Conv1d(64, 64, kernel_size=3, padding=1), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(2), nn.Dropout(0.2), nn.Conv1d(64, 128, kernel_size=3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, 128, kernel_size=3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2), nn.Dropout(0.3), nn.Conv1d(128, 256, kernel_size=3, padding=1), nn.BatchNorm1d(256), nn.ReLU(), nn.AdaptiveAvgPool1d(1), nn.Dropout(0.3), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, 2) ) def forward(self, x): return self.classifier(self.features(x.unsqueeze(1))) class CNN1D_Attack(nn.Module): def __init__(self, num_features, num_classes): super().__init__() self.features = nn.Sequential( nn.Conv1d(1, 64, kernel_size=3, padding=1), nn.BatchNorm1d(64), nn.ReLU(), nn.Conv1d(64, 64, kernel_size=3, padding=1), nn.BatchNorm1d(64), nn.ReLU(), nn.MaxPool1d(2), nn.Dropout(0.2), nn.Conv1d(64, 128, kernel_size=3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.Conv1d(128, 128, kernel_size=3, padding=1), nn.BatchNorm1d(128), nn.ReLU(), nn.MaxPool1d(2), nn.Dropout(0.3), nn.Conv1d(128, 256, kernel_size=3, padding=1), nn.BatchNorm1d(256), nn.ReLU(), nn.Conv1d(256, 256, kernel_size=3, padding=1), nn.BatchNorm1d(256), nn.ReLU(), nn.AdaptiveAvgPool1d(1), nn.Dropout(0.3), ) self.classifier = nn.Sequential( nn.Flatten(), nn.Linear(256, 256), nn.BatchNorm1d(256), nn.ReLU(), nn.Dropout(0.4), nn.Linear(256, 128), nn.BatchNorm1d(128), nn.ReLU(), nn.Dropout(0.3), nn.Linear(128, num_classes) ) def forward(self, x): return self.classifier(self.features(x.unsqueeze(1))) # ══════════════════════════════════════════════════════════════════ # TÉLÉCHARGEMENT DES MODÈLES DEPUIS HUGGING FACE # ══════════════════════════════════════════════════════════════════ def download_models(): """Télécharge les modèles depuis Hugging Face Hub""" files = { "binary": "cnn1d_binary.pth", "attack": "cnn1d_attacks_only.pth", "scaler": "scaler.pkl", "encoder": "label_encoder_attacks.pkl" } paths = {} print("=" * 50) print("Téléchargement des modèles depuis Hugging Face...") print("=" * 50) for key, filename in files.items(): print(" ↓ " + filename) paths[key] = hf_hub_download( repo_id=HF_REPO_ID, filename=filename, cache_dir=CACHE_DIR, local_dir=CACHE_DIR, local_dir_use_symlinks=False ) print(" ✓ " + paths[key]) return paths # Télécharge au démarrage du serveur paths = download_models() print("=" * 50) # ══════════════════════════════════════════════════════════════════ # CHARGEMENT DES MODÈLES EN MÉMOIRE # ══════════════════════════════════════════════════════════════════ device = torch.device('cuda' if torch.cuda.is_available() else 'cpu') print("Device: " + str(device)) scaler = joblib.load(paths["scaler"]) le = joblib.load(paths["encoder"]) num_features = scaler.n_features_in_ num_attack_classes = len(le.classes_) print("Features: " + str(num_features)) print("Classes: " + str(list(le.classes_))) # Modèle binaire binary_model = CNN1D_Binary(num_features).to(device) binary_model.load_state_dict(torch.load(paths["binary"], map_location=device, weights_only=True)) binary_model.eval() # Modèle d'attaque attack_model = CNN1D_Attack(num_features, num_attack_classes).to(device) attack_model.load_state_dict(torch.load(paths["attack"], map_location=device, weights_only=True)) attack_model.eval() print("Tous les modèles sont chargés ✓\n") # ══════════════════════════════════════════════════════════════════ # PRÉTRAITEMENT # ══════════════════════════════════════════════════════════════════ def preprocess(df): df.columns = df.columns.str.strip() cols_to_drop = ['Flow ID', 'Src IP', 'Src Port', 'Dst IP', 'Dst Port', 'Protocol', 'Timestamp', 'Label'] for col in cols_to_drop: if col in df.columns: df = df.drop(columns=[col]) rename_dict = { 'Tot Fwd Pkts': 'Total Fwd Packets', 'Tot Bwd Pkts': 'Total Backward Packets', 'TotLen Fwd Pkts': 'Total Length of Fwd Packets', 'TotLen Bwd Pkts': 'Total Length of Bwd Packets', 'Fwd Pkt Len Max': 'Fwd Packet Length Max', 'Fwd Pkt Len Min': 'Fwd Packet Length Min', 'Fwd Pkt Len Mean': 'Fwd Packet Length Mean', 'Fwd Pkt Len Std': 'Fwd Packet Length Std', 'Bwd Pkt Len Max': 'Bwd Packet Length Max', 'Fwd Header Len': 'Fwd Header Length', 'Bwd Header Len': 'Bwd Header Length', 'Fwd Pkts/s': 'Fwd Packets/s', 'Bwd Pkts/s': 'Bwd Packets/s', 'Pkt Len Min': 'Min Packet Length', 'Pkt Len Max': 'Max Packet Length', 'Pkt Len Mean': 'Packet Length Mean', 'Pkt Len Std': 'Packet Length Std', 'Pkt Len Var': 'Packet Length Variance', 'FIN Flag Cnt': 'FIN Flag Count', 'SYN Flag Cnt': 'SYN Flag Count', 'RST Flag Cnt': 'RST Flag Count', 'PSH Flag Cnt': 'PSH Flag Count', 'ACK Flag Cnt': 'ACK Flag Count', 'URG Flag Cnt': 'URG Flag Count', 'Pkt Size Avg': 'Average Packet Size', 'Fwd Seg Size Avg': 'Avg Fwd Segment Size', 'Bwd Seg Size Avg': 'Avg Bwd Segment Size', 'Fwd Byts/b Avg': 'Fwd Avg Bytes/Bulk', 'Fwd Pkts/b Avg': 'Fwd Avg Packets/Bulk', 'Fwd Blk Rate Avg': 'Fwd Avg Bulk Rate', 'Bwd Byts/b Avg': 'Bwd Avg Bytes/Bulk', 'Bwd Pkts/b Avg': 'Bwd Avg Packets/Bulk', 'Bwd Blk Rate Avg': 'Bwd Avg Bulk Rate', 'Subflow Fwd Pkts': 'Subflow Fwd Packets', 'Subflow Bwd Pkts': 'Subflow Bwd Packets', 'Init Fwd Win Byts': 'Init_Win_bytes_forward', 'Init Bwd Win Byts': 'Init_Win_bytes_backward', 'Fwd Act Data Pkts': 'act_data_pkt_fwd', 'Fwd Seg Size Min': 'min_seg_size_forward', } df = df.rename(columns=rename_dict) df = df.select_dtypes(include=[np.number]) df.replace([np.inf, -np.inf], np.nan, inplace=True) df.fillna(0, inplace=True) if hasattr(scaler, 'feature_names_in_'): for col in scaler.feature_names_in_: if col not in df.columns: df[col] = 0 df = df[scaler.feature_names_in_] else: while df.shape[1] < 78: df['missing_' + str(df.shape[1])] = 0 df = df.iloc[:, :78] return scaler.transform(df.values) # ══════════════════════════════════════════════════════════════════ # ROUTES API # ══════════════════════════════════════════════════════════════════ @app.route('/analyze', methods=['POST']) def analyze(): if 'file' not in request.files: return jsonify({'error': 'No file uploaded'}), 400 try: file = request.files['file'] df = pd.read_csv(file) if df.empty: return jsonify({'error': 'CSV file is empty'}), 400 total_flows = len(df) X_scaled = preprocess(df) X = torch.tensor(X_scaled, dtype=torch.float32).to(device) results = [] with torch.no_grad(): binary_out = binary_model(X) binary_pred = torch.argmax(binary_out, dim=1) for i in range(len(X)): if binary_pred[i] == 0: results.append('BENIGN') else: single = X[i].unsqueeze(0) attack_out = attack_model(single) attack_pred = torch.argmax(attack_out, dim=1).item() results.append(le.classes_[attack_pred]) counts = Counter(results) total = len(results) labels = list(counts.keys()) values = list(counts.values()) percentages = [round(v/total*100, 2) for v in values] attacks = {k: v for k, v in counts.items() if k != 'BENIGN'} benign = counts.get('BENIGN', 0) return jsonify({ 'total_flows': total, 'benign_count': benign, 'attack_count': total - benign, 'labels': labels, 'values': values, 'percentages': percentages, 'attack_types': attacks, 'results': results[:100] }) except Exception as e: return jsonify({'error': str(e)}), 500 @app.route('/health', methods=['GET']) def health(): return jsonify({ 'status': 'ok', 'device': str(device), 'repo': HF_REPO_ID, 'attack_classes': le.classes_.tolist() }) if __name__ == '__main__': import os port = int(os.environ.get('PORT', 5000)) app.run(debug=False, port=port, host='0.0.0.0')