""" 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:,}")