Cyberattack Detection & Response Agent

This project implements an automated, closed-loop security system designed to detect cyberattacks in network traffic and trigger automated defense responses. It utilizes Machine Learning (XGBoost) for high-accuracy threat classification and Deep Reinforcement Learning (DQN) for dynamic, context-aware incident response.

---

Project Structure

• data/raw/: Directory for storing the CICIDS2017 network traffic CSV datasets.
• src/: Core pipeline source code:
  • preprocessing.py: Data cleaning, scaling fitting/persistence, target label encoding, training balancing, and full-dimensional dummy generator.
  • detection.py: Training, evaluation, and saving of the XGBoost threat detection classifier.
  • response.py: Custom Gymnasium MDP environment (with simulated Host Health metrics) and the stable PyTorch DQN Response Agent.
• models/: Directory where trained models and preprocessor artifacts are saved:
  • xgb_model.pkl: Trained XGBoost classifier.
  • dqn_model.pth: Deep Q-Network state dictionary.
  • scaler.pkl: MinMaxScaler fit parameters.
  • label_encoder_target.pkl: Label encoder classes mapping.
  • feature_cols.pkl: Canonical feature column sequence.
• tests/: Automated unit tests suite:
  • test_pipeline.py: Testing preprocessing correctness, DQN weight shapes, and Gymnasium environment state changes.
• app.py: Streamlit-based graphical dashboard demonstration.

---

Setup & installation

1. Install Dependencies:

   pip install -r requirements.txt

2. Run Tests: Before proceeding, verify that the package logic is functional by running the unit tests suite:

   python -m unittest tests/test_pipeline.py

3. Data Setup:

   • For real-world use, download the CICIDS2017 dataset.
   • Place the CSV files (e.g. Monday-WorkingHours.pcap_ISCX.csv) inside the data/raw/ directory.
   • Note: If no files exist in data/raw/, running the scripts or dashboard will automatically generate a proper 78-column dummy dataset dummy.csv for demonstration purposes.

---

Usage

1. Train the Threat Detection Model

Run the detection pipeline script to preprocess the raw data, isolate splits cleanly, and fit the XGBoost classifier. This will save xgb_model.pkl along with preprocessor artifacts (scaler.pkl, feature_cols.pkl, and label_encoder_target.pkl):

python src/detection.py

2. Train the RL Response Agent

Train the DQN agent in the custom Gymnasium environment where actions impact simulated Host Health metrics:

python src/response.py

This saves the trained DQN model state weights to models/dqn_model.pth.

3. Run the Web Application

Launch the interactive Streamlit dashboard to simulate network stream classification and automated mitigation tracking:

streamlit run app.py

---

Key Refactoring Improvements (Upgrade to 9.5+ Rating)

• Eliminated Data Leakage: Moved splitting logic before data scaling and training downsampling.
• Robust Out-of-Sample Preprocessing: PERSISTED scalers and label encoders as artifacts, ensuring incoming inference streams are processed identically to train splits with zero shape or encoding misalignment.
• True MDP Formulation: Introduced a stateful host_health metric to the response environment. The DQN agent's choices (such as allowing attacks or triggering isolation blocks) directly influence this metric, creating a mathematically sound Markov Decision Process.
• DQN Stability Upgrades: Integrated a standard Experience Replay Buffer and a Target Network structure, solving reinforcement learning divergence issues.