Spam. Anomalies. Malware. Sentiment.
Four AI Models Built for Security. All Scored Against Live APIs.

The Module

This module sits at the intersection of machine learning and cybersecurity. The goal is direct: build real AI models that solve real security problems, then evaluate them against live API endpoints. Four labs, each ending with a working model submitted to a scoring endpoint on a live target. The module is part of the AI Red Teamer path, which covers everything from AI fundamentals through prompt injection, LLM output attacks, and AI evasion techniques.

The four labs covered in this module:

• Lab 1: Spam Detection using Naive Bayes
• Lab 2: Network Anomaly Detection using Random Forest
• Lab 3: Malware Image Classification using ResNet50 and PyTorch
• Lab 4 (Skills Assessment): Sentiment Analysis

Environment Setup

Labs 1, 2, and 4 ran locally on a Mac M2 inside a conda environment. Lab 3 ran entirely inside a JupyterLab instance on the HTB target machine because it required PyTorch and GPU-adjacent tooling pre-installed. VPN connectivity was required for all labs to reach the evaluation API endpoints.

conda activate ai
pip3 install nltk pandas scikit-learn scipy matplotlib requests joblib seaborn

For the malware lab, the target machine already had PyTorch and torchvision installed inside the JupyterLab environment. Lab 3 is the only one where you cannot work locally.

Lab 1: Spam Detection

Dataset

The SMS Spam Collection dataset contains 5,572 SMS messages labeled as either spam or ham (not spam). It is a classic NLP benchmark hosted by UCI. The class distribution is heavily skewed toward ham, which matters when choosing a scoring metric. F1 is the right choice here, not accuracy, because accuracy would reward a model that labels everything as ham.

Approach

The model uses a scikit-learn Pipeline combining CountVectorizer for feature extraction with a Multinomial Naive Bayes classifier. GridSearchCV was used to tune the alpha hyperparameter across eight candidate values with 5-fold cross-validation scored on F1.

Text preprocessing steps applied before vectorization:

1. Convert to lowercase
2. Remove punctuation and non-alphabetic characters (except $ and ! which carry signal in spam)
3. Tokenize using NLTK word tokenizer
4. Remove English stop words
5. Apply Porter Stemmer to reduce words to root form

vectorizer = CountVectorizer(min_df=1, max_df=0.9, ngram_range=(1, 2))
pipeline = Pipeline([
    ("vectorizer", vectorizer),
    ("classifier", MultinomialNB())
])

param_grid = {"classifier__alpha": [0.01, 0.1, 0.15, 0.2, 0.25, 0.5, 0.75, 1.0]}
grid_search = GridSearchCV(pipeline, param_grid, cv=5, scoring="f1")
grid_search.fit(X, y)

Results

Security Relevance

Spam detection is one of the oldest and most widely deployed AI security applications. The same Naive Bayes approach used here underpins email security products from major vendors. Understanding how it works at the code level lets you reason about where it fails. Adversaries craft phishing messages specifically to avoid token patterns that Naive Bayes flags. Knowing the feature space the classifier uses tells you which tokens to avoid in a red team phishing campaign.

Lab 2: Network Anomaly Detection

Dataset

The NSL-KDD dataset is an improved version of the KDD Cup 1999 dataset. It contains network connection records with 41 features per record, labeled as either normal traffic or one of several attack categories.

Approach

The model uses a Random Forest classifier. Categorical features (protocol_type and service) were one-hot encoded using pandas get_dummies. Numeric features were used directly without scaling. Random Forest is scale-invariant, so normalisation adds no value here. The dataset was split 80/20 for train and test, with a further 70/30 split within the training set to create a validation set.

features_to_encode = ['protocol_type', 'service']
encoded = pd.get_dummies(df[features_to_encode])
train_set = encoded.join(df[numeric_features])

rf_model = RandomForestClassifier(random_state=1337)
rf_model.fit(multi_train_X, multi_train_y)

Results

Security Relevance

