Commit 532029c9 authored by rawan.salameh's avatar rawan.salameh

Add full project: models, system code, and report

parent 25e0a018
This diff is collapsed.
{
"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
}
{
"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
}
{
"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
}
{
"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
}
{
"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
}
This diff is collapsed.
This diff is collapsed.
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"]
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}
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"]
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
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.")
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
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.")
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__), ".."))
"""
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()
"""
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)
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"]
This diff is collapsed.
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
"""
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)
"""
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)
FROM nginx:latest
COPY nginx.conf /etc/nginx/nginx.conf
COPY certs/ /etc/nginx/certs/
EXPOSE 443
-----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-----
-----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-----
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;
}
}
}
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