Skip to content

Dimertuper/NAT_Module_Framework

BranchesTags

Repository files navigation

NAT Module Framework

A modular toolkit for Network Address Translation (NAT) detection in network traffic, built on top of the NEMEA ecosystem. It provides an end-to-end pipeline — from raw flow data parsing, through feature engineering and aggregation, to training ML models and deploying them as standalone NEMEA detection modules.

Highlights

  • Declarative aggregation engine — define complex flow-level feature pipelines with a builder-pattern API; supports both Pandas and Polars backends.
  • One-command module generation — bundle any scikit-learn-compatible model into a deployable NEMEA pytrap module with a single command.
  • ML utilities — wrappers for 10+ classifiers (XGBoost, LightGBM, CatBoost, EBM …), Optuna hyperparameter search, SHAP/LIME explainability, RFECV feature selection, and EDA helpers.
  • Bidirectional flow analysis — direction-aware aggregation splits flows into forward (SRC -> DST) and backward (DST -> SRC) views before computing per-entity features.

Architecture

┌─────────────────────────────────────────────────────────────────────┐
│                        NAT Module Framework                         │
├──────────────────┬──────────────────┬───────────────┬───────────────┤
│  data_aggregator │ nat_module_      │ utils_bundle  │ example_      │
│                  │ framework        │               │ detector      │
│  Flow aggregation│ Model bundling & │ ML, dataset,  │ Sample        │
│  & feature       │ NEMEA module     │ & EDA         │ generated     │
│  extraction      │ code generation  │ utilities     │ module        │
└──────────────────┴──────────────────┴───────────────┴───────────────┘

Typical Workflow

graph LR
    A[Raw IPFIX / CSV Flows] --> B[data_aggregator]
    B -->|Aggregated Features| C[Train Model<br/>utils_bundle]
    C -->|Trained Model + Plan| D[nat_module_framework]
    D -->|Build Module| E[Deployable NEMEA Module]
Loading
  1. Parse & prepare raw flow exports with utils_bundle.dataset_utils.
  2. Define aggregations using the data_aggregator builder API and compile them into a CompiledBackendPlan.
  3. Train a classifier using utils_bundle.ml_utils wrappers (or plain scikit-learn).
  4. Bundle & build the model + plan into a standalone NEMEA module via nat_module_framework.
  5. Deploy the generated module with pytrap on a NEMEA collector.

Project Structure

NAT_Module_Framework/
│
├── data_aggregator/            # Flow aggregation & feature extraction engine
│   ├── src/data_aggregator/
│   │   ├── core.py             # Aggregator builder (add -> run -> export)
│   │   ├── backends/           # Execution backends
│   │   │   ├── _backend.py     # Abstract backend & AggregationPlan
│   │   │   ├── pandas_backend.py
│   │   │   └── polars_backend.py
│   │   ├── runtimes/           # Lightweight runtime for compiled plans
│   │   │   ├── plan.py         # CompiledBackendPlan dataclass
│   │   │   ├── pandas_runtime.py
│   │   │   └── polars_runtime.py
│   │   └── strategies/
│   │       └── strategy.py     # Aggregation & processing strategy definitions
│   ├── tests/
│   ├── pyproject.toml
│   └── requirements.txt
│
├── nat_module_framework/       # NEMEA module generator
│   ├── main.py                 # Entry-point template (pytrap loop)
│   ├── runtime.py              # Bundled runtime (plan + pandas runtime)
│   ├── bundle.py               # create_bundle / load_bundle utilities
│   ├── build_module.py         # CLI: assemble bundle -> deployable module
│   ├── examples/
│   └── tests/
│
├── utils_bundle/               # Shared ML & data-science utilities
│   ├── ml_utils/               # Model wrappers, feature selection, HPO, explainability
│   ├── dataset_utils/          # IPFIX/CSV parsing, train/test splits, resampling
│   ├── exploratory_utils/      # EDA plots, PCA / t-SNE / UMAP projections
│   ├── pyproject.toml
│   └── requirements.txt
│
├── example_detector/           # Sample auto-generated NEMEA module
│   ├── example_detector.py     # Self-contained module script
│   ├── data/                   # model.pkl, plan.pkl, meta.json
│   ├── test.py
│   └── requirements.txt
│
├── Makefile                    # Virtualenv, Jupyter & Podman automation
├── Containerfile.dev           # Dev container (Python 3.10 + all deps)
├── docker-compose.yml          # Podman Compose services
├── requirements.txt            # Pinned full dependency list
└── README.md                   # ← you are here

Getting Started

Prerequisites

Important

The NEMEA framework must be installed on your system before using this project. Follow the NEMEA installation guide to set it up, including the pytrap Python bindings.

Requirement Version
NEMEA framework + pytrap latest
Python ≥ 3.10
Podman + podman-compose latest

Quick Start (Makefile)

The included Makefile automates environment setup. Run make help to see all available targets:

make help
Available targets:
  make venv               Create virtualenv and install dependencies
  make venv-run           Run Jupyter Notebook using the virtualenv
  make venv-clean         Remove virtualenv
  make test               Run all tests
  make build_module       Build a NEMEA module (BUNDLE=path/to/bundle)
  make podman-build       Build Podman container (PROFILE=dev)
  make podman-run         Run Podman Compose (PROFILE=dev)
  make podman-down        Stop Podman Compose
  make podman-logs        Follow Podman logs (PROFILE=dev)