The NSL-KDD dataset is a staple of network security research. Random Forest performs exceptionally well on structured network flow data because attack categories cluster distinctly in feature space. In production, this type of model feeds into SIEM platforms to triage alerts by attack category rather than simply flagging anomalies. The network probe category (nmap, portsweep) is particularly relevant: the model learns the traffic signature of active reconnaissance, not just the payload.

Lab 3: Malware Image Classification

Dataset

The MalImg dataset contains grayscale visualisations of malware binaries. When a binary is rendered as an image, different malware families produce visually distinct patterns due to their code structure, data sections, and obfuscation characteristics. The dataset covers 25 malware families. This technique is used in production threat intelligence tooling because it is robust to bytecode-level obfuscation that defeats signature-based detection.

Approach

The model uses transfer learning with a pretrained ResNet50 CNN loaded from torchvision. All ResNet50 layers were frozen to preserve the pretrained ImageNet weights. Only the final fully connected layer was replaced and retrained for 25 output classes. Images were preprocessed with resize to 75x75, tensor conversion, and ImageNet normalization. Dataset split 80/20 using splitfolders.

class MalwareClassifier(nn.Module):
    def __init__(self, n_classes):
        super(MalwareClassifier, self).__init__()
        self.resnet = models.resnet50(weights='DEFAULT')

        for param in self.resnet.parameters():
            param.requires_grad = False

        num_features = self.resnet.fc.in_features
        self.resnet.fc = nn.Sequential(
            nn.Linear(num_features, HIDDEN_LAYER_SIZE),
            nn.ReLU(),
            nn.Linear(HIDDEN_LAYER_SIZE, n_classes)
        )

ResNet50 (pretrained, all layers frozen)
      |
Linear(2048, 1000)
      |
ReLU
      |
Linear(1000, 25)
      |
Predicted Malware Family (25 classes)

Challenges

This was the hardest lab in the module. Each epoch took approximately 7.5 minutes on the target machine's CPU. The original script used 10 epochs, which would take over 75 minutes. The HTB target has a time limit and expired multiple times during early attempts, losing all training progress each time.

Two things resolved it:

• Reduced epochs from 10 to 3 to finish within the target window. Three epochs were enough to achieve a passing score.
• Used curl instead of the Python requests library for the model upload. Requests timed out on the large .pth file. curl handled it cleanly.

curl -X POST http://localhost:8002/api/upload -F "model=@malware_classifier.pth"

Results

Security Relevance

Malware visualisation is a real technique used in threat intelligence. Tools like CAPE Sandbox and academic researchers use binary-to-image conversion as a feature for malware family clustering. Transfer learning with ResNet50 is practical here because the visual patterns in malware binaries share enough structure with natural images that ImageNet pretraining provides useful feature extraction without fine-tuning the full network. The tradeoff is inference accuracy, which is lower than Random Forest on tabular data. But for detection tasks where the feature space is image-based, it is the right tool.

Lab 4: Skills Assessment (Sentiment Analysis)

Dataset

25,000 labeled movie reviews from IMDB in JSON format. Labels were binary: 1 for positive and 0 for negative. The raw text contained HTML markup including br tags used for paragraph breaks, which required explicit cleaning before any vectorization step.

Approach

The model reused the same Naive Bayes pipeline architecture from Lab 1. The key additional step was a preprocessing function to strip HTML tags from the raw review text before anything else runs.

def clean_text(text):
    text = re.sub(r"<.*?>", " ", text)      # strip HTML tags
    text = re.sub(r"[^\w\s]", " ", text)    # remove punctuation
    text = re.sub(r"\s+", " ", text).strip() # normalise whitespace
    return text

CountVectorizer was configured with unigrams and bigrams, English stop words, and a word-boundary token pattern.

Results

Security Relevance

Sentiment analysis in security applies directly to threat intelligence processing. Security teams ingest massive volumes of text from threat feeds, dark web forums, and security blogs. Classifying that content by sentiment and relevance allows analysts to prioritise what they read rather than working through unstructured data. A model that labels forum posts as positive or negative can surface emerging threat discussions before they appear in formal advisories.

Key Takeaways

Repository

All scripts and write-ups from this module are available on GitHub. The repository will be updated as additional modules on the AI Red Teamer path are completed.