Skip to content

N
NIDS-Encrypted-Traffic-Anomaly-Detection
Commit 532029c9 authored by rawan.salameh's avatar rawan.salameh
Browse files
Options

Add full project: models, system code, and report

parent 25e0a018
Hide whitespace changes
Inline Side-by-side
Showing with 4045 additions and 0 deletions
models/CNN_AE.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "markdown",
"id": "8da53469",
"metadata": {},
"source": [
"# Architecture Search — CNN AutoEncoder Baseline\n",
"\n",
"## 1. Load Preprocessed Data"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "3ecb7c74",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"=== بيانات التدريب (normal) ===\n",
" X_norm_scaled: (7295027, 22)\n",
" ppi_norm: (7295027, 30, 3)\n",
"\n",
"=== بيانات الاختبار (TrickBot) ===\n",
" X_mal_scaled: (10823, 22)\n",
" ppi_mal_normalized: (10823, 30, 3)\n",
"Device: cuda\n"
]
}
],
"source": [
"import numpy as np\n",
"import joblib\n",
"import torch\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"data_dir = r'D:\\project 5\\data'\n",
"\n",
"X_norm_scaled = np.load(data_dir + r'\\normal_features_22_40k.npy')\n",
"ppi_norm = np.load(data_dir + r'\\normal_ppi_sequences_40k.npy')\n",
"\n",
"X_mal_scaled = np.load(data_dir + r'\\trickbot_features_FINAL.npy')\n",
"ppi_mal_normalized = np.load(data_dir + r'\\trickbot_ppi_FINAL.npy')\n",
"\n",
"scaler = joblib.load(data_dir + r'\\flow_scaler_40k.pkl')\n",
"ppi_norm_stats = joblib.load(data_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"\n",
"print(\"=== بيانات التدريب (normal) ===\")\n",
"print(f\" X_norm_scaled: {X_norm_scaled.shape}\")\n",
"print(f\" ppi_norm: {ppi_norm.shape}\")\n",
"print(\"\\n=== بيانات الاختبار (TrickBot) ===\")\n",
"print(f\" X_mal_scaled: {X_mal_scaled.shape}\")\n",
"print(f\" ppi_mal_normalized: {ppi_mal_normalized.shape}\")\n",
"\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"print(f\"Device: {device}\")"
]
},
{
"cell_type": "markdown",
"id": "c70de283",
"metadata": {},
"source": [
"## 3. Balanced Train/Val/Test Split"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "709fe18f",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import numpy as np\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"from sklearn.model_selection import train_test_split\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"SEED = 42\n",
"torch.manual_seed(SEED); np.random.seed(SEED)\n",
"\n",
"N_MAL = len(ppi_mal_normalized)\n",
"print(f\"عدد الـ malware: {N_MAL}\")\n",
"\n",
"ppi_norm_cnn = np.transpose(ppi_norm, (0, 2, 1))\n",
"ppi_mal_cnn = np.transpose(ppi_mal_normalized, (0, 2, 1))\n",
"\n",
"rng = np.random.RandomState(SEED)\n",
"all_idx = np.arange(len(ppi_norm_cnn))\n",
"test_norm_idx = rng.choice(all_idx, size=N_MAL, replace=False)\n",
"train_pool_idx = np.setdiff1d(all_idx, test_norm_idx)\n",
"\n",
"S_test_norm = ppi_norm_cnn[test_norm_idx]\n",
"S_pool = ppi_norm_cnn[train_pool_idx]\n",
"\n",
"S_train, S_val = train_test_split(S_pool, test_size=0.1, random_state=SEED)\n",
"\n",
"print(f\"Train: {S_train.shape}\")\n",
"print(f\"Val: {S_val.shape}\")\n",
"print(f\"Test normal: {S_test_norm.shape} | Test malware: {ppi_mal_cnn.shape}\")\n",
"\n",
"S_train_t = torch.tensor(S_train, dtype=torch.float32)\n",
"S_val_t = torch.tensor(S_val, dtype=torch.float32)\n",
"\n",
"train_loader_cnn = DataLoader(TensorDataset(S_train_t), batch_size=512, shuffle=True)\n",
"val_loader_cnn = DataLoader(TensorDataset(S_val_t), batch_size=512, shuffle=False)\n",
"\n",
"y_true_test = np.concatenate([np.zeros(N_MAL), np.ones(N_MAL)])\n",
"print(f\"\\nTest متوازن: {N_MAL} normal + {N_MAL} malware = {2*N_MAL} صف\")"
]
},
{
"cell_type": "markdown",
"id": "0bb5f289",
"metadata": {},
"source": [
"## 4. Model Definition, Training & Evaluation — ImprovedCNNAutoEncoder"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "c0c93861",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import numpy as np\n",
"from sklearn.metrics import roc_auc_score, precision_score, recall_score, f1_score\n",
"\n",
"class ImprovedCNNAutoEncoder(nn.Module):\n",
" def __init__(self, in_ch=3, seq_len=30, latent_dim=16, dropout_p=0.2):\n",
" super().__init__()\n",
" # Encoder\n",
" self.enc_conv1 = nn.Conv1d(in_ch, 16, 3, padding=1)\n",
" self.enc_bn1 = nn.BatchNorm1d(16); self.enc_pool1 = nn.MaxPool1d(2)\n",
" self.enc_conv2 = nn.Conv1d(16, 32, 3, padding=1)\n",
" self.enc_bn2 = nn.BatchNorm1d(32); self.enc_pool2 = nn.MaxPool1d(2)\n",
" self.enc_conv3 = nn.Conv1d(32, 16, 3, padding=1)\n",
" self.enc_bn3 = nn.BatchNorm1d(16)\n",
" self.relu = nn.ReLU()\n",
" # Bottleneck\n",
" self.enc_dropout = nn.Dropout(dropout_p)\n",
" self.to_latent = nn.Linear(16 * 7, latent_dim)\n",
" self.from_latent = nn.Linear(latent_dim, 16 * 7)\n",
" self.dec_dropout = nn.Dropout(dropout_p)\n",
" # Decoder\n",
" self.dec_conv1 = nn.Conv1d(16, 32, 3, padding=1)\n",
" self.dec_bn1 = nn.BatchNorm1d(32); self.dec_up1 = nn.Upsample(size=15)\n",
" self.dec_conv2 = nn.Conv1d(32, 16, 3, padding=1)\n",
" self.dec_bn2 = nn.BatchNorm1d(16); self.dec_up2 = nn.Upsample(size=30)\n",
" self.dec_out = nn.Conv1d(16, in_ch, 3, padding=1)\n",
"\n",
" def forward(self, x):\n",
" x = self.enc_pool1(self.relu(self.enc_bn1(self.enc_conv1(x)))) # (16,15)\n",
" x = self.enc_pool2(self.relu(self.enc_bn2(self.enc_conv2(x)))) # (32,7)\n",
" x = self.relu(self.enc_bn3(self.enc_conv3(x))) # (16,7)\n",
" x = x.flatten(1)\n",
" z = self.to_latent(self.enc_dropout(x))\n",
" z = self.dec_dropout(self.from_latent(z)).view(-1, 16, 7)\n",
" x = self.dec_up1(self.relu(self.dec_bn1(self.dec_conv1(z)))) # (32,15)\n",
" x = self.dec_up2(self.relu(self.dec_bn2(self.dec_conv2(x)))) # (16,30)\n",
" return self.dec_out(x) # (3,30)\n",
"\n",
"cnn_ae = ImprovedCNNAutoEncoder().to(device)\n",
"print(f\"CNN AE parameters: {sum(p.numel() for p in cnn_ae.parameters()):,}\")\n",
"\n",
"criterion = nn.MSELoss()\n",
"optimizer = torch.optim.AdamW(cnn_ae.parameters(), lr=1e-3, weight_decay=1e-4)\n",
"scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)\n",
"\n",
"def recon_err_cnn(model, X, bs=4096):\n",
" model.eval(); errs = []\n",
" with torch.no_grad():\n",
" for i in range(0, len(X), bs):\n",
" xb = torch.tensor(X[i:i+bs], dtype=torch.float32, device=device)\n",
" out = model(xb)\n",
" errs.append(((out - xb) ** 2).mean(dim=(1,2)).cpu().numpy())\n",
" return np.concatenate(errs)\n",
"\n",
"MAX_EPOCHS, WARMUP, PATIENCE, BASE_LR = 100, 3, 5, 1e-3\n",
"best_val, wait = float('inf'), 0\n",
"out_dir = r'D:\\project 5\\data'\n",
"\n",
"for epoch in range(MAX_EPOCHS):\n",
" if epoch < WARMUP:\n",
" for g in optimizer.param_groups:\n",
" g['lr'] = BASE_LR * (epoch + 1) / WARMUP\n",
"\n",
" cnn_ae.train()\n",
" tr_loss = 0.0\n",
" for (xb,) in train_loader_cnn:\n",
" xb = xb.to(device)\n",
" optimizer.zero_grad()\n",
" loss = criterion(cnn_ae(xb), xb)\n",
" loss.backward()\n",
" optimizer.step()\n",
" tr_loss += loss.item() * xb.size(0)\n",
" tr_loss /= len(train_loader_cnn.dataset)\n",
"\n",
" cnn_ae.eval()\n",
" val_loss = 0.0\n",
" with torch.no_grad():\n",
" for (xb,) in val_loader_cnn:\n",
" xb = xb.to(device)\n",
" val_loss += criterion(cnn_ae(xb), xb).item() * xb.size(0)\n",
" val_loss /= len(val_loader_cnn.dataset)\n",
"\n",
" if epoch >= WARMUP:\n",
" scheduler.step(val_loss)\n",
"\n",
" if (epoch+1) % 5 == 0 or epoch < 5:\n",
" e_norm = recon_err_cnn(cnn_ae, S_test_norm)\n",
" e_mal = recon_err_cnn(cnn_ae, ppi_mal_cnn)\n",
" auc = roc_auc_score(y_true_test, np.concatenate([e_norm, e_mal]))\n",
" print(f\"Epoch {epoch+1:3d} | train {tr_loss:.5f} | val {val_loss:.5f} | ROC-AUC {auc:.4f}\")\n",
"\n",
" if val_loss < best_val:\n",
" best_val, wait = val_loss, 0\n",
" torch.save(cnn_ae.state_dict(), out_dir + r'\\cnn_ae_40k_best.pt')\n",
" else:\n",
" if epoch >= WARMUP:\n",
" wait += 1\n",
" if wait >= PATIENCE:\n",
" print(f\"Early stopping at epoch {epoch+1}\")\n",
" break\n",
"\n",
"cnn_ae.load_state_dict(torch.load(out_dir + r'\\cnn_ae_40k_best.pt'))\n"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "63ad6dbe",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "20abce4a",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "6759c8d9",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "becf15e7",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "41307980",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
models/CNN_Conditioned.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "markdown",
"id": "db7f8fd2",
"metadata": {},
"source": [
"# Architecture Search — CNN Conditioned on Packet Count\n",
"\n",
"## 1. Load Preprocessed Data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ecb7c74",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import joblib\n",
"import torch\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"data_dir = r'D:\\project 5\\data'\n",
"\n",
"# --- بيانات الـ normal (التدريب) ---\n",
"X_norm_scaled = np.load(data_dir + r'\\normal_features_22_40k.npy')\n",
"ppi_norm = np.load(data_dir + r'\\normal_ppi_sequences_40k.npy')\n",
"\n",
"X_mal_scaled = np.load(data_dir + r'\\trickbot_features_FINAL.npy')\n",
"ppi_mal_normalized = np.load(data_dir + r'\\trickbot_ppi_FINAL.npy')\n",
"\n",
"# --- الـ scaler + الإحصائيات ---\n",
"scaler = joblib.load(data_dir + r'\\flow_scaler_40k.pkl')\n",
"ppi_norm_stats = joblib.load(data_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"\n",
"print(\"=== بيانات التدريب (normal) ===\")\n",
"print(f\" X_norm_scaled: {X_norm_scaled.shape}\")\n",
"print(f\" ppi_norm: {ppi_norm.shape}\")\n",
"print(\"\\n=== بيانات الاختبار (TrickBot) ===\")\n",
"print(f\" X_mal_scaled: {X_mal_scaled.shape}\")\n",
"print(f\" ppi_mal_normalized: {ppi_mal_normalized.shape}\")\n",
"\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"print(f\"Device: {device}\")"
]
},
{
"cell_type": "markdown",
"id": "63945810",
"metadata": {},
"source": [
"## 2. Model Definition, Training & Evaluation — CNN Conditioned on Packet Count"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "bb72e844",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import numpy as np\n",
"from sklearn.metrics import roc_auc_score, precision_score, recall_score, f1_score\n",
"from sklearn.model_selection import train_test_split\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"\n",
"SEED = 42\n",
"torch.manual_seed(SEED); np.random.seed(SEED)\n",
"\n",
"# --- تجهيز PACKETS المطبّع (العمودان 2,3 من X_norm_scaled) ---\n",
"pkt_norm_all = X_norm_scaled[:, [2, 3]].astype(np.float32) # PACKETS, PACKETS_REV\n",
"pkt_mal_all = X_mal_scaled[:, [2, 3]].astype(np.float32)\n",
"\n",
"# نفس تقسيم الـ indices (SEED=42) ليتطابق مع السكونس\n",
"rng = np.random.RandomState(SEED)\n",
"all_idx = np.arange(len(ppi_norm))\n",
"test_norm_idx = rng.choice(all_idx, size=len(ppi_mal_normalized), replace=False)\n",
"train_pool_idx = np.setdiff1d(all_idx, test_norm_idx)\n",
"\n",
"tr_idx, val_idx = train_test_split(train_pool_idx, test_size=0.1, random_state=SEED)\n",
"\n",
"# السكونس (N,3,30)\n",
"seq_tr = np.transpose(ppi_norm[tr_idx], (0, 2, 1))\n",
"seq_val = np.transpose(ppi_norm[val_idx], (0, 2, 1))\n",
"seq_test_norm = np.transpose(ppi_norm[test_norm_idx], (0, 2, 1))\n",
"seq_mal = np.transpose(ppi_mal_normalized, (0, 2, 1))\n",
"\n",
"# PACKETS المرافق (نفس الـ indices)\n",
"pkt_tr = pkt_norm_all[tr_idx]\n",
"pkt_val = pkt_norm_all[val_idx]\n",
"pkt_test_norm = pkt_norm_all[test_norm_idx]\n",
"pkt_mal = pkt_mal_all\n",
"\n",
"print(f\"Train seq: {seq_tr.shape} | pkt: {pkt_tr.shape}\")\n",
"print(f\"Test: normal {seq_test_norm.shape} | malware {seq_mal.shape}\")\n",
"\n",
"train_loader = DataLoader(TensorDataset(\n",
" torch.tensor(seq_tr, dtype=torch.float32),\n",
" torch.tensor(pkt_tr, dtype=torch.float32)), batch_size=512, shuffle=True)\n",
"val_loader = DataLoader(TensorDataset(\n",
" torch.tensor(seq_val, dtype=torch.float32),\n",
" torch.tensor(pkt_val, dtype=torch.float32)), batch_size=512, shuffle=False)\n",
"\n",
"\n",
"class CNN_Conditioned(nn.Module):\n",
" \"\"\"CNN AE على السكونس + PACKETS كمدخل مساعد. المخرج = السكونس فقط.\"\"\"\n",
" def __init__(self, in_ch=3, pkt_dim=2, latent_dim=16, dropout_p=0.2):\n",
" super().__init__()\n",
" self.enc_conv1 = nn.Conv1d(in_ch, 16, 3, padding=1); self.enc_bn1 = nn.BatchNorm1d(16)\n",
" self.enc_conv2 = nn.Conv1d(16, 32, 3, padding=1); self.enc_bn2 = nn.BatchNorm1d(32)\n",
" self.enc_conv3 = nn.Conv1d(32, 16, 3, padding=1); self.enc_bn3 = nn.BatchNorm1d(16)\n",
" self.pool = nn.MaxPool1d(2); self.relu = nn.ReLU()\n",
" self.to_latent = nn.Linear(16 * 7, latent_dim)\n",
" self.pkt_dense = nn.Sequential(nn.Linear(pkt_dim, 4), nn.ReLU())\n",
" self.merge = nn.Sequential(nn.Linear(latent_dim + 4, latent_dim), nn.ReLU())\n",
" self.dropout = nn.Dropout(dropout_p)\n",
" self.from_latent = nn.Linear(latent_dim, 16 * 7)\n",
" self.dec_conv1 = nn.Conv1d(16, 32, 3, padding=1); self.dec_bn1 = nn.BatchNorm1d(32); self.up1 = nn.Upsample(size=15)\n",
" self.dec_conv2 = nn.Conv1d(32, 16, 3, padding=1); self.dec_bn2 = nn.BatchNorm1d(16); self.up2 = nn.Upsample(size=30)\n",
" self.dec_out = nn.Conv1d(16, in_ch, 3, padding=1)\n",
"\n",
" def forward(self, x_seq, x_pkt):\n",
" z = self.pool(self.relu(self.enc_bn1(self.enc_conv1(x_seq))))\n",
" z = self.pool(self.relu(self.enc_bn2(self.enc_conv2(z))))\n",
" z = self.relu(self.enc_bn3(self.enc_conv3(z)))\n",
" z = self.to_latent(z.flatten(1))\n",
" p = self.pkt_dense(x_pkt)\n",
" merged = self.merge(torch.cat([z, p], dim=1))\n",
" merged = self.dropout(merged)\n",
" d = self.from_latent(merged).view(-1, 16, 7)\n",
" d = self.up1(self.relu(self.dec_bn1(self.dec_conv1(d))))\n",
" d = self.up2(self.relu(self.dec_bn2(self.dec_conv2(d))))\n",
" return self.dec_out(d)\n",
"\n",
"model = CNN_Conditioned().to(device)\n",
"print(f\"Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n",
"\n",
"criterion = nn.MSELoss()\n",
"optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)\n",
"scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)\n",
"\n",
"def recon_err_cond(model, seq, pkt, bs=4096):\n",
" model.eval(); errs = []\n",
" with torch.no_grad():\n",
" for i in range(0, len(seq), bs):\n",
" sb = torch.tensor(seq[i:i+bs], dtype=torch.float32, device=device)\n",
" pb = torch.tensor(pkt[i:i+bs], dtype=torch.float32, device=device)\n",
" out = model(sb, pb)\n",
" errs.append(((out - sb) ** 2).mean(dim=(1, 2)).cpu().numpy())\n",
" return np.concatenate(errs)\n",
"\n",
"y_true_test = np.concatenate([np.zeros(len(seq_test_norm)), np.ones(len(seq_mal))])\n",
"\n",
"MAX_EPOCHS, WARMUP, PATIENCE, BASE_LR = 100, 3, 5, 1e-3\n",
"best_val, wait = float('inf'), 0\n",
"out_dir = r'D:\\project 5\\data'\n",
"\n",
"for epoch in range(MAX_EPOCHS):\n",
" if epoch < WARMUP:\n",
" for g in optimizer.param_groups:\n",
" g['lr'] = BASE_LR * (epoch + 1) / WARMUP\n",
" model.train(); tr = 0.0\n",
" for sb, pb in train_loader:\n",
" sb, pb = sb.to(device), pb.to(device)\n",
" optimizer.zero_grad()\n",
" loss = criterion(model(sb, pb), sb)\n",
" loss.backward(); optimizer.step()\n",
" tr += loss.item() * sb.size(0)\n",
" tr /= len(train_loader.dataset)\n",
"\n",
" model.eval(); vl = 0.0\n",
" with torch.no_grad():\n",
" for sb, pb in val_loader:\n",
" sb, pb = sb.to(device), pb.to(device)\n",
" vl += criterion(model(sb, pb), sb).item() * sb.size(0)\n",
" vl /= len(val_loader.dataset)\n",
"\n",
" if epoch >= WARMUP:\n",
" scheduler.step(vl)\n",
"\n",
" if (epoch + 1) % 5 == 0 or epoch < 5:\n",
" en = recon_err_cond(model, seq_test_norm, pkt_test_norm)\n",
" em = recon_err_cond(model, seq_mal, pkt_mal)\n",
" auc = roc_auc_score(y_true_test, np.concatenate([en, em]))\n",
" print(f\"Epoch {epoch+1:3d} | train {tr:.5f} | val {vl:.5f} | ROC-AUC {auc:.4f}\")\n",
"\n",
" if vl < best_val:\n",
" best_val, wait = vl, 0\n",
" torch.save(model.state_dict(), out_dir + r'\\cnn_cond_best.pt')\n",
" else:\n",
" if epoch >= WARMUP:\n",
" wait += 1\n",
" if wait >= PATIENCE:\n",
" print(f\"Early stopping at epoch {epoch+1}\")\n",
" break\n",
"\n",
"model.load_state_dict(torch.load(out_dir + r'\\cnn_cond_best.pt'))\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
models/CNN_FlagCond.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "code",
"execution_count": null,
"id": "7cadab29",
"metadata": {},
"outputs": [],
"source": [
"# Architecture Search — CNN Conditioned on TCP Flags\n",
"\n",
"## 1. Load Preprocessed Data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ecb7c74",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import joblib\n",
"import torch\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"data_dir = r'D:\\project 5\\data'\n",
"\n",
"# --- بيانات الـ normal (التدريب) ---\n",
"X_norm_scaled = np.load(data_dir + r'\\normal_features_22_40k.npy')\n",
"ppi_norm = np.load(data_dir + r'\\normal_ppi_sequences_40k.npy')\n",
"\n",
"# --- بيانات TrickBot (الاختبار فقط) ---\n",
"X_mal_scaled = np.load(data_dir + r'\\trickbot_features_FINAL.npy')\n",
"ppi_mal_normalized = np.load(data_dir + r'\\trickbot_ppi_FINAL.npy')\n",
"\n",
"# --- الـ scaler + الإحصائيات ---\n",
"scaler = joblib.load(data_dir + r'\\flow_scaler_40k.pkl')\n",
"ppi_norm_stats = joblib.load(data_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"\n",
"print(\"=== بيانات التدريب (normal) ===\")\n",
"print(f\" X_norm_scaled: {X_norm_scaled.shape}\")\n",
"print(f\" ppi_norm: {ppi_norm.shape}\")\n",
"print(\"\\n=== بيانات الاختبار (TrickBot) ===\")\n",
"print(f\" X_mal_scaled: {X_mal_scaled.shape}\")\n",
"print(f\" ppi_mal_normalized: {ppi_mal_normalized.shape}\")\n",
"\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"print(f\"Device: {device}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "709fe18f",
"metadata": {},
"outputs": [],
"source": [
"## 2. Model Definition & Training — CNN Conditioned on TCP Flags"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "e9f402ac",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import numpy as np\n",
"from sklearn.metrics import roc_auc_score, precision_score, recall_score, f1_score\n",
"from sklearn.model_selection import train_test_split\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"\n",
"SEED = 42\n",
"torch.manual_seed(SEED); np.random.seed(SEED)\n",
"\n",
"# --- الـ 14 flags = الأعمدة 8..21 من X_norm_scaled (نفس ترتيب feature_22) ---\n",
"flags_norm_all = X_norm_scaled[:, 8:22].astype(np.float32)\n",
"flags_mal_all = X_mal_scaled[:, 8:22].astype(np.float32)\n",
"\n",
"# نفس تقسيم الـ indices (SEED=42) ليتطابق مع السكونس\n",
"rng = np.random.RandomState(SEED)\n",
"all_idx = np.arange(len(ppi_norm))\n",
"test_norm_idx = rng.choice(all_idx, size=len(ppi_mal_normalized), replace=False)\n",
"train_pool_idx = np.setdiff1d(all_idx, test_norm_idx)\n",
"tr_idx, val_idx = train_test_split(train_pool_idx, test_size=0.1, random_state=SEED)\n",
"\n",
"# السكونس (N,3,30)\n",
"seq_tr = np.transpose(ppi_norm[tr_idx], (0, 2, 1))\n",
"seq_val = np.transpose(ppi_norm[val_idx], (0, 2, 1))\n",
"seq_test_norm = np.transpose(ppi_norm[test_norm_idx], (0, 2, 1))\n",
"seq_mal = np.transpose(ppi_mal_normalized, (0, 2, 1))\n",
"\n",
"# الـ flags المرافقة\n",
"fl_tr = flags_norm_all[tr_idx]\n",
"fl_val = flags_norm_all[val_idx]\n",
"fl_test_norm = flags_norm_all[test_norm_idx]\n",
"fl_mal = flags_mal_all\n",
"\n",
"print(f\"Train seq: {seq_tr.shape} | flags: {fl_tr.shape}\")\n",
"print(f\"Test: normal {seq_test_norm.shape} | malware {seq_mal.shape}\")\n",
"\n",
"train_loader = DataLoader(TensorDataset(\n",
" torch.tensor(seq_tr, dtype=torch.float32),\n",
" torch.tensor(fl_tr, dtype=torch.float32)), batch_size=512, shuffle=True)\n",
"val_loader = DataLoader(TensorDataset(\n",
" torch.tensor(seq_val, dtype=torch.float32),\n",
" torch.tensor(fl_val, dtype=torch.float32)), batch_size=512, shuffle=False)\n",
"\n",
"\n",
"class CNN_FlagCond(nn.Module):\n",
" \"\"\"CNN AE على السكونس + 14 flags كمدخل مساعد. المخرج = السكونس فقط.\"\"\"\n",
" def __init__(self, in_ch=3, flag_dim=14, latent_dim=16, dropout_p=0.2):\n",
" super().__init__()\n",
" self.enc_conv1 = nn.Conv1d(in_ch, 16, 3, padding=1); self.enc_bn1 = nn.BatchNorm1d(16)\n",
" self.enc_conv2 = nn.Conv1d(16, 32, 3, padding=1); self.enc_bn2 = nn.BatchNorm1d(32)\n",
" self.enc_conv3 = nn.Conv1d(32, 16, 3, padding=1); self.enc_bn3 = nn.BatchNorm1d(16)\n",
" self.pool = nn.MaxPool1d(2); self.relu = nn.ReLU()\n",
" self.to_latent = nn.Linear(16 * 7, latent_dim)\n",
" self.flag_dense = nn.Sequential(nn.Linear(flag_dim, 8), nn.ReLU())\n",
" self.merge = nn.Sequential(nn.Linear(latent_dim + 8, latent_dim), nn.ReLU())\n",
" self.dropout = nn.Dropout(dropout_p)\n",
" self.from_latent = nn.Linear(latent_dim, 16 * 7)\n",
" self.dec_conv1 = nn.Conv1d(16, 32, 3, padding=1); self.dec_bn1 = nn.BatchNorm1d(32); self.up1 = nn.Upsample(size=15)\n",
" self.dec_conv2 = nn.Conv1d(32, 16, 3, padding=1); self.dec_bn2 = nn.BatchNorm1d(16); self.up2 = nn.Upsample(size=30)\n",
" self.dec_out = nn.Conv1d(16, in_ch, 3, padding=1)\n",
"\n",
" def forward(self, x_seq, x_flag):\n",
" z = self.pool(self.relu(self.enc_bn1(self.enc_conv1(x_seq))))\n",
" z = self.pool(self.relu(self.enc_bn2(self.enc_conv2(z))))\n",
" z = self.relu(self.enc_bn3(self.enc_conv3(z)))\n",
" z = self.to_latent(z.flatten(1))\n",
" f = self.flag_dense(x_flag)\n",
" merged = self.dropout(self.merge(torch.cat([z, f], dim=1)))\n",
" d = self.from_latent(merged).view(-1, 16, 7)\n",
" d = self.up1(self.relu(self.dec_bn1(self.dec_conv1(d))))\n",
" d = self.up2(self.relu(self.dec_bn2(self.dec_conv2(d))))\n",
" return self.dec_out(d)\n",
"\n",
"model = CNN_FlagCond().to(device)\n",
"print(f\"Parameters: {sum(p.numel() for p in model.parameters()):,}\")\n",
"\n",
"criterion = nn.MSELoss()\n",
"optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)\n",
"scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=2)\n",
"\n",
"def recon_err_fc(model, seq, fl, bs=4096):\n",
" model.eval(); errs = []\n",
" with torch.no_grad():\n",
" for i in range(0, len(seq), bs):\n",
" sb = torch.tensor(seq[i:i+bs], dtype=torch.float32, device=device)\n",
" fb = torch.tensor(fl[i:i+bs], dtype=torch.float32, device=device)\n",
" out = model(sb, fb)\n",
" errs.append(((out - sb) ** 2).mean(dim=(1, 2)).cpu().numpy())\n",
" return np.concatenate(errs)\n",
"\n",
"y_true_test = np.concatenate([np.zeros(len(seq_test_norm)), np.ones(len(seq_mal))])\n",
"\n",
"MAX_EPOCHS, WARMUP, PATIENCE, BASE_LR = 100, 3, 5, 1e-3\n",
"best_val, wait = float('inf'), 0\n",
"out_dir = r'D:\\project 5\\data'\n",
"\n",
"for epoch in range(MAX_EPOCHS):\n",
" if epoch < WARMUP:\n",
" for g in optimizer.param_groups:\n",
" g['lr'] = BASE_LR * (epoch + 1) / WARMUP\n",
" model.train(); tr = 0.0\n",
" for sb, fb in train_loader:\n",
" sb, fb = sb.to(device), fb.to(device)\n",
" optimizer.zero_grad()\n",
" loss = criterion(model(sb, fb), sb)\n",
" loss.backward(); optimizer.step()\n",
" tr += loss.item() * sb.size(0)\n",
" tr /= len(train_loader.dataset)\n",
"\n",
" model.eval(); vl = 0.0\n",
" with torch.no_grad():\n",
" for sb, fb in val_loader:\n",
" sb, fb = sb.to(device), fb.to(device)\n",
" vl += criterion(model(sb, fb), sb).item() * sb.size(0)\n",
" vl /= len(val_loader.dataset)\n",
"\n",
" if epoch >= WARMUP:\n",
" scheduler.step(vl)\n",
"\n",
" if (epoch + 1) % 5 == 0 or epoch < 5:\n",
" en = recon_err_fc(model, seq_test_norm, fl_test_norm)\n",
" em = recon_err_fc(model, seq_mal, fl_mal)\n",
" auc = roc_auc_score(y_true_test, np.concatenate([en, em]))\n",
" print(f\"Epoch {epoch+1:3d} | train {tr:.5f} | val {vl:.5f} | ROC-AUC {auc:.4f}\")\n",
"\n",
" if vl < best_val:\n",
" best_val, wait = vl, 0\n",
" torch.save(model.state_dict(), out_dir + r'\\cnn_flagcond_best.pt')\n",
" else:\n",
" if epoch >= WARMUP:\n",
" wait += 1\n",
" if wait >= PATIENCE:\n",
" print(f\"Early stopping at epoch {epoch+1}\")\n",
" break\n",
"\n",
"model.load_state_dict(torch.load(out_dir + r'\\cnn_flagcond_best.pt'))\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
models/DensePPI_AE.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "markdown",
"id": "2bff8dcc",
"metadata": {},
"source": [
"# Architecture Search — Dense PPI AutoEncoder (Kim & Kim Baseline)\n",
"\n",
"## 1. Load Preprocessed Data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ecb7c74",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import joblib\n",
"import torch\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"data_dir = r'D:\\project 5\\data'\n",
"\n",
"# --- بيانات الـ normal (التدريب) ---\n",
"X_norm_scaled = np.load(data_dir + r'\\normal_features_22_40k.npy')\n",
"ppi_norm = np.load(data_dir + r'\\normal_ppi_sequences_40k.npy')\n",
"\n",
"# --- بيانات TrickBot (الاختبار فقط) ---\n",
"X_mal_scaled = np.load(data_dir + r'\\trickbot_features_FINAL.npy')\n",
"ppi_mal_normalized = np.load(data_dir + r'\\trickbot_ppi_FINAL.npy')\n",
"\n",
"# --- الـ scaler + الإحصائيات ---\n",
"scaler = joblib.load(data_dir + r'\\flow_scaler_40k.pkl')\n",
"ppi_norm_stats = joblib.load(data_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"\n",
"print(\"=== بيانات التدريب (normal) ===\")\n",
"print(f\" X_norm_scaled: {X_norm_scaled.shape}\")\n",
"print(f\" ppi_norm: {ppi_norm.shape}\")\n",
"print(\"\\n=== بيانات الاختبار (TrickBot) ===\")\n",
"print(f\" X_mal_scaled: {X_mal_scaled.shape}\")\n",
"print(f\" ppi_mal_normalized: {ppi_mal_normalized.shape}\")\n",
"\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"print(f\"Device: {device}\")"
]
},
{
"cell_type": "markdown",
"id": "44e0c772",
"metadata": {},
"source": [
"## 2. Model Definition & Training — Dense PPI AutoEncoder (Kim & Kim Baseline)"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "95a82fdd",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import numpy as np\n",
"from sklearn.model_selection import train_test_split\n",
"from sklearn.metrics import roc_auc_score, precision_score, recall_score, f1_score\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"\n",
"SEED = 42\n",
"torch.manual_seed(SEED); np.random.seed(SEED)\n",
"N_MAL = len(ppi_mal_normalized)\n",
"\n",
"# --- تسطيح: (N,30,3) → (N,90) ---\n",
"flat_norm = ppi_norm.reshape(len(ppi_norm), -1)\n",
"flat_mal = ppi_mal_normalized.reshape(N_MAL, -1)\n",
"\n",
"rng = np.random.RandomState(SEED)\n",
"all_idx = np.arange(len(flat_norm))\n",
"test_norm_idx = rng.choice(all_idx, size=N_MAL, replace=False)\n",
"train_pool_idx = np.setdiff1d(all_idx, test_norm_idx)\n",
"tr_idx, val_idx = train_test_split(train_pool_idx, test_size=0.1, random_state=SEED)\n",
"\n",
"X_tr = flat_norm[tr_idx]\n",
"X_val = flat_norm[val_idx]\n",
"X_test_norm = flat_norm[test_norm_idx]\n",
"X_mal = flat_mal\n",
"\n",
"print(f\"Train: {X_tr.shape} | Val: {X_val.shape}\")\n",
"print(f\"Test: normal {X_test_norm.shape} | malware {X_mal.shape}\")\n",
"\n",
"train_loader = DataLoader(TensorDataset(torch.tensor(X_tr, dtype=torch.float32)),\n",
" batch_size=512, shuffle=True)\n",
"val_loader = DataLoader(TensorDataset(torch.tensor(X_val, dtype=torch.float32)),\n",
" batch_size=512, shuffle=False)\n",
"\n",
"\n",
"class DensePPI_AE(nn.Module):\n",
" \"\"\"baseline Kim & Kim: 90→67→45→67→90 مع Dropout\"\"\"\n",
" def __init__(self, input_dim=90, dropout_p=0.2):\n",
" super().__init__()\n",
" self.encoder = nn.Sequential(\n",
" nn.Linear(input_dim, 67), nn.ReLU(), nn.Dropout(dropout_p),\n",
" nn.Linear(67, 45), nn.ReLU())\n",
" self.decoder = nn.Sequential(\n",
" nn.Linear(45, 67), nn.ReLU(), nn.Dropout(dropout_p),\n",
" nn.Linear(67, input_dim))\n",
" def forward(self, x):\n",
" return self.decoder(self.encoder(x))\n",
"\n",
"model = DensePPI_AE().to(device)\n",
"print(f\"Dense PPI AE parameters: {sum(p.numel() for p in model.parameters()):,}\")\n",
"\n",
"criterion = nn.MSELoss()\n",
"optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-5)\n",
"scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=4)\n",
"\n",
"def recon_err_dense(model, X, bs=8192):\n",
" model.eval(); errs = []\n",
" with torch.no_grad():\n",
" for i in range(0, len(X), bs):\n",
" xb = torch.tensor(X[i:i+bs], dtype=torch.float32, device=device)\n",
" out = model(xb)\n",
" errs.append(((out - xb) ** 2).mean(dim=1).cpu().numpy())\n",
" return np.concatenate(errs)\n",
"\n",
"y_true_test = np.concatenate([np.zeros(N_MAL), np.ones(N_MAL)])\n",
"\n",
"MAX_EPOCHS, WARMUP, PATIENCE, BASE_LR = 100, 5, 8, 1e-3\n",
"best_val, wait = float('inf'), 0\n",
"out_dir = r'D:\\project 5\\data'\n",
"\n",
"for epoch in range(MAX_EPOCHS):\n",
" if epoch < WARMUP:\n",
" for g in optimizer.param_groups:\n",
" g['lr'] = BASE_LR * (epoch + 1) / WARMUP\n",
" model.train(); tr = 0.0\n",
" for (xb,) in train_loader:\n",
" xb = xb.to(device)\n",
" optimizer.zero_grad()\n",
" loss = criterion(model(xb), xb)\n",
" loss.backward()\n",
" torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
" optimizer.step()\n",
" tr += loss.item() * xb.size(0)\n",
" tr /= len(train_loader.dataset)\n",
"\n",
" model.eval(); vl = 0.0\n",
" with torch.no_grad():\n",
" for (xb,) in val_loader:\n",
" xb = xb.to(device)\n",
" vl += criterion(model(xb), xb).item() * xb.size(0)\n",
" vl /= len(val_loader.dataset)\n",
"\n",
" if epoch >= WARMUP:\n",
" scheduler.step(vl)\n",
"\n",
" if (epoch + 1) % 5 == 0 or epoch < 5:\n",
" en = recon_err_dense(model, X_test_norm)\n",
" em = recon_err_dense(model, X_mal)\n",
" auc = roc_auc_score(y_true_test, np.concatenate([en, em]))\n",
" print(f\"Epoch {epoch+1:3d} | train {tr:.5f} | val {vl:.5f} | ROC-AUC {auc:.4f}\")\n",
"\n",
" if vl < best_val:\n",
" best_val, wait = vl, 0\n",
" torch.save(model.state_dict(), out_dir + r'\\dense_ppi_40k_best.pt')\n",
" else:\n",
" if epoch >= WARMUP:\n",
" wait += 1\n",
" if wait >= PATIENCE:\n",
" print(f\"Early stopping at epoch {epoch+1}\")\n",
" break\n",
"\n",
"model.load_state_dict(torch.load(out_dir + r'\\dense_ppi_40k_best.pt'))\n"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
models/FlowAE.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "markdown",
"id": "5e48da0c",
"metadata": {},
"source": [
"# Architecture Search — Flow-Level AutoEncoder (22 Features)\n",
"\n",
"## 1. Load Preprocessed Data"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "3ecb7c74",
"metadata": {},
"outputs": [],
"source": [
"import numpy as np\n",
"import joblib\n",
"import torch\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"data_dir = r'D:\\project 5\\data'\n",
"\n",
"# --- بيانات الـ normal (التدريب) ---\n",
"X_norm_scaled = np.load(data_dir + r'\\normal_features_22_40k.npy')\n",
"ppi_norm = np.load(data_dir + r'\\normal_ppi_sequences_40k.npy')\n",
"\n",
"# --- بيانات TrickBot (الاختبار فقط) ---\n",
"X_mal_scaled = np.load(data_dir + r'\\trickbot_features_FINAL.npy')\n",
"ppi_mal_normalized = np.load(data_dir + r'\\trickbot_ppi_FINAL.npy')\n",
"\n",
"# --- الـ scaler + الإحصائيات ---\n",
"scaler = joblib.load(data_dir + r'\\flow_scaler_40k.pkl')\n",
"ppi_norm_stats = joblib.load(data_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"\n",
"print(\"=== بيانات التدريب (normal) ===\")\n",
"print(f\" X_norm_scaled: {X_norm_scaled.shape}\")\n",
"print(f\" ppi_norm: {ppi_norm.shape}\")\n",
"print(\"\\n=== بيانات الاختبار (TrickBot) ===\")\n",
"print(f\" X_mal_scaled: {X_mal_scaled.shape}\")\n",
"print(f\" ppi_mal_normalized: {ppi_mal_normalized.shape}\")\n",
"\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"print(f\"Device: {device}\")"
]
},
{
"cell_type": "markdown",
"id": "d151b505",
"metadata": {},
"source": [
"## 2. Model Definition & Training — Flow AutoEncoder"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0e399b2e",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import numpy as np\n",
"from sklearn.model_selection import train_test_split\n",
"from sklearn.metrics import roc_auc_score, precision_score, recall_score, f1_score\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"\n",
"SEED = 42\n",
"torch.manual_seed(SEED); np.random.seed(SEED)\n",
"\n",
"N_MAL = len(X_mal_scaled)\n",
"\n",
"# --- تقسيم متوازن ---\n",
"rng = np.random.RandomState(SEED)\n",
"all_idx = np.arange(len(X_norm_scaled))\n",
"test_norm_idx = rng.choice(all_idx, size=N_MAL, replace=False)\n",
"train_pool_idx = np.setdiff1d(all_idx, test_norm_idx)\n",
"tr_idx, val_idx = train_test_split(train_pool_idx, test_size=0.1, random_state=SEED)\n",
"\n",
"X_train = X_norm_scaled[tr_idx]\n",
"X_val = X_norm_scaled[val_idx]\n",
"X_test_norm = X_norm_scaled[test_norm_idx]\n",
"\n",
"print(f\"Train: {X_train.shape} | Val: {X_val.shape}\")\n",
"print(f\"Test: normal {X_test_norm.shape} | malware {X_mal_scaled.shape}\")\n",
"\n",
"train_loader = DataLoader(TensorDataset(torch.tensor(X_train, dtype=torch.float32)),\n",
" batch_size=512, shuffle=True)\n",
"val_loader = DataLoader(TensorDataset(torch.tensor(X_val, dtype=torch.float32)),\n",
" batch_size=512, shuffle=False)\n",
"\n",
"class FlowAE(nn.Module):\n",
" def __init__(self, input_dim=22):\n",
" super().__init__()\n",
" self.encoder = nn.Sequential(\n",
" nn.Linear(input_dim, 16), nn.ReLU(),\n",
" nn.Linear(16, 8), nn.ReLU(),\n",
" nn.Linear(8, 4), nn.ReLU())\n",
" self.decoder = nn.Sequential(\n",
" nn.Linear(4, 8), nn.ReLU(),\n",
" nn.Linear(8, 16), nn.ReLU(),\n",
" nn.Linear(16, input_dim))\n",
" def forward(self, x):\n",
" return self.decoder(self.encoder(x))\n",
"\n",
"flow_ae = FlowAE(22).to(device)\n",
"print(f\"Flow AE parameters: {sum(p.numel() for p in flow_ae.parameters()):,}\")\n",
"\n",
"criterion = nn.MSELoss()\n",
"optimizer = torch.optim.AdamW(flow_ae.parameters(), lr=1e-3, weight_decay=1e-5)\n",
"scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=4)\n",
"\n",
"def recon_err_flow(model, X, bs=8192):\n",
" model.eval(); errs = []\n",
" with torch.no_grad():\n",
" for i in range(0, len(X), bs):\n",
" xb = torch.tensor(X[i:i+bs], dtype=torch.float32, device=device)\n",
" out = model(xb)\n",
" errs.append(((out - xb) ** 2).mean(dim=1).cpu().numpy())\n",
" return np.concatenate(errs)\n",
"\n",
"y_true_test = np.concatenate([np.zeros(N_MAL), np.ones(N_MAL)])\n",
"\n",
"MAX_EPOCHS, WARMUP, PATIENCE, BASE_LR = 100, 5, 8, 1e-3\n",
"best_val, wait = float('inf'), 0\n",
"out_dir = r'D:\\project 5\\data'\n",
"\n",
"for epoch in range(MAX_EPOCHS):\n",
" if epoch < WARMUP:\n",
" for g in optimizer.param_groups:\n",
" g['lr'] = BASE_LR * (epoch + 1) / WARMUP\n",
" flow_ae.train(); tr = 0.0\n",
" for (xb,) in train_loader:\n",
" xb = xb.to(device)\n",
" optimizer.zero_grad()\n",
" loss = criterion(flow_ae(xb), xb)\n",
" loss.backward()\n",
" torch.nn.utils.clip_grad_norm_(flow_ae.parameters(), 1.0)\n",
" optimizer.step()\n",
" tr += loss.item() * xb.size(0)\n",
" tr /= len(train_loader.dataset)\n",
"\n",
" flow_ae.eval(); vl = 0.0\n",
" with torch.no_grad():\n",
" for (xb,) in val_loader:\n",
" xb = xb.to(device)\n",
" vl += criterion(flow_ae(xb), xb).item() * xb.size(0)\n",
" vl /= len(val_loader.dataset)\n",
"\n",
" if epoch >= WARMUP:\n",
" scheduler.step(vl)\n",
"\n",
" if (epoch + 1) % 5 == 0 or epoch < 5:\n",
" en = recon_err_flow(flow_ae, X_test_norm)\n",
" em = recon_err_flow(flow_ae, X_mal_scaled)\n",
" auc = roc_auc_score(y_true_test, np.concatenate([en, em]))\n",
" print(f\"Epoch {epoch+1:3d} | train {tr:.5f} | val {vl:.5f} | ROC-AUC {auc:.4f}\")\n",
"\n",
" if vl < best_val:\n",
" best_val, wait = vl, 0\n",
" torch.save(flow_ae.state_dict(), out_dir + r'\\flow_ae_40k_best.pt')\n",
" else:\n",
" if epoch >= WARMUP:\n",
" wait += 1\n",
" if wait >= PATIENCE:\n",
" print(f\"Early stopping at epoch {epoch+1}\")\n",
" break\n",
"\n",
"flow_ae.load_state_dict(torch.load(out_dir + r'\\flow_ae_40k_best.pt'))"
]
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
models/StackedLSTM_AE.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "markdown",
"id": "d2ccce72",
"metadata": {},
"source": [
"# Architecture Search — Stacked LSTM AutoEncoder (L=30 Baseline)\n",
"\n",
"## 1. Load Preprocessed Data"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "3ecb7c74",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"=== بيانات التدريب (normal) ===\n",
" X_norm_scaled: (7295027, 22)\n",
" ppi_norm: (7295027, 30, 3)\n",
"\n",
"=== بيانات الاختبار (TrickBot) ===\n",
" X_mal_scaled: (10823, 22)\n",
" ppi_mal_normalized: (10823, 30, 3)\n",
"Device: cuda\n"
]
}
],
"source": [
"import numpy as np\n",
"import joblib\n",
"import torch\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"data_dir = r'D:\\project 5\\data'\n",
"\n",
"X_norm_scaled = np.load(data_dir + r'\\normal_features_22_40k.npy')\n",
"ppi_norm = np.load(data_dir + r'\\normal_ppi_sequences_40k.npy')\n",
"\n",
"X_mal_scaled = np.load(data_dir + r'\\trickbot_features_FINAL.npy')\n",
"ppi_mal_normalized = np.load(data_dir + r'\\trickbot_ppi_FINAL.npy')\n",
"\n",
"scaler = joblib.load(data_dir + r'\\flow_scaler_40k.pkl')\n",
"ppi_norm_stats = joblib.load(data_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"\n",
"print(\"=== بيانات التدريب (normal) ===\")\n",
"print(f\" X_norm_scaled: {X_norm_scaled.shape}\")\n",
"print(f\" ppi_norm: {ppi_norm.shape}\")\n",
"print(\"\\n=== بيانات الاختبار (TrickBot) ===\")\n",
"print(f\" X_mal_scaled: {X_mal_scaled.shape}\")\n",
"print(f\" ppi_mal_normalized: {ppi_mal_normalized.shape}\")\n",
"\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"print(f\"Device: {device}\")"
]
},
{
"cell_type": "markdown",
"id": "90ec4790",
"metadata": {},
"source": [
"## 2. Model Definition & Training — Stacked LSTM AutoEncoder"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "54449f49",
"metadata": {},
"outputs": [],
"source": [
"import torch\n",
"import torch.nn as nn\n",
"import numpy as np\n",
"from sklearn.model_selection import train_test_split\n",
"from sklearn.metrics import roc_auc_score, precision_score, recall_score, f1_score\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"\n",
"SEED = 42\n",
"torch.manual_seed(SEED); np.random.seed(SEED)\n",
"N_MAL = len(ppi_mal_normalized)\n",
"\n",
"rng = np.random.RandomState(SEED)\n",
"all_idx = np.arange(len(ppi_norm))\n",
"test_norm_idx = rng.choice(all_idx, size=N_MAL, replace=False)\n",
"train_pool_idx = np.setdiff1d(all_idx, test_norm_idx)\n",
"tr_idx, val_idx = train_test_split(train_pool_idx, test_size=0.1, random_state=SEED)\n",
"\n",
"seq_tr = ppi_norm[tr_idx]\n",
"seq_val = ppi_norm[val_idx]\n",
"seq_test_norm = ppi_norm[test_norm_idx]\n",
"seq_mal = ppi_mal_normalized\n",
"\n",
"print(f\"Train: {seq_tr.shape} | Val: {seq_val.shape}\")\n",
"print(f\"Test: normal {seq_test_norm.shape} | malware {seq_mal.shape}\")\n",
"\n",
"train_loader = DataLoader(TensorDataset(torch.tensor(seq_tr, dtype=torch.float32)),\n",
" batch_size=512, shuffle=True)\n",
"val_loader = DataLoader(TensorDataset(torch.tensor(seq_val, dtype=torch.float32)),\n",
" batch_size=512, shuffle=False)\n",
"\n",
"\n",
"class StackedLSTM_AE(nn.Module):\n",
" def __init__(self, n_features=3, seq_len=30, hidden=32, latent=16, dropout_p=0.2):\n",
" super().__init__()\n",
" self.seq_len = seq_len\n",
" self.enc_lstm = nn.LSTM(n_features, hidden, num_layers=2, batch_first=True, dropout=dropout_p)\n",
" self.enc_norm = nn.LayerNorm(hidden)\n",
" self.to_latent = nn.Linear(hidden, latent)\n",
" self.from_latent = nn.Linear(latent, hidden)\n",
" self.dec_lstm = nn.LSTM(hidden, hidden, num_layers=2, batch_first=True, dropout=dropout_p)\n",
" self.dec_norm = nn.LayerNorm(hidden)\n",
" self.out = nn.Linear(hidden, n_features)\n",
"\n",
" def forward(self, x):\n",
" _, (h, _) = self.enc_lstm(x)\n",
" h_last = self.enc_norm(h[-1])\n",
" z = self.to_latent(h_last)\n",
" d = self.from_latent(z).unsqueeze(1).repeat(1, self.seq_len, 1)\n",
" d, _ = self.dec_lstm(d)\n",
" d = self.dec_norm(d)\n",
" return self.out(d)\n",
"\n",
"model = StackedLSTM_AE().to(device)\n",
"print(f\"LSTM AE parameters: {sum(p.numel() for p in model.parameters()):,}\")\n",
"\n",
"criterion = nn.MSELoss()\n",
"optimizer = torch.optim.AdamW(model.parameters(), lr=1e-3, weight_decay=1e-4)\n",
"scheduler = torch.optim.lr_scheduler.ReduceLROnPlateau(optimizer, mode='min', factor=0.5, patience=3)\n",
"\n",
"def recon_err_lstm(model, X, bs=4096):\n",
" model.eval(); errs = []\n",
" with torch.no_grad():\n",
" for i in range(0, len(X), bs):\n",
" xb = torch.tensor(X[i:i+bs], dtype=torch.float32, device=device)\n",
" out = model(xb)\n",
" errs.append(((out - xb) ** 2).mean(dim=(1, 2)).cpu().numpy())\n",
" return np.concatenate(errs)\n",
"\n",
"y_true_test = np.concatenate([np.zeros(N_MAL), np.ones(N_MAL)])\n",
"\n",
"MAX_EPOCHS, WARMUP, PATIENCE, BASE_LR = 100, 3, 6, 1e-3\n",
"best_val, wait = float('inf'), 0\n",
"out_dir = r'D:\\project 5\\data'\n",
"\n",
"for epoch in range(MAX_EPOCHS):\n",
" if epoch < WARMUP:\n",
" for g in optimizer.param_groups:\n",
" g['lr'] = BASE_LR * (epoch + 1) / WARMUP\n",
" model.train(); tr = 0.0\n",
" for (xb,) in train_loader:\n",
" xb = xb.to(device)\n",
" optimizer.zero_grad()\n",
" loss = criterion(model(xb), xb)\n",
" loss.backward()\n",
" torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)\n",
" optimizer.step()\n",
" tr += loss.item() * xb.size(0)\n",
" tr /= len(train_loader.dataset)\n",
"\n",
" model.eval(); vl = 0.0\n",
" with torch.no_grad():\n",
" for (xb,) in val_loader:\n",
" xb = xb.to(device)\n",
" vl += criterion(model(xb), xb).item() * xb.size(0)\n",
" vl /= len(val_loader.dataset)\n",
"\n",
" if epoch >= WARMUP:\n",
" scheduler.step(vl)\n",
"\n",
" if (epoch + 1) % 5 == 0 or epoch < 5:\n",
" en = recon_err_lstm(model, seq_test_norm)\n",
" em = recon_err_lstm(model, seq_mal)\n",
" auc = roc_auc_score(y_true_test, np.concatenate([en, em]))\n",
" print(f\"Epoch {epoch+1:3d} | train {tr:.5f} | val {vl:.5f} | ROC-AUC {auc:.4f}\")\n",
"\n",
" if vl < best_val:\n",
" best_val, wait = vl, 0\n",
" torch.save(model.state_dict(), out_dir + r'\\lstm_ae_40k_best.pt')\n",
" else:\n",
" if epoch >= WARMUP:\n",
" wait += 1\n",
" if wait >= PATIENCE:\n",
" print(f\"Early stopping at epoch {epoch+1}\")\n",
" break\n",
"\n",
"model.load_state_dict(torch.load(out_dir + r'\\lstm_ae_40k_best.pt'))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "0b7b20be",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "bad853de",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "a18c8219",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "eb976ae0",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
models/data_preprocessing.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "markdown",
"id": "64438eeb",
"metadata": {},
"source": [
"# Normal Traffic Preprocessing — CESNET-TLS22\n",
"\n",
"## 1. Load Stratified Sample"
]
},
{
"cell_type": "code",
"execution_count": 1,
"id": "34ee5c7c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Shape: (7295027, 25)\n",
"Unique apps: 191\n"
]
}
],
"source": [
"import polars as pl\n",
"\n",
"path = r'D:\\project 5\\data\\tls22_sample_40k.parquet'\n",
"df = pl.read_parquet(path)\n",
"\n",
"print(f\"Shape: {df.shape}\")\n",
"print(f\"Unique apps: {df['APP'].n_unique()}\")"
]
},
{
"cell_type": "markdown",
"id": "e48c5672",
"metadata": {},
"source": [
"## 2. Flow-Level Feature Engineering (22 Features)"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "10c9a789",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Normal flow features: (7295027, 22)\n"
]
}
],
"source": [
"import numpy as np\n",
"import joblib\n",
"from sklearn.preprocessing import StandardScaler\n",
"\n",
"log_cols = ['BYTES', 'BYTES_REV', 'PACKETS', 'PACKETS_REV', 'DURATION',\n",
" 'PPI_LEN', 'PPI_DURATION', 'PPI_ROUNDTRIPS']\n",
"flag_cols = ['FLAG_ECE', 'FLAG_ECE_REV', 'FLAG_URG', 'FLAG_URG_REV',\n",
" 'FLAG_ACK', 'FLAG_ACK_REV', 'FLAG_PSH', 'FLAG_PSH_REV',\n",
" 'FLAG_RST', 'FLAG_RST_REV', 'FLAG_SYN', 'FLAG_SYN_REV',\n",
" 'FLAG_FIN', 'FLAG_FIN_REV']\n",
"\n",
"X_log = np.log1p(df.select(log_cols).to_numpy().astype(np.float64))\n",
"\n",
"scaler = StandardScaler()\n",
"X_log_scaled = scaler.fit_transform(X_log) \n",
"\n",
"X_flags = df.select(flag_cols).to_numpy().astype(np.float64)\n",
"X_norm_scaled = np.hstack([X_log_scaled, X_flags]).astype(np.float32)\n",
"\n",
"out_dir = r'D:\\project 5\\data'\n",
"joblib.dump(scaler, out_dir + r'\\flow_scaler_40k.pkl')\n",
"np.save(out_dir + r'\\normal_features_22_40k.npy', X_norm_scaled)\n",
"\n",
"print(f\"Normal flow features: {X_norm_scaled.shape}\")"
]
},
{
"cell_type": "markdown",
"id": "3aa904f5",
"metadata": {},
"source": [
"## 3. Parse PPI Sequences to Fixed-Length Matrix (N, 30, 3) :\n",
"[0]=Timing/IAT، [1]=Direction، [2]=Size."
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "9b0ccdcf",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Processed 1,000,000/7,295,027\n",
"Processed 2,000,000/7,295,027\n",
"Processed 3,000,000/7,295,027\n",
"Processed 4,000,000/7,295,027\n",
"Processed 5,000,000/7,295,027\n",
"Processed 6,000,000/7,295,027\n",
"Processed 7,000,000/7,295,027\n",
"ppi_norm shape: (7295027, 30, 3)\n"
]
}
],
"source": [
"import ast\n",
"\n",
"MAX_PACKETS = 30\n",
"N = len(df)\n",
"ppi_norm = np.zeros((N, MAX_PACKETS, 3), dtype=np.float32)\n",
"\n",
"for i, s in enumerate(df['PPI'].to_list()):\n",
" channels = ast.literal_eval(s)\n",
" timing, direction, size = channels[0], channels[1], channels[2]\n",
" n_pkt = min(len(timing), MAX_PACKETS)\n",
" ppi_norm[i, :n_pkt, 0] = timing[:n_pkt]\n",
" ppi_norm[i, :n_pkt, 1] = direction[:n_pkt]\n",
" ppi_norm[i, :n_pkt, 2] = size[:n_pkt]\n",
" if (i + 1) % 1_000_000 == 0:\n",
" print(f\"Processed {i+1:,}/{N:,}\")\n",
"\n",
"print(f\"ppi_norm shape: {ppi_norm.shape}\")"
]
},
{
"cell_type": "markdown",
"id": "9ee26846",
"metadata": {},
"source": [
"## 4. Normalize PPI Channels (Timing & Size)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "b09fc5f7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Channel 0: mean=2.9660, std=2.0991\n",
"Channel 2: mean=5.9385, std=1.4185\n",
"Normal PPI normalized: (7295027, 30, 3)\n"
]
}
],
"source": [
"ppi_norm_stats = {}\n",
"\n",
"for col in [0, 2]:\n",
" vals = ppi_norm[:, :, col]\n",
" vals = np.log1p(np.abs(vals)) * np.sign(vals)\n",
" nonzero = vals[vals != 0]\n",
" mean, std = float(nonzero.mean()), float(nonzero.std())\n",
" ppi_norm[:, :, col] = (vals - mean) / (std + 1e-8)\n",
" ppi_norm_stats[col] = (mean, std)\n",
" print(f\"Channel {col}: mean={mean:.4f}, std={std:.4f}\")\n",
"\n",
"joblib.dump(ppi_norm_stats, out_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"np.save(out_dir + r'\\normal_ppi_sequences_40k.npy', ppi_norm)\n",
"\n",
"print(f\"Normal PPI normalized: {ppi_norm.shape}\")"
]
},
{
"cell_type": "markdown",
"id": "6bfcbcdb",
"metadata": {},
"source": [
"## "
]
},
{
"cell_type": "markdown",
"id": "4ccdfc7e",
"metadata": {},
"source": [
"# Malware Traffic Preprocessing — TrickBot\n",
"\n",
"## 1. Load & Clean Raw CSV"
]
},
{
"cell_type": "code",
"execution_count": 5,
"id": "df8cdda7",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Loaded: (31157, 34)\n"
]
}
],
"source": [
"import pandas as pd\n",
"\n",
"csv_file = r'D:\\project 5\\ipfixprobe\\malware_pcap\\2017-3-30_win6.pcap.csv'\n",
"df_mal = pd.read_csv(csv_file)\n",
"df_mal.columns = [c.split(' ', 1)[1] if ' ' in c else c for c in df_mal.columns]\n",
"\n",
"print(f\"Loaded: {df_mal.shape}\")"
]
},
{
"cell_type": "markdown",
"id": "12f56ceb",
"metadata": {},
"source": [
"## 2. Filter TLS Flows & Compute Derived Features\n",
"معالجة CSV الخام لتدفقات TrickBot عبر عدة خطوات متتالية:\n",
"1. **فلترة TLS**: الاحتفاظ فقط بالتدفقات على المنفذ 443 (مصدر أو وجهة)\n",
"2. **فلترة PPI_LEN ≥ 3**: استبعاد التدفقات القصيرة جدًا (أقل من 3 حزم)\n",
"3. **فلترة SYN**: الاحتفاظ فقط بالتدفقات التي تحمل راية SYN — هذا هو الفلتر الحاسم الذي يحدد العدد النهائي للصفوف\n",
"4. **حساب الأعمدة المشتقة** : `PPI_ROUNDTRIPS` (عدد تبدلات الاتجاه)، `PPI_DURATION` (المدة الزمنية من نمط الحزم)، `DURATION` (الفرق بين أول وآخر حزمة)، و14 راية TCP ثنائية (اتجاهين × 7 رايات)\n",
"5. بناء مصفوفة الـ 22 سمة النهائية "
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "ce93708b",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"بعد PPI_LEN>=3: 10,997\n",
"بعد فلترة SYN: 10,823 ( العدد النهائي لكل الملفات)\n",
"TrickBot feature matrix: (10823, 22) \n"
]
}
],
"source": [
"def parse_ppi_field(s):\n",
" s = str(s).strip('[]')\n",
" return s.split('|') if s != '' else []\n",
"\n",
"def get_roundtrips(s):\n",
" dirs = parse_ppi_field(s)\n",
" if len(dirs) < 2:\n",
" return 0\n",
" dirs = [int(d) for d in dirs]\n",
" return sum(1 for i in range(1, len(dirs)) if dirs[i] != dirs[i-1])\n",
"\n",
"def get_ppi_duration(s):\n",
" times = parse_ppi_field(s)\n",
" if len(times) < 2:\n",
" return 0.0\n",
" try:\n",
" t = pd.to_datetime(times)\n",
" return (t[-1] - t[0]).total_seconds()\n",
" except:\n",
" return 0.0\n",
"\n",
"def flag_bit(series, bit):\n",
" return (series.astype(int) & bit).astype(bool).astype(int)\n",
"\n",
"def has_syn(flags_val):\n",
" try:\n",
" return int(flags_val) & 2 == 2\n",
" except:\n",
" return False\n",
"\n",
"# فلترة TLS\n",
"df_tls = df_mal[(df_mal['DST_PORT'] == 443) | (df_mal['SRC_PORT'] == 443)].copy()\n",
"df_tls = df_tls.reset_index(drop=True)\n",
"\n",
"# فلترة PPI_LEN >= 3\n",
"df_tls['PPI_LEN'] = df_tls['PPI_PKT_DIRECTIONS'].apply(lambda x: len(parse_ppi_field(x)))\n",
"df_final = df_tls[df_tls['PPI_LEN'] >= 3].reset_index(drop=True)\n",
"print(f\"بعد PPI_LEN>=3: {len(df_final):,}\")\n",
"\n",
"# فلترة SYN \n",
"syn_mask = df_final['TCP_FLAGS'].apply(has_syn)\n",
"df_clean = df_final[syn_mask].reset_index(drop=True)\n",
"print(f\"بعد فلترة SYN: {len(df_clean):,} ( العدد النهائي لكل الملفات)\")\n",
"\n",
"df_clean['PPI_ROUNDTRIPS'] = df_clean['PPI_PKT_DIRECTIONS'].apply(get_roundtrips)\n",
"df_clean['PPI_DURATION'] = df_clean['PPI_PKT_TIMES'].apply(get_ppi_duration)\n",
"df_clean['DURATION'] = (pd.to_datetime(df_clean['TIME_LAST']) -\n",
" pd.to_datetime(df_clean['TIME_FIRST'])).dt.total_seconds()\n",
"\n",
"for name, bit in [('FIN', 1), ('SYN', 2), ('RST', 4), ('PSH', 8), ('ACK', 16), ('URG', 32), ('ECE', 64)]:\n",
" df_clean[f'FLAG_{name}'] = flag_bit(df_clean['TCP_FLAGS'], bit)\n",
" df_clean[f'FLAG_{name}_REV'] = flag_bit(df_clean['TCP_FLAGS_REV'], bit)\n",
"\n",
"feature_22 = [\n",
" 'BYTES', 'BYTES_REV', 'PACKETS', 'PACKETS_REV', 'DURATION',\n",
" 'PPI_LEN', 'PPI_DURATION', 'PPI_ROUNDTRIPS',\n",
" 'FLAG_ECE', 'FLAG_ECE_REV', 'FLAG_URG', 'FLAG_URG_REV',\n",
" 'FLAG_ACK', 'FLAG_ACK_REV', 'FLAG_PSH', 'FLAG_PSH_REV',\n",
" 'FLAG_RST', 'FLAG_RST_REV', 'FLAG_SYN', 'FLAG_SYN_REV',\n",
" 'FLAG_FIN', 'FLAG_FIN_REV'\n",
"]\n",
"df_features_mal = df_clean[feature_22].copy()\n",
"\n",
"print(f\"TrickBot feature matrix: {df_features_mal.shape} \")"
]
},
{
"cell_type": "markdown",
"id": "7c466bf7",
"metadata": {},
"source": [
"## 3. Normalize Flow-Level Features (Transform Only)"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "97ef436c",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"TrickBot flow features (scaled): (10823, 22)\n"
]
}
],
"source": [
"import joblib\n",
"scaler = joblib.load(out_dir + r'\\flow_scaler_40k.pkl')\n",
"\n",
"X_mal_raw = df_features_mal.values.astype(np.float64)\n",
"log_cols_idx = list(range(0, 8))\n",
"flag_cols_idx = list(range(8, 22))\n",
"\n",
"X_mal_log = np.log1p(X_mal_raw[:, log_cols_idx])\n",
"X_mal_log_scaled = scaler.transform(X_mal_log) \n",
"\n",
"X_mal_flags = X_mal_raw[:, flag_cols_idx]\n",
"X_mal_scaled = np.hstack([X_mal_log_scaled, X_mal_flags]).astype(np.float32)\n",
"\n",
"np.save(out_dir + r'\\trickbot_features_FINAL.npy', X_mal_scaled)\n",
"print(f\"TrickBot flow features (scaled): {X_mal_scaled.shape}\")"
]
},
{
"cell_type": "markdown",
"id": "729993d4",
"metadata": {},
"source": [
"## 4. Build & Normalize PPI Sequences"
]
},
{
"cell_type": "code",
"execution_count": 8,
"id": "95b6e361",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"TrickBot PPI (normalized): (10823, 30, 3)\n"
]
}
],
"source": [
"SEQ_LEN = 30\n",
"\n",
"def build_ppi_sequence(directions_str, lengths_str, times_str):\n",
" dirs = parse_ppi_field(directions_str)\n",
" lens = parse_ppi_field(lengths_str)\n",
" times = parse_ppi_field(times_str)\n",
" n = min(len(dirs), len(lens), len(times))\n",
"\n",
" seq = np.zeros((SEQ_LEN, 3), dtype=np.float64)\n",
" t_seconds = []\n",
" for t in times[:n]:\n",
" try:\n",
" t_seconds.append(pd.to_datetime(t).timestamp())\n",
" except:\n",
" t_seconds.append(0.0)\n",
"\n",
" for i in range(min(n, SEQ_LEN)):\n",
" seq[i, 0] = 0.0 if i == 0 else (t_seconds[i] - t_seconds[i-1]) * 1000.0\n",
" seq[i, 1] = float(dirs[i])\n",
" seq[i, 2] = float(lens[i])\n",
" return seq\n",
"\n",
"ppi_trick_raw = np.stack([\n",
" build_ppi_sequence(r['PPI_PKT_DIRECTIONS'], r['PPI_PKT_LENGTHS'], r['PPI_PKT_TIMES'])\n",
" for _, r in df_clean.iterrows()\n",
"])\n",
"\n",
"ppi_norm_stats = joblib.load(out_dir + r'\\ppi_norm_stats_40k.pkl')\n",
"ppi_trick_final = ppi_trick_raw.copy().astype(np.float32)\n",
"\n",
"for col in [0, 2]:\n",
" vals = ppi_trick_final[:, :, col]\n",
" vals = np.log1p(np.abs(vals)) * np.sign(vals)\n",
" mean, std = ppi_norm_stats[col]\n",
" ppi_trick_final[:, :, col] = (vals - mean) / (std + 1e-8)\n",
"\n",
"np.save(out_dir + r'\\trickbot_ppi_FINAL.npy', ppi_trick_final)\n",
"print(f\"TrickBot PPI (normalized): {ppi_trick_final.shape}\")"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "aaa26bbb",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "2178b4cb",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "75aef49d",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "713a6a6b",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
models/fine-tuning.ipynb 0 → 100644
View file @ 532029c9
{
"cells": [
{
"cell_type": "markdown",
"id": "3428dd04",
"metadata": {},
"source": [
"# Fine-tuning — StackedLSTM_AE (L=8) on CTU-Normal-20\n",
"\n",
"## 1. Setup & Load Base Model\n"
]
},
{
"cell_type": "code",
"execution_count": 11,
"id": "787be4f3",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"الموديل الأصلي L=8 اتحمّل \n",
"norm_L shape: (7295027, 8, 3)\n",
"trick_L shape: (10823, 8, 3)\n"
]
}
],
"source": [
"import numpy as np\n",
"import torch\n",
"import torch.nn as nn\n",
"import polars as pl\n",
"import shutil\n",
"import random\n",
"from datetime import datetime\n",
"from torch.utils.data import DataLoader, TensorDataset\n",
"from sklearn.model_selection import train_test_split\n",
"from sklearn.metrics import roc_auc_score\n",
"\n",
"data_dir = r'D:\\project 5\\data'\n",
"device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')\n",
"SEED = 42\n",
"L = 8\n",
"\n",
"random.seed(SEED)\n",
"np.random.seed(SEED)\n",
"torch.manual_seed(SEED)\n",
"torch.cuda.manual_seed_all(SEED)\n",
"\n",
"\n",
"class StackedLSTM_AE(nn.Module):\n",
" def __init__(self, seq_len, n_features=3, hidden=32, latent=16, dropout_p=0.2):\n",
" super().__init__()\n",
" self.seq_len = seq_len\n",
" self.enc_lstm = nn.LSTM(n_features, hidden, 2, batch_first=True, dropout=dropout_p)\n",
" self.enc_norm = nn.LayerNorm(hidden)\n",
" self.to_latent = nn.Linear(hidden, latent)\n",
" self.from_latent = nn.Linear(latent, hidden)\n",
" self.dec_lstm = nn.LSTM(hidden, hidden, 2, batch_first=True, dropout=dropout_p)\n",
" self.dec_norm = nn.LayerNorm(hidden)\n",
" self.out = nn.Linear(hidden, n_features)\n",
"\n",
" def forward(self, x):\n",
" _, (h, _) = self.enc_lstm(x)\n",
" z = self.to_latent(self.enc_norm(h[-1]))\n",
" d = self.from_latent(z).unsqueeze(1).repeat(1, self.seq_len, 1)\n",
" d, _ = self.dec_lstm(d)\n",
" return self.out(self.dec_norm(d))\n",
"\n",
"\n",
"def recon_err(m, X, bs=4096):\n",
" m.eval()\n",
" e = []\n",
" with torch.no_grad():\n",
" for i in range(0, len(X), bs):\n",
" xb = torch.tensor(np.array(X[i:i+bs]), dtype=torch.float32, device=device)\n",
" e.append(((m(xb) - xb) ** 2).mean(dim=(1, 2)).cpu().numpy())\n",
" return np.concatenate(e)\n",
"\n",
"\n",
"model = StackedLSTM_AE(seq_len=L).to(device)\n",
"model.load_state_dict(torch.load(data_dir + f'\\\\lstm_trained_L{L}.pt', map_location=device))\n",
"print(f\"الموديل الأصلي L={L} اتحمّل \")\n",
"\n",
"ppi_norm = np.load(data_dir + r'\\normal_ppi_sequences_40k.npy', mmap_mode='r')\n",
"norm_L = np.array(ppi_norm[:, :L, :])\n",
"print(\"norm_L shape:\", norm_L.shape)\n",
"\n",
"ppi_trick = np.load(data_dir + r'\\trickbot_ppi_FINAL.npy')\n",
"trick_L = ppi_trick[:, :L, :]\n",
"print(\"trick_L shape:\", trick_L.shape)"
]
},
{
"cell_type": "markdown",
"id": "d6ed92f5",
"metadata": {},
"source": [
"## 2. Build CTU-Normal-20 Data — File 1 (2017-04-30)"
]
},
{
"cell_type": "code",
"execution_count": 2,
"id": "21961fd6",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"CTU-Normal-20 SYN=1: 2486\n",
"X_ctu20_full shape: (2458, 30, 3)\n"
]
}
],
"source": [
"import polars as pl\n",
"from datetime import datetime\n",
"\n",
"IAT_MEAN, IAT_STD = 2.965984344482422, 2.0991218090057373\n",
"SIZE_MEAN, SIZE_STD = 5.938459873199463, 1.4185206890106201\n",
"SYN_BIT = 0x02\n",
"SEQ_LEN_FULL = 30\n",
"\n",
"df_ctu = pl.read_csv(r'D:\\project 5\\ipfixprobe\\normal_test_pcap\\2017-04-30_win-normal.pcap.csv')\n",
"df_tls = df_ctu.filter(\n",
" (pl.col(\"uint8 PROTOCOL\") == 6) &\n",
" ((pl.col(\"uint16 DST_PORT\") == 443) | (pl.col(\"uint16 SRC_PORT\") == 443))\n",
")\n",
"df_ctu_syn1 = df_tls.filter(\n",
" (df_tls[\"uint8 TCP_FLAGS\"].cast(pl.Int64) & SYN_BIT) > 0\n",
")\n",
"print(\"CTU-Normal-20 SYN=1:\", df_ctu_syn1.height)\n",
"\n",
"def parse_pipe_ints(s):\n",
" if s is None or s == \"\":\n",
" return []\n",
" return [int(x) for x in s.strip(\"[]\").split(\"|\") if x != \"\"]\n",
"\n",
"def parse_pipe_times_ms(s):\n",
" if s is None or s == \"\":\n",
" return []\n",
" parts = [p for p in s.strip(\"[]\").split(\"|\") if p != \"\"]\n",
" ts = [datetime.fromisoformat(p) for p in parts]\n",
" iats = [0.0]\n",
" for i in range(1, len(ts)):\n",
" iats.append((ts[i] - ts[i-1]).total_seconds() * 1000.0)\n",
" return iats\n",
"\n",
"def build_normalized_matrix_correct(lengths, directions, iats_ms):\n",
" n = min(len(lengths), SEQ_LEN_FULL)\n",
" mat = np.zeros((SEQ_LEN_FULL, 3), dtype=np.float32)\n",
" for i in range(n):\n",
" iat_norm = (np.log1p(max(iats_ms[i], 0)) - IAT_MEAN) / IAT_STD\n",
" size_norm = (np.log1p(lengths[i]) - SIZE_MEAN) / SIZE_STD\n",
" mat[i] = [iat_norm, directions[i], size_norm]\n",
" if n < SEQ_LEN_FULL:\n",
" pad_iat = (np.log1p(0) - IAT_MEAN) / IAT_STD\n",
" pad_size = (np.log1p(0) - SIZE_MEAN) / SIZE_STD\n",
" mat[n:] = [pad_iat, 0.0, pad_size]\n",
" return mat\n",
"\n",
"X_ctu20_full = []\n",
"for row in df_ctu_syn1.iter_rows(named=True):\n",
" lengths = parse_pipe_ints(row[\"uint16* PPI_PKT_LENGTHS\"])\n",
" directions = parse_pipe_ints(row[\"int8* PPI_PKT_DIRECTIONS\"])\n",
" iats_ms = parse_pipe_times_ms(row[\"time* PPI_PKT_TIMES\"])\n",
" if len(lengths) == 0:\n",
" continue\n",
" X_ctu20_full.append(build_normalized_matrix_correct(lengths, directions, iats_ms))\n",
"\n",
"X_ctu20_full = np.array(X_ctu20_full)\n",
"print(\"X_ctu20_full shape:\", X_ctu20_full.shape)"
]
},
{
"cell_type": "markdown",
"id": "1b5860c8",
"metadata": {},
"source": [
"## 3. Build CTU-Normal-20 Data — File 2 (2017-04-25)"
]
},
{
"cell_type": "code",
"execution_count": 3,
"id": "52d0fbca",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"04-25 TLS + SYN=1: 886\n",
"X_ctu2_full shape: (874, 30, 3)\n",
"X_ctu2_L shape: (874, 8, 3)\n"
]
}
],
"source": [
"df_ctu2 = pl.read_csv(r'D:\\project 5\\ipfixprobe\\normal_test2_pcap\\2017-04-25_win-normal.pcap.csv')\n",
"df_tls2 = df_ctu2.filter(\n",
" (pl.col(\"uint8 PROTOCOL\") == 6) &\n",
" ((pl.col(\"uint16 DST_PORT\") == 443) | (pl.col(\"uint16 SRC_PORT\") == 443))\n",
")\n",
"df_ctu2_syn1 = df_tls2.filter(\n",
" (df_tls2[\"uint8 TCP_FLAGS\"].cast(pl.Int64) & SYN_BIT) > 0\n",
")\n",
"print(\"04-25 TLS + SYN=1:\", df_ctu2_syn1.height)\n",
"\n",
"X_ctu2_full = []\n",
"for row in df_ctu2_syn1.iter_rows(named=True):\n",
" lengths = parse_pipe_ints(row[\"uint16* PPI_PKT_LENGTHS\"])\n",
" directions = parse_pipe_ints(row[\"int8* PPI_PKT_DIRECTIONS\"])\n",
" iats_ms = parse_pipe_times_ms(row[\"time* PPI_PKT_TIMES\"])\n",
" if len(lengths) == 0:\n",
" continue\n",
" X_ctu2_full.append(build_normalized_matrix_correct(lengths, directions, iats_ms))\n",
"\n",
"X_ctu2_full = np.array(X_ctu2_full)\n",
"print(\"X_ctu2_full shape:\", X_ctu2_full.shape)\n",
"\n",
"X_ctu2_L = X_ctu2_full[:, :L, :]\n",
"print(\"X_ctu2_L shape:\", X_ctu2_L.shape)"
]
},
{
"cell_type": "markdown",
"id": "14cd788f",
"metadata": {},
"source": [
"## 4. Build Fine-tuning Mix (35% CTU-Normal-20 + 65% CESNET)"
]
},
{
"cell_type": "code",
"execution_count": 4,
"id": "d1c73750",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"X_ctu20_L shape: (2458, 8, 3)\n",
"CTU-20 للـ fine-tuning: 1720\n",
"CTU-20 للتيست النهائي: 738\n",
"إجمالي CTU-Normal الجديد: 2594\n",
"إجمالي عيّنة الـ fine-tuning: 7411\n",
"نسبة CTU-Normal فيها: 35.0 بالمية\n",
"TrickBot AUC قبل الـ fine-tuning: 0.994\n"
]
}
],
"source": [
"X_ctu20_L = X_ctu20_full[:, :L, :]\n",
"print(\"X_ctu20_L shape:\", X_ctu20_L.shape)\n",
"\n",
"ctu20_ft, ctu20_test = train_test_split(X_ctu20_L, test_size=0.3, random_state=SEED)\n",
"print(\"CTU-20 للـ fine-tuning:\", len(ctu20_ft))\n",
"print(\"CTU-20 للتيست النهائي:\", len(ctu20_test))\n",
"\n",
"n_new_1 = len(ctu20_ft) # من 04-30\n",
"n_new_2 = len(X_ctu2_L) # من 04-25\n",
"n_new_total = n_new_1 + n_new_2\n",
"print(\"إجمالي CTU-Normal الجديد:\", n_new_total)\n",
"\n",
"# 35%\n",
"target_ratio = 0.35\n",
"n_old = int(n_new_total * (1 - target_ratio) / target_ratio)\n",
"\n",
"rng = np.random.RandomState(SEED)\n",
"old_idx = rng.choice(len(norm_L), n_old, replace=False)\n",
"\n",
"X_finetune = np.concatenate([norm_L[old_idx], ctu20_ft, X_ctu2_L], axis=0).astype(np.float32)\n",
"np.random.RandomState(SEED).shuffle(X_finetune)\n",
"\n",
"print(\"إجمالي عيّنة الـ fine-tuning:\", len(X_finetune))\n",
"print(\"نسبة CTU-Normal فيها:\", round(n_new_total / len(X_finetune) * 100, 1), \"بالمية\")\n",
"\n",
"monitor_norm_idx = rng.choice(len(norm_L), len(trick_L), replace=False)\n",
"monitor_norm = norm_L[monitor_norm_idx]\n",
"\n",
"def check_trickbot_auc(m):\n",
" e_norm = recon_err(m, monitor_norm)\n",
" e_trick = recon_err(m, trick_L)\n",
" y = np.concatenate([np.zeros(len(e_norm)), np.ones(len(e_trick))])\n",
" s = np.concatenate([e_norm, e_trick])\n",
" return roc_auc_score(y, s)\n",
"\n",
"auc_before = check_trickbot_auc(model)\n",
"print(\"TrickBot AUC قبل الـ fine-tuning:\", round(auc_before, 4))"
]
},
{
"cell_type": "markdown",
"id": "624ba9ba",
"metadata": {},
"source": [
"## 5. Fine-tuning Attempt 1 — lr=5e-5, weight_decay=1e-4 (Unsuccessful)"
]
},
{
"cell_type": "code",
"execution_count": 6,
"id": "acd038d5",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 1 | Loss: 0.01994 | TrickBot AUC: 0.993\n",
"Epoch 2 | Loss: 0.01801 | TrickBot AUC: 0.9934\n",
"Epoch 3 | Loss: 0.01724 | TrickBot AUC: 0.9931\n",
"Epoch 4 | Loss: 0.01644 | TrickBot AUC: 0.9933\n",
"Epoch 5 | Loss: 0.01606 | TrickBot AUC: 0.9934\n",
"Epoch 6 | Loss: 0.01547 | TrickBot AUC: 0.9934\n",
"Epoch 7 | Loss: 0.01522 | TrickBot AUC: 0.9934\n",
"Epoch 8 | Loss: 0.01488 | TrickBot AUC: 0.994\n",
"Epoch 9 | Loss: 0.01463 | TrickBot AUC: 0.9935\n",
"Epoch 10 | Loss: 0.01462 | TrickBot AUC: 0.9935\n",
"Epoch 11 | Loss: 0.01434 | TrickBot AUC: 0.9932\n",
"Epoch 12 | Loss: 0.01405 | TrickBot AUC: 0.9939\n",
"Epoch 13 | Loss: 0.01416 | TrickBot AUC: 0.9934\n",
"Epoch 14 | Loss: 0.01372 | TrickBot AUC: 0.9931\n",
"Epoch 15 | Loss: 0.01387 | TrickBot AUC: 0.9932\n",
"Epoch 16 | Loss: 0.01351 | TrickBot AUC: 0.9933\n",
"Epoch 17 | Loss: 0.01343 | TrickBot AUC: 0.9934\n",
"Epoch 18 | Loss: 0.01356 | TrickBot AUC: 0.9931\n",
"Epoch 19 | Loss: 0.01333 | TrickBot AUC: 0.9933\n",
"Epoch 20 | Loss: 0.01326 | TrickBot AUC: 0.9934\n",
"Epoch 21 | Loss: 0.01307 | TrickBot AUC: 0.9927\n",
"Epoch 22 | Loss: 0.01304 | TrickBot AUC: 0.9929\n",
"Epoch 23 | Loss: 0.01305 | TrickBot AUC: 0.9933\n",
"Epoch 24 | Loss: 0.01293 | TrickBot AUC: 0.9931\n",
"Epoch 25 | Loss: 0.01281 | TrickBot AUC: 0.9929\n",
"Epoch 26 | Loss: 0.01284 | TrickBot AUC: 0.993\n",
"Epoch 27 | Loss: 0.01275 | TrickBot AUC: 0.9929\n",
"Epoch 28 | Loss: 0.01285 | TrickBot AUC: 0.9935\n",
"Epoch 29 | Loss: 0.01305 | TrickBot AUC: 0.9931\n",
"Epoch 30 | Loss: 0.0128 | TrickBot AUC: 0.9931\n",
"Epoch 31 | Loss: 0.01271 | TrickBot AUC: 0.9933\n",
"Epoch 32 | Loss: 0.01273 | TrickBot AUC: 0.9934\n",
"Epoch 33 | Loss: 0.01254 | TrickBot AUC: 0.9934\n",
"Epoch 34 | Loss: 0.01252 | TrickBot AUC: 0.993\n",
"Epoch 35 | Loss: 0.01266 | TrickBot AUC: 0.9934\n",
"Epoch 36 | Loss: 0.01245 | TrickBot AUC: 0.993\n",
"Epoch 37 | Loss: 0.01247 | TrickBot AUC: 0.9931\n",
"Epoch 38 | Loss: 0.01223 | TrickBot AUC: 0.9933\n",
"Epoch 39 | Loss: 0.01234 | TrickBot AUC: 0.9929\n",
"Epoch 40 | Loss: 0.01244 | TrickBot AUC: 0.9932\n",
"الموديل المحسّن - أفضل AUC وصلها: 0.994\n"
]
}
],
"source": [
"opt = torch.optim.AdamW(model.parameters(), lr=5e-5, weight_decay=1e-4)\n",
"crit = nn.MSELoss()\n",
"loader = DataLoader(TensorDataset(torch.tensor(X_finetune)), batch_size=256, shuffle=True)\n",
"\n",
"MIN_ACCEPTABLE_AUC = 0.97\n",
"best_state = {k: v.clone() for k, v in model.state_dict().items()}\n",
"best_auc = auc_before\n",
"\n",
"for epoch in range(40):\n",
" model.train()\n",
" total_loss = 0\n",
" for (xb,) in loader:\n",
" xb = xb.to(device)\n",
" opt.zero_grad()\n",
" loss = crit(model(xb), xb)\n",
" loss.backward()\n",
" opt.step()\n",
" total_loss += loss.item()\n",
"\n",
" auc = check_trickbot_auc(model)\n",
" print(\"Epoch\", epoch + 1, \"| Loss:\", round(total_loss / len(loader), 5), \"| TrickBot AUC:\", round(auc, 4))\n",
"\n",
" if auc < MIN_ACCEPTABLE_AUC:\n",
" print(\"AUC نزل تحت\", MIN_ACCEPTABLE_AUC)\n",
" model.load_state_dict(best_state)\n",
" break\n",
" if auc >= best_auc:\n",
" best_auc = auc\n",
" best_state = {k: v.clone() for k, v in model.state_dict().items()}\n",
"\n",
"model.load_state_dict(best_state)\n",
"torch.save(model.state_dict(), data_dir + f'\\\\lstm_trained_L{L}_finetuned.pt')\n",
"print(\"الموديل المحسّن - أفضل AUC وصلها:\", round(best_auc, 4))"
]
},
{
"cell_type": "markdown",
"id": "f1fe7158",
"metadata": {},
"source": [
"## 6. Evaluate Attempt 1 — FP Rate on Held-out CTU-20"
]
},
{
"cell_type": "code",
"execution_count": 7,
"id": "8fb84c11",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"العتبة الجديدة بعد الـ fine-tuning: 0.014\n",
"FP على CTU-20 (held-out, 738 تدفق): 34.0 بالمية\n",
"كانت قبل الـ fine-tuning: 48.9 بالمية\n",
"TrickBot AUC النهائي: 0.994 كانت قبل: 0.994\n"
]
}
],
"source": [
"thr_new = np.percentile(recon_err(model, monitor_norm), 97)\n",
"err_test = recon_err(model, ctu20_test)\n",
"fp_rate = (err_test > thr_new).mean()\n",
"\n",
"print(\"العتبة الجديدة بعد الـ fine-tuning:\", round(thr_new, 4))\n",
"print(\"FP على CTU-20 (held-out, 738 تدفق):\", round(fp_rate * 100, 1), \"بالمية\")\n",
"print(\"كانت قبل الـ fine-tuning: 48.9 بالمية\")\n",
"print(\"TrickBot AUC النهائي:\", round(best_auc, 4), \" كانت قبل:\", round(auc_before, 4))"
]
},
{
"cell_type": "markdown",
"id": "88b8725c",
"metadata": {},
"source": [
"## 8. Fine-tuning Attempt 2 — lr=3e-4, weight_decay=1e-6 (Successful)"
]
},
{
"cell_type": "code",
"execution_count": 12,
"id": "c8836238",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"Epoch 1 | Loss: 0.01125 | TrickBot AUC: 0.993\n",
"Epoch 2 | Loss: 0.01117 | TrickBot AUC: 0.9935\n",
"Epoch 3 | Loss: 0.01098 | TrickBot AUC: 0.994\n",
"Epoch 4 | Loss: 0.01095 | TrickBot AUC: 0.9922\n",
"Epoch 5 | Loss: 0.01088 | TrickBot AUC: 0.9933\n",
"Epoch 6 | Loss: 0.01085 | TrickBot AUC: 0.9922\n",
"Epoch 7 | Loss: 0.0109 | TrickBot AUC: 0.993\n",
"Epoch 8 | Loss: 0.01089 | TrickBot AUC: 0.9922\n",
"Epoch 9 | Loss: 0.01074 | TrickBot AUC: 0.9927\n",
"Epoch 10 | Loss: 0.01069 | TrickBot AUC: 0.9927\n",
"Epoch 11 | Loss: 0.01066 | TrickBot AUC: 0.9933\n",
"Epoch 12 | Loss: 0.01175 | TrickBot AUC: 0.9911\n",
"Epoch 13 | Loss: 0.01107 | TrickBot AUC: 0.9922\n",
"Epoch 14 | Loss: 0.01097 | TrickBot AUC: 0.9926\n",
"Epoch 15 | Loss: 0.01056 | TrickBot AUC: 0.9923\n",
"Epoch 16 | Loss: 0.01067 | TrickBot AUC: 0.9918\n",
"Epoch 17 | Loss: 0.01044 | TrickBot AUC: 0.9927\n",
"Epoch 18 | Loss: 0.01052 | TrickBot AUC: 0.9922\n",
"Epoch 19 | Loss: 0.01068 | TrickBot AUC: 0.9931\n",
"Epoch 20 | Loss: 0.01056 | TrickBot AUC: 0.9915\n",
"Epoch 21 | Loss: 0.01054 | TrickBot AUC: 0.9921\n",
"Epoch 22 | Loss: 0.0105 | TrickBot AUC: 0.9926\n",
"Epoch 23 | Loss: 0.01043 | TrickBot AUC: 0.9917\n",
"Epoch 24 | Loss: 0.01069 | TrickBot AUC: 0.9932\n",
"Epoch 25 | Loss: 0.01069 | TrickBot AUC: 0.9917\n",
"Epoch 26 | Loss: 0.01048 | TrickBot AUC: 0.992\n",
"Epoch 27 | Loss: 0.01045 | TrickBot AUC: 0.9918\n",
"Epoch 28 | Loss: 0.0105 | TrickBot AUC: 0.9933\n",
"Epoch 29 | Loss: 0.01036 | TrickBot AUC: 0.992\n",
"Epoch 30 | Loss: 0.01036 | TrickBot AUC: 0.9922\n",
"Epoch 31 | Loss: 0.01043 | TrickBot AUC: 0.9932\n",
"Epoch 32 | Loss: 0.01033 | TrickBot AUC: 0.9924\n",
"Epoch 33 | Loss: 0.01048 | TrickBot AUC: 0.9917\n",
"Epoch 34 | Loss: 0.01052 | TrickBot AUC: 0.9917\n",
"Epoch 35 | Loss: 0.01033 | TrickBot AUC: 0.9918\n",
"Epoch 36 | Loss: 0.01034 | TrickBot AUC: 0.9925\n",
"Epoch 37 | Loss: 0.0104 | TrickBot AUC: 0.9917\n",
"Epoch 38 | Loss: 0.01045 | TrickBot AUC: 0.9909\n",
"Epoch 39 | Loss: 0.01038 | TrickBot AUC: 0.9919\n",
"Epoch 40 | Loss: 0.01041 | TrickBot AUC: 0.9931\n",
"الموديل - أفضل AUC وصل لها: 0.994\n"
]
}
],
"source": [
"opt = torch.optim.AdamW(model.parameters(), lr=3e-4, weight_decay=1e-6)\n",
"crit = nn.MSELoss()\n",
"loader = DataLoader(TensorDataset(torch.tensor(X_finetune)), batch_size=256, shuffle=True)\n",
"\n",
"MIN_ACCEPTABLE_AUC = 0.97\n",
"best_state = {k: v.clone() for k, v in model.state_dict().items()}\n",
"best_auc = auc_before\n",
"\n",
"for epoch in range(40):\n",
" model.train()\n",
" total_loss = 0\n",
" for (xb,) in loader:\n",
" xb = xb.to(device)\n",
" opt.zero_grad()\n",
" loss = crit(model(xb), xb)\n",
" loss.backward()\n",
" opt.step()\n",
" total_loss += loss.item()\n",
"\n",
" auc = check_trickbot_auc(model)\n",
" print(\"Epoch\", epoch + 1, \"| Loss:\", round(total_loss / len(loader), 5), \"| TrickBot AUC:\", round(auc, 4))\n",
"\n",
" if auc < MIN_ACCEPTABLE_AUC:\n",
" print(\"AUC نزل تحت\", MIN_ACCEPTABLE_AUC, )\n",
" model.load_state_dict(best_state)\n",
" break\n",
" if auc >= best_auc:\n",
" best_auc = auc\n",
" best_state = {k: v.clone() for k, v in model.state_dict().items()}\n",
"\n",
"model.load_state_dict(best_state)\n",
"torch.save(model.state_dict(), data_dir + f'\\\\lstm_trained_L{L}_finetuned_lr3e4.pt')\n",
"print(\"الموديل - أفضل AUC وصل لها:\", round(best_auc, 4))"
]
},
{
"cell_type": "markdown",
"id": "2dda7347",
"metadata": {},
"source": [
"## 9. Evaluate & Confirm Final "
]
},
{
"cell_type": "code",
"execution_count": 16,
"id": "a3b6c492",
"metadata": {},
"outputs": [
{
"name": "stdout",
"output_type": "stream",
"text": [
"العتبة النهائية: 0.0155\n",
"FP على CTU-20 (held-out, 738 تدفق): 12.1 بالمية\n",
"TrickBot AUC النهائي: 0.994 كانت قبل: 0.994\n"
]
}
],
"source": [
"thr_new = np.percentile(recon_err(model, monitor_norm), 97)\n",
"err_test = recon_err(model, ctu20_test)\n",
"fp_rate = (err_test > thr_new).mean()\n",
"\n",
"print(\"العتبة النهائية:\", round(thr_new, 4))\n",
"print(\"FP على CTU-20 (held-out, 738 تدفق):\", round(fp_rate * 100, 1), \"بالمية\")\n",
"print(\"TrickBot AUC النهائي:\", round(best_auc, 4), \" كانت قبل:\", round(auc_before, 4))"
]
},
{
"cell_type": "code",
"execution_count": null,
"id": "efb7a2cc",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "50f67c71",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "75f2e38c",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "6da57fd3",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "d4145554",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "d3a63e85",
"metadata": {},
"outputs": [],
"source": []
},
{
"cell_type": "code",
"execution_count": null,
"id": "bf47253e",
"metadata": {},
"outputs": [],
"source": []
}
],
"metadata": {
"kernelspec": {
"display_name": "Python 3 (ipykernel)",
"language": "python",
"name": "python3"
},
"language_info": {
"codemirror_mode": {
"name": "ipython",
"version": 3
},
"file_extension": ".py",
"mimetype": "text/x-python",
"name": "python",
"nbconvert_exporter": "python",
"pygments_lexer": "ipython3",
"version": "3.11.5"
}
},
"nbformat": 4,
"nbformat_minor": 5
}
nids-system-code-final/api-service/Dockerfile 0 → 100644
View file @ 532029c9
FROM python:3.11-slim
WORKDIR /app
RUN pip install --no-cache-dir fastapi "uvicorn[standard]"
COPY main.py .
EXPOSE 8000
CMD ["uvicorn", "main:app", "--host", "0.0.0.0", "--port", "8000"]
nids-system-code-final/api-service/main.py 0 → 100644
View file @ 532029c9
from fastapi import FastAPI
import sqlite3
DB_PATH = "/shared_data/detections.db"
app = FastAPI(title="NIDS Detections API")
def get_connection():
conn = sqlite3.connect(DB_PATH)
conn.row_factory = sqlite3.Row
return conn
@app.get("/health")
def health():
return {"status": "ok"}
@app.get("/detections")
def get_detections():
"""Return all detection rows, oldest first (same shape as the old direct SQLite read)."""
conn = get_connection()
cursor = conn.cursor()
cursor.execute("SELECT * FROM detections ORDER BY id ASC")
rows = [dict(row) for row in cursor.fetchall()]
conn.close()
return rows
@app.get("/active_flows")
def get_active_flows():
"""Return the current in-progress flow count."""
conn = get_connection()
cursor = conn.cursor()
cursor.execute("SELECT active_flows_count FROM system_status WHERE id = 1")
row = cursor.fetchone()
conn.close()
return {"active_flows_count": row["active_flows_count"] if row else 0}
nids-system-code-final/capture-service/Dockerfile 0 → 100644
View file @ 532029c9
FROM python:3.11-slim
RUN apt-get update && apt-get install -y --no-install-recommends \
libpcap-dev \
tcpdump \
&& rm -rf /var/lib/apt/lists/*
WORKDIR /app
RUN pip install --no-cache-dir scapy joblib numpy pytest \
--index-url https://download.pytorch.org/whl/cpu torch \
--extra-index-url https://pypi.org/simple
COPY sniffer.py flow_aggregator.py inference_engine.py db_writer.py ./
COPY tests/ ./tests/
COPY models/ ./models/
CMD ["python3", "-u", "flow_aggregator.py"]
nids-system-code-final/capture-service/db_writer.py 0 → 100644
View file @ 532029c9
import sqlite3
from datetime import datetime
DB_PATH = "/shared_data/detections.db"
def init_db():
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
# WAL mode
cursor.execute("PRAGMA journal_mode=WAL;")
cursor.execute("""
CREATE TABLE IF NOT EXISTS detections (
id INTEGER PRIMARY KEY AUTOINCREMENT,
timestamp TEXT NOT NULL,
src_ip TEXT NOT NULL,
dst_ip TEXT NOT NULL,
dst_port INTEGER,
stage TEXT NOT NULL,
decision TEXT NOT NULL,
reconstruction_error REAL NOT NULL,
threshold REAL NOT NULL,
processing_latency_ms REAL,
flow_duration_sec REAL
)
""")
# Migration:
try:
cursor.execute("ALTER TABLE detections ADD COLUMN processing_latency_ms REAL")
except sqlite3.OperationalError:
pass
try:
cursor.execute("ALTER TABLE detections ADD COLUMN flow_duration_sec REAL")
except sqlite3.OperationalError:
pass
cursor.execute("""
CREATE TABLE IF NOT EXISTS system_status (
id INTEGER PRIMARY KEY CHECK (id = 1),
active_flows_count INTEGER NOT NULL,
last_updated TEXT NOT NULL
)
""")
cursor.execute("""
INSERT OR IGNORE INTO system_status (id, active_flows_count, last_updated)
VALUES (1, 0, ?)
""", (datetime.now().isoformat(),))
conn.commit()
conn.close()
print("[DB] Database initialized successfully.")
def log_detection(src_ip, dst_ip, dst_port, stage, decision, reconstruction_error, threshold, processing_latency_ms=None, flow_duration_sec=None):
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("""
INSERT INTO detections
(timestamp, src_ip, dst_ip, dst_port, stage, decision, reconstruction_error, threshold, processing_latency_ms, flow_duration_sec)
VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?)
""", (
datetime.now().isoformat(),
src_ip,
dst_ip,
dst_port,
stage,
decision,
float(reconstruction_error),
float(threshold),
float(processing_latency_ms) if processing_latency_ms is not None else None,
float(flow_duration_sec) if flow_duration_sec is not None else None
))
conn.commit()
conn.close()
def update_active_flow_count(count):
conn = sqlite3.connect(DB_PATH)
cursor = conn.cursor()
cursor.execute("""
UPDATE system_status
SET active_flows_count = ?, last_updated = ?
WHERE id = 1
""", (count, datetime.now().isoformat()))
conn.commit()
conn.close()
\ No newline at end of file
nids-system-code-final/capture-service/flow_aggregator.py 0 → 100644
View file @ 532029c9
import queue
import threading
import time
import numpy as np
from db_writer import init_db, log_detection, update_active_flow_count
from inference_engine import classify_flow, THRESHOLD
SEQ_LEN = 8
FLOW_TIMEOUT = 30
REPEAT_LOG_INTERVAL_SEC = 5
active_flows = {}
flows_lock = threading.Lock()
def get_flow_key(packet):
if not packet.haslayer("IP") or not packet.haslayer("TCP"):
return None
ip_layer = packet["IP"]
tcp_layer = packet["TCP"]
src_ip = ip_layer.src
dst_ip = ip_layer.dst
src_port = tcp_layer.sport
dst_port = tcp_layer.dport
if (src_ip, src_port) < (dst_ip, dst_port):
forward = True
key = (src_ip, dst_ip, src_port, dst_port)
else:
forward = False
key = (dst_ip, src_ip, dst_port, src_port)
return key, forward
def create_new_flow_state(timestamp, forward):
return {
"first_seen": timestamp,
"last_seen": timestamp,
"bytes_fwd": 0,
"bytes_rev": 0,
"packets_fwd": 0,
"packets_rev": 0,
"ppi_sizes": [],
"ppi_directions": [],
"ppi_times": [],
"tcp_flags_fwd": 0,
"tcp_flags_rev": 0,
"final_features_sent": False,
"client_is_forward": forward,
"last_decision": None,
"last_error": None,
"last_repeat_logged_at": 0.0
}
def compute_ppi_matrix(flow):
seq = np.zeros((SEQ_LEN, 3), dtype=np.float64)
sizes = flow["ppi_sizes"]
directions = flow["ppi_directions"]
times = flow["ppi_times"]
n = len(sizes)
for i in range(n):
if i == 0:
seq[i, 0] = 0.0
else:
seq[i, 0] = (times[i] - times[i - 1]) * 1000.0
seq[i, 1] = directions[i]
seq[i, 2] = sizes[i]
return seq
def update_flow(packet):
result = get_flow_key(packet)
if result is None:
return None
key, forward = result
timestamp = float(packet.time)
packet_size = len(packet)
with flows_lock:
if key not in active_flows:
active_flows[key] = create_new_flow_state(timestamp, forward)
flow = active_flows[key]
flow["last_seen"] = timestamp
if forward:
flow["bytes_fwd"] += packet_size
flow["packets_fwd"] += 1
else:
flow["bytes_rev"] += packet_size
flow["packets_rev"] += 1
if len(flow["ppi_sizes"]) < SEQ_LEN:
flow["ppi_sizes"].append(packet_size)
is_client_packet = (forward == flow["client_is_forward"])
flow["ppi_directions"].append(1 if is_client_packet else -1)
flow["ppi_times"].append(timestamp)
tcp_layer = packet["TCP"]
flags_int = int(tcp_layer.flags)
if forward:
flow["tcp_flags_fwd"] |= flags_int
else:
flow["tcp_flags_rev"] |= flags_int
return key
def check_flow_completion(key, packet):
with flows_lock:
if key not in active_flows:
return
flow = active_flows[key]
ppi_count = len(flow["ppi_sizes"])
# key = (ip_a, ip_b, port_a, port_b) بترتيب موحّد - منستخدمها للعرض بس
src_ip, dst_ip, _, dst_port = key
# ---------- تدفق سبق وتصنّف: نسجل تكرار برأيه القديم بدل ما نتجاهله بصمت ----------
# محدود بفاصل زمني (REPEAT_LOG_INTERVAL_SEC) لتفادي إغراق قاعدة البيانات
if flow["final_features_sent"]:
now = flow["last_seen"]
if now - flow["last_repeat_logged_at"] >= REPEAT_LOG_INTERVAL_SEC:
flow["last_repeat_logged_at"] = now
flow_duration_sec = flow["last_seen"] - flow["first_seen"]
log_detection(
src_ip=src_ip,
dst_ip=dst_ip,
dst_port=dst_port,
stage="REPEATED",
decision=flow["last_decision"],
reconstruction_error=flow["last_error"],
threshold=THRESHOLD,
processing_latency_ms=None,
flow_duration_sec=flow_duration_sec
)
return
if ppi_count == SEQ_LEN and not flow["final_features_sent"]:
flow["final_features_sent"] = True
ppi_matrix = compute_ppi_matrix(flow)
decision, error, latency_ms = classify_flow(ppi_matrix)
flow["last_decision"] = decision
flow["last_error"] = error
flow_duration_sec = flow["last_seen"] - flow["first_seen"]
print(f"[FINAL DETECTION] Flow {key} reached {SEQ_LEN} packets (full window)")
print(f" Decision: {decision} | Reconstruction Error: {error:.4f} | Latency: {latency_ms:.3f}ms")
print(f" FULL RAW MATRIX:\n{ppi_matrix}")
log_detection(
src_ip=src_ip,
dst_ip=dst_ip,
dst_port=dst_port,
stage="FINAL_CONFIRMED",
decision=decision,
reconstruction_error=error,
threshold=THRESHOLD,
processing_latency_ms=latency_ms,
flow_duration_sec=flow_duration_sec
)
tcp_layer = packet["TCP"]
flags_str = str(tcp_layer.flags)
if ("F" in flags_str or "R" in flags_str) and not flow["final_features_sent"]:
flow["final_features_sent"] = True
ppi_matrix = compute_ppi_matrix(flow)
decision, error, latency_ms = classify_flow(ppi_matrix)
flow["last_decision"] = decision
flow["last_error"] = error
flow_duration_sec = flow["last_seen"] - flow["first_seen"]
print(f"[FINAL DETECTION] Flow {key} ended early via FIN/RST at {ppi_count} packets")
print(f" Decision: {decision} | Reconstruction Error: {error:.4f} | Latency: {latency_ms:.3f}ms")
log_detection(
src_ip=src_ip,
dst_ip=dst_ip,
dst_port=dst_port,
stage="FINAL_CONFIRMED",
decision=decision,
reconstruction_error=error,
threshold=THRESHOLD,
processing_latency_ms=latency_ms,
flow_duration_sec=flow_duration_sec
)
del active_flows[key]
def check_timeouts():
current_time = time.time()
with flows_lock:
expired_keys = []
for key, flow in active_flows.items():
if current_time - flow["last_seen"] > FLOW_TIMEOUT:
expired_keys.append(key)
for key in expired_keys:
flow = active_flows[key]
if not flow["final_features_sent"]:
ppi_matrix = compute_ppi_matrix(flow)
ppi_count = len(flow["ppi_sizes"])
print(f"[FINAL DETECTION] Flow {key} timed out at {ppi_count} packets")
print(f" PPI matrix shape: {ppi_matrix.shape}")
print(f" First 3 rows:\n{ppi_matrix[:3]}")
decision, error, latency_ms = classify_flow(ppi_matrix)
flow["last_decision"] = decision
flow["last_error"] = error
flow_duration_sec = flow["last_seen"] - flow["first_seen"]
src_ip, dst_ip, _, dst_port = key
log_detection(
src_ip=src_ip,
dst_ip=dst_ip,
dst_port=dst_port,
stage="FINAL_CONFIRMED",
decision=decision,
reconstruction_error=error,
threshold=THRESHOLD,
processing_latency_ms=latency_ms,
flow_duration_sec=flow_duration_sec
)
del active_flows[key]
def process_packets(packet_queue, stop_event):
while not stop_event.is_set():
try:
packet = packet_queue.get(timeout=1)
except queue.Empty:
continue
try:
key = update_flow(packet)
if key is not None:
check_flow_completion(key, packet)
except Exception as e:
print(f"[ERROR] Failed to process packet: {e}")
def timeout_checker_loop(stop_event):
while not stop_event.is_set():
time.sleep(1)
try:
check_timeouts()
with flows_lock:
current_count = len(active_flows)
update_active_flow_count(current_count)
except Exception as e:
print(f"[ERROR] Timeout checker failed: {e}")
if __name__ == "__main__":
from sniffer import sniffer, packet_queue
init_db()
stop_event = threading.Event()
processing_thread = threading.Thread(
target=process_packets,
args=(packet_queue, stop_event)
)
processing_thread.start()
timeout_thread = threading.Thread(
target=timeout_checker_loop,
args=(stop_event,)
)
timeout_thread.start()
print("[*] Starting live packet capture and flow aggregation...")
sniffer.start()
try:
while True:
time.sleep(1)
except KeyboardInterrupt:
print("\n[*] Stopping...")
sniffer.stop()
stop_event.set()
processing_thread.join()
timeout_thread.join()
print("[*] Stopped cleanly.")
nids-system-code-final/capture-service/inference_engine.py 0 → 100644
View file @ 532029c9
import torch
import torch.nn as nn
import numpy as np
import joblib
import time
MODEL_PATH = "/app/models/lstm_trained_L8_finetuned_FINAL_12pct.pt"
NORM_STATS_PATH = "/app/models/ppi_norm_stats_40k.pkl"
THRESHOLD = 0.0155
=class StackedLSTM_AE(nn.Module):
def __init__(self, n_features=3, seq_len=30, hidden=32, latent=16, dropout_p=0.2):
super().__init__()
self.seq_len = seq_len
self.enc_lstm = nn.LSTM(n_features, hidden, num_layers=2, batch_first=True, dropout=dropout_p)
self.enc_norm = nn.LayerNorm(hidden)
self.to_latent = nn.Linear(hidden, latent)
self.from_latent = nn.Linear(latent, hidden)
self.dec_lstm = nn.LSTM(hidden, hidden, num_layers=2, batch_first=True, dropout=dropout_p)
self.dec_norm = nn.LayerNorm(hidden)
self.out = nn.Linear(hidden, n_features)
def forward(self, x):
_, (h, _) = self.enc_lstm(x)
h_last = self.enc_norm(h[-1])
z = self.to_latent(h_last)
d = self.from_latent(z).unsqueeze(1).repeat(1, self.seq_len, 1)
d, _ = self.dec_lstm(d)
d = self.dec_norm(d)
return self.out(d)
=device = torch.device("cpu")
model = StackedLSTM_AE(seq_len=8).to(device)
model.load_state_dict(torch.load(MODEL_PATH, map_location=device))
model.eval()
norm_stats = joblib.load(NORM_STATS_PATH)
MEAN_TIMING, STD_TIMING = norm_stats[0]
MEAN_SIZE, STD_SIZE = norm_stats[2]
print(" Inference Engine: model and scaler loaded successfully")
def preprocess_ppi_matrix(raw_matrix):
processed = raw_matrix.copy()
processed[:, 0] = np.log1p(processed[:, 0])
processed[:, 2] = np.log1p(processed[:, 2])
processed[:, 0] = (processed[:, 0] - MEAN_TIMING) / STD_TIMING
processed[:, 2] = (processed[:, 2] - MEAN_SIZE) / STD_SIZE
return processed
def classify_flow(raw_ppi_matrix):
start_time = time.perf_counter()
processed = preprocess_ppi_matrix(raw_ppi_matrix)
input_tensor = torch.tensor(processed, dtype=torch.float32).unsqueeze(0).to(device)
with torch.no_grad():
reconstructed = model(input_tensor)
error = torch.mean((reconstructed - input_tensor) ** 2).item()
decision = "MALICIOUS" if error > THRESHOLD else "BENIGN"
latency_ms = (time.perf_counter() - start_time) * 1000
return decision, error, latency_ms
nids-system-code-final/capture-service/sniffer.py 0 → 100644
View file @ 532029c9
from scapy.all import AsyncSniffer
import queue
import threading
packet_queue = queue.Queue()
def packet_callback(packet):
print(f"[+] Captured packet: {packet.summary()}")
packet_queue.put(packet)
sniffer = AsyncSniffer(
iface="br-12fb6e5258fa",
filter="tcp port 443",
prn=packet_callback,
store=False
)
if __name__ == "__main__":
print("[*] Starting live packet capture on ens33...")
sniffer.start()
try:
while True:
pass
except KeyboardInterrupt:
print("\n[*] Stopping capture...")
sniffer.stop()
print("[*] Capture stopped.")
nids-system-code-final/capture-service/tests/conftest.py 0 → 100644
View file @ 532029c9
import sys
import os
# Add the parent project directory (/app) to the search path so pytest
# can import flow_aggregator and inference_engine.
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
nids-system-code-final/capture-service/tests/load_test_synthetic.py 0 → 100644
View file @ 532029c9
"""
Synthetic load test for the NIDS inference pipeline.
Simulates increasing numbers of concurrent flows hitting classify_flow()
and log_detection() directly (no real network traffic needed), to measure
how latency and throughput behave under increasing pressure.
The test doubles the concurrency each phase (10 -> 20 -> 40 -> ...) until
an actual failure occurs (a real exception while processing a flow) or a
safety cap is reached.
Usage (inside the capture-service container):
python3 tests/load_test_synthetic.py
"""
import random
import sys
import os
import time
from concurrent.futures import ThreadPoolExecutor, as_completed
import numpy as np
sys.path.insert(0, os.path.join(os.path.dirname(__file__), ".."))
import db_writer
import inference_engine as ie
db_writer.init_db()
SEQ_LEN = ie.model.seq_len
# Roughly realistic ranges for TLS packet size / inter-arrival time
SIZE_RANGE = (60, 1500) # bytes, typical range for TLS packets
IAT_RANGE_MS = (0.5, 300) # milliseconds, typical gap between packets
# Safety cap to avoid an infinite loop if the system never actually fails
MAX_CONCURRENT_FLOWS = 20000
def build_random_ppi_matrix():
"""Build one synthetic but realistically-shaped PPI matrix."""
matrix = np.zeros((SEQ_LEN, 3), dtype=np.float64)
n_packets = random.randint(3, SEQ_LEN) # realistic flow length (some flows shorter than SEQ_LEN)
for i in range(n_packets):
matrix[i, 0] = 0.0 if i == 0 else random.uniform(*IAT_RANGE_MS)
matrix[i, 1] = random.choice([1, -1])
matrix[i, 2] = random.uniform(*SIZE_RANGE)
return matrix
def simulate_one_flow(flow_index):
"""Classify one synthetic flow and log it, returning the measured latency."""
matrix = build_random_ppi_matrix()
decision, error, latency_ms = ie.classify_flow(matrix)
db_writer.log_detection(
src_ip=f"198.51.100.{flow_index % 255}", # reserved TEST-NET-2 range (RFC 5737)
dst_ip="203.0.113.1",
dst_port=443,
stage="SYNTHETIC_LOAD_TEST",
decision=decision,
reconstruction_error=error,
threshold=ie.THRESHOLD,
processing_latency_ms=latency_ms,
flow_duration_sec=random.uniform(0.1, 5.0),
)
return latency_ms
def run_phase(name, n_concurrent):
print("=" * 70)
print(f"PHASE: {name}")
print("=" * 70)
start = time.perf_counter()
latencies = []
failures = 0
with ThreadPoolExecutor(max_workers=min(n_concurrent, 64)) as executor:
futures = [executor.submit(simulate_one_flow, i) for i in range(n_concurrent)]
for future in as_completed(futures):
try:
latencies.append(future.result())
except Exception as e:
failures += 1
print(f" ERROR processing flow: {e}")
total_time = time.perf_counter() - start
throughput = n_concurrent / total_time if total_time > 0 else 0
latencies.sort()
n = len(latencies)
print(f"Flows submitted: {n_concurrent} | Succeeded: {n} | Failed: {failures}")
print(f"Total wall time: {total_time:.3f}s | Throughput: {throughput:.1f} flows/sec")
if n > 0:
print(
f"Per-flow latency (ms) -> min: {latencies[0]:.2f} | "
f"p90: {latencies[int(n * 0.9)]:.2f} | "
f"max: {latencies[-1]:.2f}"
)
print()
return {
"phase": name,
"n_concurrent": n_concurrent,
"succeeded": n,
"failed": failures,
"total_time_s": round(total_time, 3),
"throughput_fps": round(throughput, 1),
"p90_ms": round(latencies[int(n * 0.9)], 2) if n else None,
"max_ms": round(latencies[-1], 2) if n else None,
}
def main():
results = []
n_concurrent = 10
while True:
phase_name = f"{n_concurrent} concurrent flows"
result = run_phase(phase_name, n_concurrent)
results.append(result)
failed = result.get("failed", 0)
# Stop only on an actual failure (a real exception), not on slowness
if failed > 0:
print(f"!!! STOPPING: {failed} flow(s) failed at {n_concurrent} concurrent flows.")
break
if n_concurrent >= MAX_CONCURRENT_FLOWS:
print(f"!!! STOPPING: reached safety cap of {MAX_CONCURRENT_FLOWS} concurrent flows without failure.")
break
n_concurrent *= 2 # double the load each phase (10 -> 20 -> 40 -> ...)
print("=" * 70)
print("SUMMARY TABLE")
print("=" * 70)
header = f"{'Phase':<32}{'Throughput':<14}{'P90(ms)':<12}{'Max(ms)':<12}{'Failed':<8}"
print(header)
print("-" * len(header))
for r in results:
print(
f"{r['phase']:<32}{str(r['throughput_fps']) + ' f/s':<14}"
f"{str(r['p90_ms']):<12}{str(r['max_ms']):<12}"
f"{str(r['failed']):<8}"
)
if __name__ == "__main__":
main()
nids-system-code-final/capture-service/tests/unit_test.py 0 → 100644
View file @ 532029c9
"""
Unit tests for the NIDS live pipeline (capture-service).
Covers flow key derivation, PPI matrix construction/padding, and model
classification/preprocessing logic — all without requiring live network
traffic or the running Docker network.
"""
import numpy as np
from scapy.all import IP, TCP
import flow_aggregator as fa
import inference_engine as ie
# ---------------------------------------------------------------------------
# flow_aggregator.py
# ---------------------------------------------------------------------------
def test_get_flow_key_same_for_both_directions():
"""Same conversation (two directions) must yield the same flow key."""
pkt_forward = IP(src="10.0.0.1", dst="10.0.0.2") / TCP(sport=5000, dport=443)
pkt_backward = IP(src="10.0.0.2", dst="10.0.0.1") / TCP(sport=443, dport=5000)
key_fwd, _ = fa.get_flow_key(pkt_forward)
key_bwd, _ = fa.get_flow_key(pkt_backward)
assert key_fwd == key_bwd
def test_get_flow_key_none_for_non_tcp():
"""A packet without a TCP layer must return None."""
pkt = IP(src="10.0.0.1", dst="10.0.0.2")
assert fa.get_flow_key(pkt) is None
def test_compute_ppi_matrix_padding():
"""A flow shorter than SEQ_LEN must be zero-padded for the remaining rows."""
flow = fa.create_new_flow_state(timestamp=1000.0, forward=True)
flow["ppi_sizes"] = [100, 200, 150]
flow["ppi_directions"] = [1, -1, 1]
flow["ppi_times"] = [1000.0, 1000.05, 1000.12]
matrix = fa.compute_ppi_matrix(flow)
assert matrix.shape == (fa.SEQ_LEN, 3)
# First packet: IAT is always 0 (no preceding packet)
assert matrix[0, 0] == 0.0
assert matrix[0, 1] == 1
assert matrix[0, 2] == 100
# Remaining rows (3 onward, past the actual sequence) must stay zero (padding)
assert np.all(matrix[3:] == 0.0)
def test_compute_ppi_matrix_iat_calculation():
"""Verify IAT is computed in milliseconds between consecutive packets."""
flow = fa.create_new_flow_state(timestamp=0.0, forward=True)
flow["ppi_sizes"] = [100, 200]
flow["ppi_directions"] = [1, 1]
flow["ppi_times"] = [0.0, 0.5] # 0.5 second gap = 500 milliseconds
matrix = fa.compute_ppi_matrix(flow)
assert matrix[1, 0] == 500.0
# ---------------------------------------------------------------------------
# inference_engine.py
# ---------------------------------------------------------------------------
def test_classify_flow_returns_three_values():
"""classify_flow must return (decision, error, latency_ms) with correct types."""
zero_matrix = np.zeros((ie.model.seq_len, 3), dtype=np.float64)
decision, error, latency_ms = ie.classify_flow(zero_matrix)
assert decision in ("BENIGN", "MALICIOUS")
assert isinstance(error, float)
assert isinstance(latency_ms, float)
assert latency_ms >= 0
def test_classify_flow_decision_matches_threshold():
"""The decision must logically match comparing the error against the threshold."""
zero_matrix = np.zeros((ie.model.seq_len, 3), dtype=np.float64)
decision, error, _ = ie.classify_flow(zero_matrix)
expected = "MALICIOUS" if error > ie.THRESHOLD else "BENIGN"
assert decision == expected
def test_preprocess_ppi_matrix_shape_preserved():
"""Preprocessing must preserve the matrix shape."""
raw = np.random.rand(ie.model.seq_len, 3) * 100
processed = ie.preprocess_ppi_matrix(raw)
assert processed.shape == raw.shape
def test_preprocess_ppi_matrix_does_not_mutate_input():
"""preprocess_ppi_matrix must operate on a copy, not mutate the original matrix in place."""
raw = np.ones((ie.model.seq_len, 3), dtype=np.float64) * 5.0
raw_copy = raw.copy()
ie.preprocess_ppi_matrix(raw)
assert np.array_equal(raw, raw_copy)
nids-system-code-final/dashboard/Dockerfile 0 → 100644
View file @ 532029c9
FROM python:3.11-slim
WORKDIR /app
RUN pip install --no-cache-dir streamlit pandas altair requests
COPY dashboard.py .
EXPOSE 8501
CMD ["streamlit", "run", "dashboard.py", "--server.address=0.0.0.0", "--server.port=8501"]
nids-system-code-final/dashboard/dashboard.py 0 → 100644
View file @ 532029c9
import streamlit as st
import requests
import os
import pandas as pd
import sqlite3
import altair as alt
from datetime import datetime, timedelta
API_BASE_URL = os.environ.get("API_BASE_URL", "http://api-service:8000")
MODEL_NAME = "Stacked LSTM Autoencoder (L=8, Fine-tuned)"
MODEL_THRESHOLD_DISPLAY = 0.0155
EVAL_ROC_AUC = 0.994
EVAL_FPR = "12.1%"
EVAL_RECALL = "99.7%"
st.set_page_config(
page_title="NIDS - Real-Time Detection Dashboard",
page_icon="🛡️",
layout="wide"
)
def load_data():
"""Read all detections from the FastAPI service into a DataFrame."""
try:
resp = requests.get(f"{API_BASE_URL}/detections", timeout=5)
resp.raise_for_status()
rows = resp.json()
df = pd.DataFrame(rows)
if not df.empty:
df["timestamp"] = pd.to_datetime(df["timestamp"])
return df
except Exception as e:
st.error(f"Error reading from API: {e}")
return pd.DataFrame()
def load_active_flow_count():
"""Read the current count of in-progress (not-yet-decided) flows from the API."""
try:
resp = requests.get(f"{API_BASE_URL}/active_flows", timeout=5)
resp.raise_for_status()
return resp.json().get("active_flows_count", 0)
except Exception:
return 0
def render_header():
"""Static header: title only."""
st.markdown(
"""
<div style="background-color:#1a1a2e; border:1px solid #2d2d44; border-radius:12px;
padding:20px 24px; margin-bottom:20px;">
<div style="font-size:2em; font-weight:700; color:white;">
🛡️ Real-Time Malicious Activity Detection
</div>
</div>
""",
unsafe_allow_html=True
)
def render_metric_card(label, value, accent):
"""Colored KPI/performance box."""
st.markdown(
f'''
<div style="background-color:{accent}22; border:1px solid {accent}55;
border-radius:10px; padding:14px 16px; text-align:center;">
<div style="color:#ccc; font-size:1.1em; margin-bottom:10px;">{label}</div>
<div style="color:{accent}; font-size:2.4em; font-weight:700;">{value}</div>
</div>
''',
unsafe_allow_html=True
)
def render_flows_table(display_df):
"""Custom HTML table: colored badges, monospace IPs, fixed columns, right-aligned error."""
rows_html = ""
for _, row in display_df.iterrows():
is_malicious = row["decision"] == "MALICIOUS"
if is_malicious:
badge = (
'<span style="background-color:#4a1414; color:#ff6b6b; padding:5px 14px; '
'border-radius:14px; font-size:1em; white-space:nowrap;">⚠ Malicious</span>'
)
error_color = "#ff6b6b"
else:
badge = (
'<span style="background-color:#14301c; color:#4ade80; padding:5px 14px; '
'border-radius:14px; font-size:1em; white-space:nowrap;">✓ Benign</span>'
)
error_color = "#4ade80"
rows_html += f"""<tr>
<td style="padding:10px 8px; border-bottom:1px solid #222; color:#ddd; font-size:1.05em;">{row['timestamp'].strftime('%Y-%m-%d %H:%M:%S')}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222; font-family:monospace; color:#ddd; font-size:1.05em;">{row['src_ip']}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222; font-family:monospace; color:#ddd; font-size:1.05em;">{row['dst_ip']}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222;">{badge}</td>
<td style="padding:10px 8px; border-bottom:1px solid #222; text-align:right; color:{error_color}; font-family:monospace; font-size:1.1em; font-weight:600;">{row['reconstruction_error']:.3f}</td>
</tr>"""
table_html = f"""<style>
.flows-table {{ width:100%; border-collapse:collapse; table-layout:fixed; }}
.flows-table th {{ text-align:left; padding:10px; border-bottom:1px solid #333; color:#aaa; font-weight:600; font-size:1.05em; }}
</style>
<table class="flows-table">
<thead>
<tr>
<th style="width:20%;">Time</th>
<th style="width:24%;">Source</th>
<th style="width:24%;">Destination</th>
<th style="width:16%;">Decision</th>
<th style="width:16%; text-align:right;">Error</th>
</tr>
</thead>
<tbody>{rows_html}</tbody>
</table>"""
st.markdown(
f'''<div style="max-height:480px; overflow-y:auto; border:1px solid #222; border-radius:8px;">{table_html}</div>''',
unsafe_allow_html=True
)
@st.fragment(run_every=2)
def live_section():
df = load_data()
if df.empty:
st.info("Waiting for first data... Run flow_aggregator.py and generate some traffic (curl https://example.com).")
return
recent_cutoff = datetime.now() - timedelta(seconds=10)
recent_malicious = df[(df["decision"] == "MALICIOUS") & (df["timestamp"] >= recent_cutoff)]
if not recent_malicious.empty:
last_alert = recent_malicious.iloc[-1]
st.markdown(
f"""
<div style="background-color:#8B0000; padding:18px; border-radius:10px;
text-align:center; animation: pulse 1.2s infinite; margin-bottom:20px;">
<h2 style="color:white; margin:0;">
🚨 Malicious activity detected! {last_alert['src_ip']} → {last_alert['dst_ip']}
(error: {last_alert['reconstruction_error']:.3f})
</h2>
</div>
<style>
@keyframes pulse {{
0% {{ opacity: 1; }}
50% {{ opacity: 0.6; }}
100% {{ opacity: 1; }}
}}
</style>
""",
unsafe_allow_html=True
)
total = len(df)
benign_count = len(df[df["decision"] == "BENIGN"])
malicious_count = len(df[df["decision"] == "MALICIOUS"])
detection_rate = (malicious_count / total * 100) if total > 0 else 0
active_flows_now = load_active_flow_count()
if "processing_latency_ms" in df.columns and df["processing_latency_ms"].notna().any():
median_latency = df["processing_latency_ms"].dropna().median()
latency_display = f"{median_latency:.2f} ms"
else:
latency_display = "— ms"
kpi_cols = st.columns(6)
with kpi_cols[0]: render_metric_card("Total Decisions", total, "#60a5fa")
with kpi_cols[1]: render_metric_card("Benign", benign_count, "#4ade80")
with kpi_cols[2]: render_metric_card("Malicious", malicious_count, "#ff6b6b")
with kpi_cols[3]: render_metric_card("Detection Rate", f"{detection_rate:.1f}%", "#c084fc")
with kpi_cols[4]: render_metric_card("Active Flows Now", active_flows_now, "#94a3b8")
with kpi_cols[5]: render_metric_card("Avg Processing Latency", latency_display, "#2dd4bf")
st.write("")
chart_left, chart_right = st.columns(2)
with chart_left:
st.markdown("### Model Performance")
perf_df = pd.DataFrame({
"Metric": ["ROC-AUC", "Recall", "FPR"],
"Value": [EVAL_ROC_AUC, 0.997, 0.121]
})
perf_chart = alt.Chart(perf_df).mark_bar(size=45, cornerRadiusEnd=6).encode(
x=alt.X("Metric:N", title=None, sort=None),
y=alt.Y("Value:Q", title=None, scale=alt.Scale(domain=[0, 1])),
color=alt.Color(
"Metric:N",
scale=alt.Scale(
domain=["ROC-AUC", "Recall", "FPR"],
range=["#60a5fa", "#4ade80", "#fbbf24"]
),
legend=None
),
tooltip=["Metric:N", alt.Tooltip("Value:Q", format=".3f")]
).properties(height=280)
text_labels = perf_chart.mark_text(
align="center", baseline="bottom", dy=-5, color="white", fontSize=14
).encode(text=alt.Text("Value:Q", format=".1%"))
st.altair_chart(perf_chart + text_labels, use_container_width=True)
with chart_right:
st.markdown("### Decision Breakdown")
donut_df = pd.DataFrame({
"Decision": ["Benign", "Malicious"],
"Count": [benign_count, malicious_count]
})
donut_chart = alt.Chart(donut_df).mark_arc(innerRadius=65, outerRadius=115).encode(
theta=alt.Theta("Count:Q"),
color=alt.Color(
"Decision:N",
scale=alt.Scale(domain=["Benign", "Malicious"], range=["#4ade80", "#ff6b6b"]),
legend=alt.Legend(title=None, orient="bottom")
),
tooltip=["Decision:N", "Count:Q"]
).properties(height=280)
st.altair_chart(donut_chart, use_container_width=True)
st.divider()
if "flow_filter" not in st.session_state:
st.session_state.flow_filter = "All"
title_col, btn_col1, btn_col2, btn_col3 = st.columns([4, 1, 1, 1])
with title_col:
st.subheader("🗒️ Latest Flows")
with btn_col1:
if st.button("All", use_container_width=True):
st.session_state.flow_filter = "All"
with btn_col2:
if st.button("Malicious", use_container_width=True):
st.session_state.flow_filter = "Malicious"
with btn_col3:
if st.button("Benign", use_container_width=True):
st.session_state.flow_filter = "Benign"
filtered_df = df.copy()
if st.session_state.flow_filter == "Malicious":
filtered_df = filtered_df[filtered_df["decision"] == "MALICIOUS"]
elif st.session_state.flow_filter == "Benign":
filtered_df = filtered_df[filtered_df["decision"] == "BENIGN"]
display_df = filtered_df.tail(20).sort_values("id", ascending=False)[
["timestamp", "src_ip", "dst_ip", "dst_port", "stage", "decision", "reconstruction_error"]
]
if display_df.empty:
st.caption("No flows match this filter yet.")
else:
render_flows_table(display_df)
st.divider()
st.subheader("📈 Reconstruction Error Over Time")
threshold_value = df["threshold"].iloc[-1]
chart_df = df.tail(200).copy()
base = alt.Chart(chart_df).mark_point(
filled=True, size=28, opacity=0.75, stroke=None
).encode(
x=alt.X("timestamp:T", title="Time"),
y=alt.Y("reconstruction_error:Q", title="Reconstruction Error"),
color=alt.Color(
"decision:N",
scale=alt.Scale(domain=["BENIGN", "MALICIOUS"], range=["#2ECC71", "#E74C3C"]),
legend=alt.Legend(title="Decision")
),
tooltip=["timestamp:T", "src_ip", "dst_ip", "stage", "decision", "reconstruction_error"]
).properties(height=350)
threshold_line = alt.Chart(pd.DataFrame({"y": [threshold_value]})).mark_rule(
color="orange", strokeDash=[6, 4], size=1.5
).encode(y="y:Q")
st.altair_chart((base + threshold_line).interactive(), use_container_width=True)
render_header()
live_section()
nids-system-code-final/docker-compose.yml 0 → 100644
View file @ 532029c9
services:
victim:
build: ./victim
container_name: victim
networks:
lab_network:
ipv4_address: 172.28.0.10
ports:
- "443:443"
restart: unless-stopped
capture-service:
build: ./capture-service
container_name: capture-service
network_mode: host
cap_add:
- NET_RAW
- NET_ADMIN
volumes:
- ./shared_data:/shared_data
depends_on:
- victim
restart: unless-stopped
api-service:
build: ./api-service
container_name: api-service
ports:
- "8000:8000"
volumes:
- ./shared_data:/shared_data
depends_on:
- capture-service
restart: unless-stopped
dashboard:
build: ./dashboard
container_name: dashboard
ports:
- "8501:8501"
environment:
- API_BASE_URL=http://api-service:8000
depends_on:
- api-service
restart: unless-stopped
networks:
lab_network:
driver: bridge
ipam:
config:
- subnet: 172.28.0.0/16
gateway: 172.28.0.1
nids-system-code-final/tests/integration_test.py 0 → 100644
View file @ 532029c9
"""
Integration tests for the NIDS system.
Runs directly on the VM host against the live, running Docker Compose stack
(capture-service, api-service, dashboard) to verify the components talk to
each other correctly across process/container boundaries.
Usage:
docker compose up -d
python3 tests/integration_test.py
"""
import json
import os
import sqlite3
import sys
import time
import urllib.request
PROJECT_ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, os.path.join(PROJECT_ROOT, "capture-service"))
import db_writer # noqa: E402
# Host-side path to the same SQLite file the containers see via the shared volume
HOST_DB_PATH = os.path.join(PROJECT_ROOT, "shared_data", "detections.db")
db_writer.DB_PATH = HOST_DB_PATH
API_URL = "http://localhost:8000"
DASHBOARD_URL = "http://localhost:8501"
# Unique, reserved test IP (TEST-NET-3, RFC 5737) so it never collides with
# real captured traffic and is easy to find/clean up afterwards.
TEST_SRC_IP = "203.0.113.99"
passed = 0
failed = 0
def check(condition, description):
global passed, failed
if condition:
print(f" PASSED {description}")
passed += 1
else:
print(f" FAILED {description}")
failed += 1
def http_get_json(url):
with urllib.request.urlopen(url, timeout=5) as resp:
return json.loads(resp.read().decode())
def http_get_status(url):
try:
with urllib.request.urlopen(url, timeout=5) as resp:
return resp.status
except Exception as e:
return str(e)
def cleanup_test_rows():
conn = sqlite3.connect(HOST_DB_PATH)
cursor = conn.cursor()
cursor.execute("DELETE FROM detections WHERE src_ip = ?", (TEST_SRC_IP,))
conn.commit()
conn.close()
# ---------------------------------------------------------------------------
print("=" * 70)
print("Test 1: db_writer -> SQLite (write path)")
print("=" * 70)
db_writer.init_db()
db_writer.log_detection(
src_ip=TEST_SRC_IP,
dst_ip="10.0.0.1",
dst_port=443,
stage="INTEGRATION_TEST",
decision="BENIGN",
reconstruction_error=0.0123,
threshold=0.0155,
processing_latency_ms=12.5,
flow_duration_sec=1.5,
)
conn = sqlite3.connect(HOST_DB_PATH)
cursor = conn.cursor()
cursor.execute(
"SELECT decision, reconstruction_error, threshold FROM detections "
"WHERE src_ip = ? ORDER BY id DESC LIMIT 1",
(TEST_SRC_IP,),
)
row = cursor.fetchone()
conn.close()
check(row is not None, "test row exists in detections.db after log_detection()")
if row is not None:
check(row[0] == "BENIGN", "decision stored correctly")
check(abs(row[1] - 0.0123) < 1e-6, "reconstruction_error stored correctly")
check(abs(row[2] - 0.0155) < 1e-6, "threshold stored correctly")
# ---------------------------------------------------------------------------
print()
print("=" * 70)
print("Test 2: api-service -> database (read path)")
print("=" * 70)
try:
detections = http_get_json(f"{API_URL}/detections")
check(isinstance(detections, list), "GET /detections returns a JSON list")
test_rows = [d for d in detections if d.get("src_ip") == TEST_SRC_IP]
check(len(test_rows) >= 1, "test row written in Test 1 is visible via the API")
if test_rows:
check(test_rows[-1]["decision"] == "BENIGN", "API-returned decision matches what was written")
active = http_get_json(f"{API_URL}/active_flows")
check("active_flows_count" in active, "GET /active_flows returns expected shape")
except Exception as e:
check(False, f"api-service reachable and responding ({e})")
# ---------------------------------------------------------------------------
print()
print("=" * 70)
print("Test 3: dashboard -> api-service connectivity")
print("=" * 70)
status = http_get_status(DASHBOARD_URL)
check(status == 200, f"dashboard is reachable on {DASHBOARD_URL} (status={status})")
# ---------------------------------------------------------------------------
cleanup_test_rows()
print()
print("=" * 70)
print(f"SUMMARY: {passed} passed, {failed} failed")
print("=" * 70)
sys.exit(0 if failed == 0 else 1)
nids-system-code-final/tests/security_test.py 0 → 100644
View file @ 532029c9
"""
Security tests for the NIDS system.
Runs directly on the VM host against the live Docker Compose stack.
Covers: TLS certificate/protocol hygiene, container capability minimization,
API error information leakage, CORS exposure, and dependency vulnerability
scanning across the three Python containers.
Usage:
docker compose up -d
python3 tests/security_test.py
"""
import json
import subprocess
import sys
import urllib.request
VICTIM_HOST = "192.168.52.129"
VICTIM_PORT = 443
API_URL = "http://localhost:8000"
passed = 0
failed = 0
info = 0
def check(condition, description):
global passed, failed
if condition:
print(f" PASSED {description}")
passed += 1
else:
print(f" FAILED {description}")
failed += 1
def note(description):
global info
print(f" INFO {description}")
info += 1
def run(cmd, timeout=10):
"""Run a shell command and return (returncode, stdout+stderr).
Passes empty stdin so interactive tools like `openssl s_client`
close automatically after the handshake instead of hanging."""
result = subprocess.run(
cmd, capture_output=True, text=True, timeout=timeout, input=""
)
return result.returncode, result.stdout + result.stderr
# ---------------------------------------------------------------------------
print("=" * 70)
print("Test 1: TLS certificate and protocol hygiene (victim)")
print("=" * 70)
# 1a. Confirm the server rejects the obsolete SSLv3 protocol
rc, out = run(
["openssl", "s_client", "-connect", f"{VICTIM_HOST}:{VICTIM_PORT}", "-ssl3"],
timeout=5,
)
check(rc != 0 or "handshake failure" in out.lower() or "wrong version" in out.lower(),
"server rejects the obsolete SSLv3 protocol")
# 1b. Confirm the server accepts a modern protocol (TLS 1.2+)
rc, out = run(
["openssl", "s_client", "-connect", f"{VICTIM_HOST}:{VICTIM_PORT}", "-tls1_2"],
timeout=5,
)
check("CONNECTED" in out and "Cipher is" in out, "server accepts modern TLS 1.2")
# 1c. Retrieve and display certificate subject (self-signed, expected for a lab setup)
rc, out = run(
["openssl", "s_client", "-connect", f"{VICTIM_HOST}:{VICTIM_PORT}",
"-servername", VICTIM_HOST],
timeout=5,
)
if "subject=" in out:
subject_line = [l for l in out.splitlines() if "subject=" in l][0]
note(f"certificate {subject_line.strip()} (self-signed, expected in this lab setup)")
else:
note("could not retrieve certificate subject")
# ---------------------------------------------------------------------------
print()
print("=" * 70)
print("Test 2: Container capability minimization (capture-service)")
print("=" * 70)
rc, out = run(["docker", "inspect", "capture-service",
"--format", "{{.HostConfig.Privileged}}"])
is_privileged = out.strip().lower() == "true"
check(not is_privileged, "capture-service does NOT run in privileged mode")
rc, out = run(["docker", "inspect", "capture-service",
"--format", "{{.HostConfig.CapAdd}}"])
cap_add = out.strip()
note(f"capture-service CapAdd: {cap_add}")
check("NET_RAW" in cap_add and "NET_ADMIN" in cap_add,
"capture-service has the expected packet-capture capabilities")
check("SYS_ADMIN" not in cap_add and "ALL" not in cap_add,
"capture-service does not request unnecessary broad capabilities")
# ---------------------------------------------------------------------------
print()
print("=" * 70)
print("Test 3: API error information leakage")
print("=" * 70)
try:
req = urllib.request.Request(f"{API_URL}/this-endpoint-does-not-exist")
try:
urllib.request.urlopen(req, timeout=5)
body = ""
except urllib.error.HTTPError as e:
body = e.read().decode(errors="ignore")
leaks_path = "/app" in body or "Traceback" in body or ".py" in body
check(not leaks_path, "404 response does not leak internal file paths or stack traces")
except Exception as e:
check(False, f"could not test API error handling ({e})")
# ---------------------------------------------------------------------------
print()
print("=" * 70)
print(f"SUMMARY: {passed} passed, {failed} failed, {info} informational notes")
print("=" * 70)
sys.exit(0 if failed == 0 else 1)
nids-system-code-final/victim/Dockerfile 0 → 100644
View file @ 532029c9
FROM nginx:latest
COPY nginx.conf /etc/nginx/nginx.conf
COPY certs/ /etc/nginx/certs/
EXPOSE 443
nids-system-code-final/victim/certs/server.crt 0 → 100644
View file @ 532029c9
-----BEGIN CERTIFICATE-----
MIIDmTCCAoGgAwIBAgIUEWnmultbR5uoOIIGFEuTJzIwvKIwDQYJKoZIhvcNAQEL
BQAwXDELMAkGA1UEBhMCU1kxETAPBgNVBAgMCERhbWFzY3VzMREwDwYDVQQHDAhE
YW1hc2N1czEQMA4GA1UECgwHTklEU0xhYjEVMBMGA1UEAwwMdmljdGltLmxvY2Fs
MB4XDTI2MDcxNjA2NTkyNVoXDTI3MDcxNjA2NTkyNVowXDELMAkGA1UEBhMCU1kx
ETAPBgNVBAgMCERhbWFzY3VzMREwDwYDVQQHDAhEYW1hc2N1czEQMA4GA1UECgwH
TklEU0xhYjEVMBMGA1UEAwwMdmljdGltLmxvY2FsMIIBIjANBgkqhkiG9w0BAQEF
AAOCAQ8AMIIBCgKCAQEAy7GuWkUxqkf9jlOnh5jUNaM9fIYWUfp5pdNrUS8gjxc/
Q/uuCwn+MwRR3/XWMh36xRfSCNtUbE3Xt0lVH2n1yDExGzwAWspuhXHetKm9ME5D
SSkoxYwK0qA7PYzBXeDr5wfrqzMI82bhURDt0h7oskzDEVhmkQbvrM8JJHpAkS6M
SvYqRCDw/swR56ATlLG4FFeUSwLXI72tGbVJOp5AU/7h1ue1q927QjvvCc6ivNom
r2E2N69izCae7urK+Zfrh60pzXvgw2YVlNnF0yq661dCpiGmKEBHjdNRcT98q26S
4ih9dQKa7Mlll1Ujj1mCzpAB6uDW2hTXmg4q0oRfqwIDAQABo1MwUTAdBgNVHQ4E
FgQU5tJVm9235/jX1pikQDrTb1uhA40wHwYDVR0jBBgwFoAU5tJVm9235/jX1pik
QDrTb1uhA40wDwYDVR0TAQH/BAUwAwEB/zANBgkqhkiG9w0BAQsFAAOCAQEAIS3f
kgdhiBDz0Lm9dAm3oT53hp5iJW8SsE3p9oCdsrIXgYFLxJf/qzibd5rlN0+1DcaP
iKywoVOfSnanZgt1BpvTQAHyQzH4sF5nBWdT6tJ+QCYlfBN1PXUjVT5UDSjyY+8d
q6yzPngXoZpSycTfPIud2wIWvBk53h1AryqjDzn8QB9FzEUxSzCmOumZPw4Vcg7x
9iCSoPkOPUdg/syuyzDioSmNL+eATIoewJk1F4J4K1qyrXiWg/MIUXvhcKBYT1zo
EVN73NOGy1sy9VfqFYPbefv0tMs5vtYXTTuLeN959Au94VyZ2J1/UYzQYBtFrzg0
rq4Ys31c6Y9FAPapTQ==
-----END CERTIFICATE-----
nids-system-code-final/victim/certs/server.key 0 → 100644
View file @ 532029c9
-----BEGIN PRIVATE KEY-----
MIIEvQIBADANBgkqhkiG9w0BAQEFAASCBKcwggSjAgEAAoIBAQDLsa5aRTGqR/2O
U6eHmNQ1oz18hhZR+nml02tRLyCPFz9D+64LCf4zBFHf9dYyHfrFF9II21RsTde3
SVUfafXIMTEbPABaym6Fcd60qb0wTkNJKSjFjArSoDs9jMFd4OvnB+urMwjzZuFR
EO3SHuiyTMMRWGaRBu+szwkkekCRLoxK9ipEIPD+zBHnoBOUsbgUV5RLAtcjva0Z
tUk6nkBT/uHW57Wr3btCO+8JzqK82iavYTY3r2LMJp7u6sr5l+uHrSnNe+DDZhWU
2cXTKrrrV0KmIaYoQEeN01FxP3yrbpLiKH11AprsyWWXVSOPWYLOkAHq4NbaFNea
DirShF+rAgMBAAECggEAP+oc305e6yNzuDkflDisPbhfDre0QTzq9tGEFxFGkh+l
RluunOAh1zg/yJqG/2rQyQ9UZFsN/n3edjCYeSbybEKjOfx10Ac87rpwOxgXuRSV
oFs9496xdq24aPPFqQNugivGEw4MfPOHN7ysIwZ31ZqoisTGSjN2+S5tm/iQUf1N
zNiLd0oHirJ5puLLsAPhbbRKEB7VOInS5TybJvHkqlMSeEqzVZuuJMg1CIaGl1sn
zUenCBMduTGOKtgc0nnjCG0wi68jAY/2UQxdU8DXE03fFoevwtC749VD9VIUhlyl
nmvKG1bXPgIUt4m7r0oGmUJdNMAJP1+l2pSsQwJqIQKBgQDoEJU44MURZKHByXoB
C8q69YOhQL0pvGcvUKip1c5Ml0BG9JiT/XQvS4F0EtsBNKcaS0MkXviY/RAC/9m0
sEXrFbsZV80+2f+iBXl2q27Mm/dJvA9sTqLumi+H4p4IDhpmLzd9FnXlJ24d0wke
wapNZfujuHqpqvl/oRkISqP8vQKBgQDgtAAi1l5T90Icaa7tzOD+yzMP6qQGDEX3
WTbYN2K0fnQNBhkCG2CvqV3JrsDZrZ/DVdU28qGAzfO10prgRow+djmokypKqS0w
EwC4uINgo6rV2maZ7LrkeO28k/u2VucsA3GRSSzrsG0GBCeHNLDHUalZ4Cu9v5pb
evFdZK/4hwKBgC3Jz9eoKdoTtALrgQIh1H9tH+q6Pli0zts8G/ter1TXZxxY+J4C
nWETkdhYPBQR0oLSpM3RW1W8YONmngsT2MfYwtDuTRYAp3rNf7jVfIgsoe9Phu+E
2Ishqw0xvqKyPa051cX9n4d2bzi9Qghqh/K7k5orNTCzBoWTQMAtEgw1AoGBAIZo
BODRuzCcaMYYDROfBkJCB3QQ1TXs53jxu1WPo5H1RbN4n0LcX3QTac82oZJTQYoU
8GlirmoeRxXOOOfIKyE9lHWaOaTZFvd3lmnnrGYE4C306oRNzujU0NeI6PqDdtmj
4a2WyLDeIgiEOaStRsSiuxrQdTmVv13S50cG8B9HAoGAQ5/QJe7lZ14SD6BbwW6P
XlQxedQJ5+ajnhH+4nxZxYlSkuClVFvJZU+K8gVOJK13X6y9tuzmoxaEMABQNqzp
9U2+ki26bSPGv1UUwWQ2l+ZVpsqc3aL8GrrfTLvX44BIlkB2ktWm2j+TThzKId/E
fW61SWlLnrL0lgvqamcpS64=
-----END PRIVATE KEY-----
nids-system-code-final/victim/nginx.conf 0 → 100644
View file @ 532029c9
events {}
http {
server {
listen 443 ssl;
server_name victim.local;
ssl_certificate /etc/nginx/certs/server.crt;
ssl_certificate_key /etc/nginx/certs/server.key;
location / {
root /usr/share/nginx/html;
index index.html;
}
}
}
report.docx 0 → 100644
View file @ 532029c9
File added
Markdown is supported
0% or
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment

HIAST - GitLab