Set up with one command

# Create .venv, install all packages (data_aggregator + utils_bundle), and configure Jupyter
make venv

# Launch Jupyter Notebook (port 8888 by default)
make venv-run

The venv target will:

  1. Create a Python virtual environment in .venv/.
  2. Install all pinned dependencies from the root requirements.txt.
  3. Install data_aggregator and utils_bundle in editable mode.
  4. Set up an IPython startup hook so Jupyter notebooks always use the project root as the working directory.

Build a NEMEA module

make build_module BUNDLE=my_detector

Run tests

make test

Podman workflow

# Build and run the dev container (Jupyter on port 8888)
make podman-build

# Stop containers
make podman-down

# View logs
make podman-logs

Manual Installation

If you prefer not to use the Makefile:

# Clone the repository
git clone https://github.com/<your-username>/NAT_Module_Framework.git
cd NAT_Module_Framework

# Create and activate a virtual environment
python3 -m venv .venv
source .venv/bin/activate

# Install the data aggregator (editable)
pip install -e data_aggregator/

# Install the utils bundle (editable)
pip install -e utils_bundle/

# (Optional) Install additional ML dependencies
pip install -r utils_bundle/requirements.txt

Usage

1. Define an Aggregation Plan

from data_aggregator import Aggregator
from data_aggregator.backends.pandas_backend import PandasBackend

aggregator = (
    Aggregator(backend=PandasBackend())
    .add_preprocessing(my_preprocessing_step)
    .add(my_aggregation_strategy)
    .add_postprocessing(my_postprocessing_step)
)

# Run directly on a DataFrame
result = aggregator.run(df, time_window="30s")

# Or export a compiled plan for deployment
compiled_plan = aggregator.export(time_window="30s", path="plan.pkl")

2. Train a Model

from utils_bundle.ml_utils.classifiers import RandomForestClassifierWrapper

clf = RandomForestClassifierWrapper(n_estimators=200, max_depth=10)
clf.fit(X_train, y_train)
clf.score(X_test, y_test)

3. Bundle Model + Plan

from nat_module_framework.bundle import create_bundle

create_bundle(
    model=clf.estimator,      # sklearn-compatible model
    plan=compiled_plan,       # CompiledBackendPlan
    output_dir="my_detector",
    dependencies=["scikit-learn"],
)

4. Build the NEMEA Module

python nat_module_framework/build_module.py my_detector

This generates a self-contained module directory:

my_detector/
├── my_detector.py      # Standalone pytrap module
├── data/
│   ├── model.pkl
│   ├── plan.pkl
│   └── meta.json
├── requirements.txt
├── README.md
└── tests/

5. Deploy

pip install -r my_detector/requirements.txt
python3 my_detector/my_detector.py -i u:input_socket -o u:output_socket

Add --proba to output probability scores alongside binary predictions (requires a model with predict_proba()).

Components

data_aggregator

The aggregation engine for transforming raw network flows into per-entity feature vectors. It supports:

  • Preprocessing — type conversion, filtering, feature construction.
  • Direction-aware aggregation — splits flows into forward and backward views before grouping by entity (IP address).
  • Postprocessing — derived ratios, normalization, protocol-specific features.
  • Dual backend — choose between Pandas and Polars for execution.
  • Plan compilation — export a CompiledBackendPlan that can be serialized and used at runtime.

nat_module_framework

The module generator that takes a trained model and a compiled aggregation plan and produces a standalone NEMEA detection module:

File Role
bundle.py Create/load bundle directories (model + plan + metadata)
build_module.py CLI script that assembles the final deployable module
runtime.py Lightweight runtime bundled into generated modules
main.py Entry-point template with the pytrap receive loop

utils_bundle

A collection of reusable ML and data-science utilities, organized into three sub-packages:

Sub-package Purpose Key Features
ml_utils Model training & evaluation Classifier wrappers (DT, RF, XGB, LGBM, CatBoost, EBM), RFECV feature selection, Optuna HPO, SHAP/LIME explainability
dataset_utils Data preparation IPFIX/CSV parsing pipeline, train/test splitting, class balancing (undersample, oversample, SMOTE)
exploratory_utils EDA & visualization Descriptive stats, correlation matrices, distribution plots, PCA/t-SNE/UMAP projections

example_detector

A pre-built example of an auto-generated NEMEA module using a Decision Tree classifier. Serves as a reference for the output of the build pipeline.

Testing

# Run data_aggregator tests
cd data_aggregator && pytest tests/

# Run utils_bundle tests
cd utils_bundle && pytest tests/

A Podman test environment is also available for data_aggregator:

podman build -f data_aggregator/Dockerfile.test -t data_aggregator_test .
podman run data_aggregator_test

Author

Filip Chodura — choduraf@gmail.com

About

NAT module framework for NEMEA Framework

Resources

Readme
Activity

Stars

0 stars

Watchers

0 watching

Forks

0 forks
Report repository

Releases

Packages

Contributors

Languages