src/model.py

"""
model.py
========
LSTM-Autoencoder for REST API anomaly detection.

Supports two input modes:
  - 'embedding' : token sequences (CSIC 2010) — integers → learned vectors
  - 'numeric'   : feature sequences (CIC-IDS2018, UNSW-NB15) — floats directly

The model trains only on normal traffic.
At inference, high reconstruction error = anomaly.

Author : K.A.D.S.D. Kandanaarachchi (2020/ICT/19)
Project: Detecting Anomalous REST API Traffic — IT4216
"""

import torch
import torch.nn as nn


class Encoder(nn.Module):
    """
    Reads a session sequence and compresses it into a fixed-size context vector (the bottleneck).

    Input  : (batch, seq_len, input_size)
    Output : (batch, hidden_size)  — last hidden state only
    """

    def __init__(
        self,
        input_size: int,
        hidden_size: int,
        num_layers: int = 2,
        dropout: float = 0.2,
    ):
        super().__init__()
        self.lstm = nn.LSTM(
            input_size=input_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0.0,
        )

    def forward(self, x):
        _, (hidden, _) = self.lstm(x)
        # hidden shape: (num_layers, batch, hidden_size)
        # take only the top layer's hidden state
        return hidden[-1]


class Decoder(nn.Module):
    """
    Takes the bottleneck vector and reconstructs the original sequence step by step.

    Input  : (batch, hidden_size)
    Output : (batch, seq_len, input_size)
    """

    def __init__(
        self,
        hidden_size: int,
        output_size: int,
        seq_len: int,
        num_layers: int = 2,
        dropout: float = 0.2,
    ):
        super().__init__()
        self.seq_len = seq_len

        self.lstm = nn.LSTM(
            input_size=hidden_size,
            hidden_size=hidden_size,
            num_layers=num_layers,
            batch_first=True,
            dropout=dropout if num_layers > 1 else 0.0,
        )
        # Project LSTM output back to original feature size
        self.output_layer = nn.Linear(hidden_size, output_size)

    def forward(self, context):
        # Repeat the context vector for each time step
        # shape: (batch, seq_len, hidden_size)
        repeated = context.unsqueeze(1).repeat(1, self.seq_len, 1)
        out, _ = self.lstm(repeated)
        # Project to output size at each timestep
        return self.output_layer(out)


class LSTMAutoencoder(nn.Module):
    """
    Full LSTM-Autoencoder for sequence anomaly detection.

    Parameters
    ----------
    input_mode  : 'embedding' for token sequences (CSIC 2010)
                  'numeric'   for float sequences (CIC-IDS2018, UNSW-NB15)
    vocab_size  : required when input_mode='embedding'
    embed_dim   : embedding dimension when input_mode='embedding'
    input_size  : number of features when input_mode='numeric'
    hidden_size : LSTM hidden state size (bottleneck width)
    seq_len     : number of timesteps per session (window size)
    num_layers  : number of stacked LSTM layers
    dropout     : dropout between LSTM layers
    """

    def __init__(
        self,
        input_mode: str = "numeric",
        vocab_size: int = None,
        embed_dim: int = 32,
        input_size: int = 23,
        hidden_size: int = 64,
        seq_len: int = 5,
        num_layers: int = 2,
        dropout: float = 0.2,
    ):
        super().__init__()

        self.input_mode = input_mode
        self.seq_len = seq_len

        # Input layer depends on mode
        if input_mode == "embedding":
            assert vocab_size is not None, "vocab_size required for embedding mode"
            self.embedding = nn.Embedding(vocab_size, embed_dim, padding_idx=0)
            encoder_input_sz = embed_dim
            decoder_output_sz = embed_dim
        else:
            self.embedding = None
            encoder_input_sz = input_size
            decoder_output_sz = input_size

        self.encoder = Encoder(encoder_input_sz, hidden_size, num_layers, dropout)
        self.decoder = Decoder(
            hidden_size, decoder_output_sz, seq_len, num_layers, dropout
        )

    def forward(self, x):
        """
        x shape:
          embedding mode : (batch, seq_len)          int64 token ids
          numeric mode   : (batch, seq_len, features) float32
        """
        if self.input_mode == "embedding":
            x = self.embedding(x)  # → (batch, seq_len, embed_dim)

        context = self.encoder(x)
        reconstruction = self.decoder(context)
        return reconstruction

    def reconstruction_error(self, x):
        """
        Compute per-sample mean squared error between input and reconstruction. Used as anomaly score.

        Returns: (batch,) float tensor of error scores
        """
        with torch.no_grad():
            if self.input_mode == "embedding":
                x_float = self.embedding(x).detach()
            else:
                x_float = x

            recon = self.forward(x)
            # MSE per sample: mean over seq_len and features
            error = ((recon - x_float) ** 2).mean(dim=(1, 2))
        return error


def build_model_csic2010(
    vocab_size: int,
    embed_dim: int = 32,
    hidden_size: int = 64,
    num_layers: int = 2,
    seq_len: int = 5,
) -> LSTMAutoencoder:
    """Instantiate model configured for CSIC 2010 token sequences."""
    return LSTMAutoencoder(
        input_mode="embedding",
        vocab_size=vocab_size,
        embed_dim=embed_dim,
        hidden_size=hidden_size,
        seq_len=seq_len,
        num_layers=num_layers,
    )


def build_model_cicids2018(
    n_features: int = 23,
    hidden_size: int = 64,
    num_layers: int = 2,
    seq_len: int = 5,
) -> LSTMAutoencoder:
    """Instantiate model configured for CIC-IDS2018 flow features."""
    return LSTMAutoencoder(
        input_mode="numeric",
        input_size=n_features,
        hidden_size=hidden_size,
        seq_len=seq_len,
        num_layers=num_layers,
    )


def build_model_unsw(
    n_features: int = 20,
    hidden_size: int = 64,
    num_layers: int = 2,
    seq_len: int = 5,
) -> LSTMAutoencoder:
    """Instantiate model configured for UNSW-NB15 flow features."""
    return LSTMAutoencoder(
        input_mode="numeric",
        input_size=n_features,
        hidden_size=hidden_size,
        seq_len=seq_len,
        num_layers=num_layers,
    )


if __name__ == "__main__":
    print("=== Smoke test — all three model configs ===\n")

    # CSIC 2010
    m1 = build_model_csic2010(vocab_size=51)
    x1 = torch.randint(0, 51, (8, 5))
    r1 = m1(x1)
    e1 = m1.reconstruction_error(x1)
    print(
        f"CSIC 2010   input={tuple(x1.shape)}  "
        f"recon={tuple(r1.shape)}  error={e1.mean():.4f}"
    )

    # CIC-IDS2018
    m2 = build_model_cicids2018(n_features=23)
    x2 = torch.randn(8, 5, 23)
    r2 = m2(x2)
    e2 = m2.reconstruction_error(x2)
    print(
        f"CIC-IDS2018 input={tuple(x2.shape)}  "
        f"recon={tuple(r2.shape)}  error={e2.mean():.4f}"
    )

    # UNSW-NB15
    m3 = build_model_unsw(n_features=20)
    x3 = torch.randn(8, 5, 20)
    r3 = m3(x3)
    e3 = m3.reconstruction_error(x3)
    print(
        f"UNSW-NB15   input={tuple(x3.shape)}  "
        f"recon={tuple(r3.shape)}  error={e3.mean():.4f}"
    )

    total = sum(p.numel() for p in m1.parameters())
    print(f"\nCSIC 2010 model parameters: {total:,}")
    total = sum(p.numel() for p in m2.parameters())
    print(f"CIC-IDS2018 model parameters: {total:,}")