HTB AI Evasion Library

A focused toolkit used across Hack The Box's AI evasion content. It bundles reproducibility helpers, dataset utilities, reference models, training loops, evaluation metrics, and Hack The Box–styled visualizations for building and analyzing adversarial ML experiments.

Table of Contents

• Getting Started
  • Prerequisites
  • Installation
• Quickstart
• Module Overview
  • Core Utilities
  • Data Utilities
  • Reference Models
  • Training Helpers
  • Metrics
  • Visualization
• Documentation
• Testing
  • Unit Tests
• Dataset Management
• Backwards Compatibility
• Contributing
• License

Version: 0.4.0

Getting Started

Prerequisites

• Python 3.10+
• pip (or another PEP 517–compatible installer)
• PyTorch and TorchVision (install via the official instructions for your OS/CUDA)

Installation

A new venv is recommended. First install PyTorch following the official guide for your platform, then install the library.

From PyPI (recommended for users of the library):

# Install PyTorch + TorchVision per https://pytorch.org/get-started/locally/
pip install torch torchvision

pip install htb-ai-library

From source (recommended for contributors and for running tests/docs locally):

git clone https://github.com/PandaSt0rm/htb-ai-library.git
cd htb-ai-library

# Install PyTorch + TorchVision per https://pytorch.org/get-started/locally/
pip install torch torchvision

pip install -e .

Quickstart

import torch

from htb_ai_library.utils import set_reproducibility
from htb_ai_library.data import get_mnist_loaders
from htb_ai_library.models import SimpleLeNet, SimpleCNN
from htb_ai_library.training import train_model
from htb_ai_library.evaluation import evaluate_attack_effectiveness
from htb_ai_library.visualization import use_htb_style, visualize_attack

# Apply HTB theme globally to all plots
use_htb_style()

# Reproducibility and device selection
set_reproducibility(1337)
device = "cuda" if torch.cuda.is_available() else "cpu"

# Data + model
train_loader, test_loader = get_mnist_loaders(batch_size=128, data_dir="./data")
model = SimpleCNN().to(device)  # or SimpleLeNet() for a LeNet-style alternative

# Train for a couple of epochs
train_model(model, train_loader, test_loader, device=device, epochs=2)

# Evaluate an example attack (dummy perturbation shown for brevity)
images, labels = next(iter(test_loader))
images, labels = images.to(device), labels.to(device)
adversarial = images + 0.01 * torch.sign(torch.randn_like(images))

results = evaluate_attack_effectiveness(model, images, adversarial, labels)
print(results)

# Visualize the difference for the first sample
visualize_attack(
    model,
    images[0].cpu(),
    labels[0].item(),
    adversarial[0].cpu(),
    title="Demo Attack",
    num_classes=10,
)

This snippet trains the reference MNIST CNN, computes attack metrics for a mock perturbation, and opens the Hack The Box–styled visualization.

Note: get_mnist_loaders returns MNIST tensors in [0, 1] by default. Pass normalize=True to apply the canonical (mean=0.1307, std=0.3081) normalization, and use mnist_denormalize when you need to revert normalized tensors for visualization.

Module Overview

The library ships as focused subpackages. Import from htb_ai_library.utils, htb_ai_library.data, htb_ai_library.models, htb_ai_library.training, htb_ai_library.evaluation, and htb_ai_library.visualization. Legacy modules (htb_ai_library.core, htb_ai_library.metrics, htb_ai_library.viz) proxy to the new layout for backwards compatibility.

Core Utilities (htb_ai_library.utils)

Function	Description
set_reproducibility(seed)	Align RNGs across random, NumPy, and PyTorch (CPU + CUDA).
save_model(model, filepath)	Persist model with version metadata.
load_model(model_or_filepath, ...)	Load model with version safeguards and friendly error messages.
get_color(name)	Retrieve a single HTB brand color by name.
get_color_palette()	Get the full HTB color palette dictionary.

Color Constants: HTB_GREEN, NODE_BLACK, MALWARE_RED, VIVID_PURPLE, AQUAMARINE, AZURE, NUGGET_YELLOW, HACKER_GREY, WHITE.

Data Utilities (htb_ai_library.data)

Function	Description
get_mnist_loaders(batch_size, data_dir, normalize, seed)	MNIST train/test loaders. Pixels in [0,1] by default; pass normalize=True for canonical normalization.
mnist_denormalize(x)	Invert MNIST normalization back to [0,1].
get_cifar10_loaders(batch_size, download, data_root)	CIFAR-10 datasets and loaders with standard normalization applied.
get_cifar10_transform()	Standard CIFAR-10 normalization transform.
cifar_normalize(x)	Apply CIFAR-10 mean/std normalization to a tensor batch.
load_adult_census(test_size, random_state, data_dir)	Load Adult Census dataset from UCI ML Repository with train/shadow/eval splits for privacy attacks.
download_sms_spam_dataset(data_dir)	Fetch and cache the UCI SMS Spam Collection as sms_spam.csv.

Reference Models (htb_ai_library.models)

Image Classifiers

Model	Description
SimpleCNN	Fast MNIST-ready CNN for demonstrations.
SimpleLeNet	LeNet-style classifier with optional forward_log_probs for log-softmax outputs.
MNISTClassifierWithDropout	Dropout-regularized MNIST classifier.
ResNetCIFAR	Lightweight ResNet-18 style network for CIFAR-10.
CIFAR10CNN	Simple CNN for CIFAR-10, compatible with Opacus DP-SGD training.

