text_plain_output_5.png application_vnd.jupyter.stderr_output_2.png text_plain_output_9.png application_vnd.jupyter.stderr_output_4.png application_vnd.jupyter.stderr_output_6.png application_vnd.jupyter.stderr_output_12.png application_vnd.jupyter.stderr_output_8.png application_vnd.jupyter.stderr_output_10.png text_plain_output_3.png text_plain_output_7.png text_plain_output_1.png text_plain_output_11.png
from sklearn.metrics import accuracy_score, precision_score, recall_score, confusion_matrix

from sklearn.model_selection import train_test_split

from sklearn.preprocessing import MinMaxScaler

from tensorflow.python.framework.ops import disable_eager_execution

import numpy as np # linear algebra

import pandas as pd # data processing, CSV file I/O (e.g. pd.read_csv)

import tensorflow as tf

df = pd.read_parquet('/kaggle/input/sampled-datasets-v2/NF-UNSW-NB15-V2.parquet')
columns_to_remove = ['L4_SRC_PORT', 'L4_DST_PORT', 'Label', 'Attack']
train, test = train_test_split(df, test_size=0.2, shuffle=True)
y_test = np.array(test['Label'], dtype=np.uint8)
train_benign_idx = train['Label'] == 0
train_attack_idx = train['Label'] == 1
train.drop(columns=columns_to_remove, axis=1, inplace=True)
test.drop(columns=columns_to_remove, axis=1, inplace=True)
train_normal = train[train_benign_idx].values
train_attack = train[train_attack_idx].values
scaler = MinMaxScaler()
train = scaler.fit_transform(train_normal)
train_attack = scaler.transform(train_attack)
train, validation = train_test_split(train, test_size=0.2)

"""
Function to test a model
========================
    - threshold_quantile: IF reconstruction loss > specified quantile of validation losses => attack! ELSE => benign!
    - validation_benign: benign samples used for validation (and threshold determination)
    - validation_attack: attack samples used for validation
    - test: test data (both benign and attack samples)
    - y_test: ground truth of the test data
    - mae: IF true => loss function equals Mean Absolute Error ELSE loss function equals Mean Squared Error
"""
def test_model(model, threshold_quantile, validation_benign, validation_attack, test, y_test, mae=True):
    val_losses = None
    if mae:
        val_losses = np.mean(abs(validation_benign - model.predict(validation_benign)), axis=1)
    else:
        val_losses = np.mean((validation_benign - model.predict(validation_benign)) ** 2, axis=1)
    val_losses = pd.DataFrame({'benign': val_losses})
    attack_losses = None
    if mae:
        attack_losses = np.mean(abs(validation_attack - model.predict(validation_attack)), axis=1)
    else:
        attack_losses = np.mean((validation_attack - model.predict(validation_attack)) ** 2, axis=1)
    attack_losses = pd.DataFrame({'attack': attack_losses})
    threshold = np.quantile(val_losses, 0.99)
    test_losses = None
    recons = model.predict(test)
    if mae:
        test_losses = np.mean(abs(test - recons), axis=1)
    else:
        test_losses = np.mean((test - recons) ** 2, axis=1)
    preds = np.array(test_losses > threshold, dtype=np.uint8)
    tn, fp, fn, tp = confusion_matrix(y_test, preds).ravel()

disable_eager_execution()
original_dim = train.shape[1]
input_shape = (original_dim,)
intermediate_dim = int(original_dim / 2)
latent_dim = int(original_dim / 3)
def sample(args):
    z_mean, z_log_var = args
    batch = K.shape(z_mean)[0]
    dim = K.int_shape(z_mean)[1]
    epsilon = K.random_normal(shape=(batch, dim))
    return z_mean + K.exp(0.5 * z_log_var) * epsilon
inputs = tf.keras.layers.Input(shape=input_shape, name='encoder_input')
x = tf.keras.layers.Dense(int(original_dim / 2), activation='relu')(inputs)
x = tf.keras.layers.Dense(int(original_dim / 3), activation='relu')(x)
z_mean = tf.keras.layers.Dense(int(original_dim / 4), name='z_mean')(x)
z_log_var = tf.keras.layers.Dense(int(original_dim / 4), name='z_log_var')(x)
z = tf.keras.layers.Lambda(sample, output_shape=(input_shape,), name='z')([z_mean, z_log_var])
encoder = tf.keras.Model(inputs, z, name='encoder')
encoder.summary()

latent_inputs = tf.keras.layers.Input(shape=(int(original_dim / 4),), name='z_sampling')
x = tf.keras.layers.Dense(int(original_dim / 3), activation='relu')(latent_inputs)
x = tf.keras.layers.Dense(int(original_dim / 2), activation='relu')(x)
outputs = tf.keras.layers.Dense(original_dim, activation='sigmoid')(x)
decoder = tf.keras.Model(latent_inputs, outputs, name='decoder')
decoder.summary()

outputs = decoder(encoder(inputs))
vae_model = tf.keras.Model(inputs, outputs, name='vae')
vae_model.summary()

def vae_loss(x, x_decoded_mean):
    reconstruction_loss = K.sum(K.square(x - x_decoded_mean))
    kl_loss = -0.5 * K.sum(1 + z_log_var - K.square(z_mean) - K.square(K.exp(z_log_var)), axis=-1)
    total_loss = K.mean(reconstruction_loss + kl_loss)
    return total_loss
vae_model.compile(optimizer='adam', loss=vae_loss)
es = tf.keras.callbacks.EarlyStopping(monitor='val_loss', patience=3)
vae_model.fit(train, train, shuffle=True, epochs=50, batch_size=64, validation_split=0.1, callbacks=[es])
print('\tEVALUATE WITH MAE & QUANTILE 0.95:')
test_model(vae_model, 0.95, validation, train_attack, test, y_test)
print('\tEVALUATE WITH MAE & QUANTILE 0.98:')
test_model(vae_model, 0.98, validation, train_attack, test, y_test)
print('\tEVALUATE WITH MSE & QUANTILE 0.95:')
test_model(vae_model, 0.95, validation, train_attack, test, y_test, mae=False)
print('\tEVALUATE WITH MSE & QUANTILE 0.98:')
test_model(vae_model, 0.98, validation, train_attack, test, y_test, mae=False)