Tabular & Attack Models

Model	Description
MLP(input_size, hidden_layers, num_classes, dropout)	Multi-layer perceptron for tabular classification tasks. Includes predict_proba() method.
AttackModel(input_size, hidden_layers, dropout)	Binary classifier for membership inference attacks. Takes prediction vectors concatenated with one-hot labels.

Training Helpers (htb_ai_library.training)

Standard Training

Function	Description
train_one_epoch(model, loader, ...)	Single-epoch loop with averaged loss reporting.
train_model(model, train_loader, test_loader, ...)	Multi-epoch trainer with Adam optimizer and validation accuracy logging.
evaluate_accuracy(model, loader, device)	Compute dataset accuracy with safe handling of empty loaders.

Privacy Training

Function	Description
train_with_early_stopping(model, train_loader, val_loader, device, epochs, patience, ...)	Training loop with early stopping. Returns history dict with losses and accuracies.
get_model_predictions(model, X, device, batch_size)	Get softmax probability predictions for a numpy array.
compute_mia_advantage(model, train_loader, test_loader, device)	Perform confidence-based Membership Inference Attack. Returns (attack_accuracy, advantage).
prepare_attack_data(predictions_in, predictions_out, labels_in, labels_out)	Prepare feature/label arrays for training an attack model (predictions + one-hot labels).
create_dataloader(X, y, batch_size, shuffle)	Create a PyTorch DataLoader from numpy arrays.

Metrics (htb_ai_library.evaluation)

Function	Description
evaluate_attack_effectiveness(model, clean, adv, labels)	Compute clean/adversarial accuracy, attack success rate, confidence deltas, and perturbation norms.
analyze_model_confidence(model, loader, device, n_samples)	Batched confidence statistics for a configurable number of samples.
print_attack_summary(results)	Formatted console report of the metrics dictionary.

Visualization (htb_ai_library.visualization)

Global Styling

Function	Description
use_htb_style()	Apply HTB theme globally to all matplotlib plots. Call once at script start.

Adversarial Attack Plots

Function	Description
visualize_attack(model, clean, label, adv, ...)	Side-by-side clean/adversarial comparison with perturbation heatmap and probability histograms.
plot_attack_effectiveness(results)	Effort vs. evasion-rate plot for tabular attack results.
visualize_perturbation_analysis(perturbations)	Histograms of L2/L∞ norms and sparsity.

Privacy Attack Plots

Function	Description
plot_training_history(history, title, save_path)	Training/validation loss and accuracy curves.
plot_overfitting_gap(train_acc, test_acc, save_path)	Visualize the overfitting gap that enables membership inference.
plot_confidence_distributions(probs_members, probs_non_members, save_path)	Distribution of prediction confidence for members vs non-members.
plot_attack_roc_curve(y_true, y_scores, save_path)	ROC curve for membership inference attack. Returns AUC.
plot_precision_recall_curve(y_true, y_scores, save_path)	Precision-recall curve for membership inference attack.
plot_attack_metrics(results, save_path)	Bar chart comparing attack accuracy, AUC, precision, and recall.
plot_accuracy_comparison(baseline, dp10, dp3, save_path)	Compare model accuracies across privacy levels (baseline, ε=10, ε=3).
plot_privacy_utility_tradeoff(test_accs, mia_advs, epsilons, save_path)	Dual-axis plot showing test accuracy and MIA advantage vs. privacy budget.

Documentation

The repository ships with Sphinx docs that combine this README with an autogenerated API reference.

1. Install the optional docs dependencies:

   pip install -e .[docs]

   Or install them directly from docs/requirements.txt.
2. Build the HTML documentation:

   make -C docs html

   The rendered site is written to docs/_build/html/index.html. Pass additional options through SPHINXOPTS, e.g. make -C docs html SPHINXOPTS="-W --keep-going" to fail on warnings.
3. (Optional) Validate external links:

   make -C docs linkcheck

Testing

Unit Tests

A pytest suite lives in tests/ and covers the public API with fast, CPU‑only tests. Run it after installing dev dependencies:

pytest

For coverage:

coverage run -m pytest
coverage report -m --include='src/htb_ai_library/*'

Dataset Management

• MNIST caches under ./data/MNIST by default.
• CIFAR-10 caches under ./data/cifar-10-batches-py by default.
• Adult Census caches as ./data/adult.data and ./data/adult.test (downloaded from UCI ML Repository).
• The SMS Spam Collection is written to ./data/sms_spam.csv after downloading; intermediate archives are removed.
• To track datasets elsewhere, change data_dir arguments or update .gitignore.

Backwards Compatibility

• htb_ai_library.core, htb_ai_library.metrics, and htb_ai_library.viz remain as compatibility shims that re‑export the new utils, evaluation, and visualization modules. New code should import from the subpackages directly.
• Saved models and generic Python objects use save_model/load_model. When loading generic objects, the loader warns on Python/PyTorch version mismatches and prints actionable messages for corrupted files.

Contributing

Contributions are welcome. Please open an issue or pull request with a clear description of the change. If you add or modify functionality, include tests.

License

This project is licensed under the MIT License.