03_classifier.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Classification Prefilter Training & Usage\n",
    "Trained on subset from GEE train split is binary or multiclass manner, using various combinations of attacks"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Download Extracted Features"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install gdown\n",
    "import gdown"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "file_id=\"1pmOhZ3S2v7Fn-uEgI9BuSk9Zj-uxUaFg\"\n",
    "url = f'https://drive.google.com/uc?id={file_id}'\n",
    "output = 'feature.tar.gz'\n",
    "gdown.download(url, output, quiet=False)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "!tar -xf feature.tar.gz"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "snUK8M8iHxa-"
   },
   "source": [
    "## Utils"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E4hifMQyHzlS"
   },
   "outputs": [],
   "source": [
    "import os\n",
    "import sys\n",
    "from pathlib import Path\n",
    "\n",
    "import numpy as np\n",
    "import psutil\n",
    "import pyspark.sql.dataframe\n",
    "from petastorm.etl.dataset_metadata import materialize_dataset\n",
    "from petastorm.unischema import Unischema, dict_to_spark_row\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import col, from_unixtime, unix_timestamp\n",
    "from pyspark.sql.types import StructType, StructField, StringType, LongType, DoubleType\n",
    "\n",
    "feature_min_max = {\n",
    "    'mean_duration': (0.0, 2042.86),\n",
    "    'mean_packet': (1.0, 109214.27272727272),\n",
    "    'mean_num_of_bytes': (28.0, 163795638.0909091),\n",
    "    'mean_packet_rate': (0.0, 17224.14377310265),\n",
    "    'mean_byte_rate': (0.0, 13902452.340182647),\n",
    "    'std_duration': (0.0, 562.7625560888366),\n",
    "    'std_packet': (0.0, 370614.95468242496),\n",
    "    'std_num_of_bytes': (0.0, 543247494.7844237),\n",
    "    'std_packet_rate': (0.0, 15783.66319664221),\n",
    "    'std_byte_rate': (0.0, 16441139.793386225),\n",
    "    'entropy_protocol': (0.0, 2.260220915066596),\n",
    "    'entropy_dst_ip': (0.0, 13.787687869067254),\n",
    "    'entropy_src_port': (0.0, 14.206227931544092),\n",
    "    'entropy_dst_port': (0.0, 14.027301292191831),\n",
    "    'entropy_flags': (0.0, 4.631615665225586)\n",
    "}\n",
    "\n",
    "\n",
    "def read_csv(spark: SparkSession, path: str) -> pyspark.sql.dataframe:\n",
    "    \"\"\"\n",
    "    Read csv files as spark dataframe\n",
    "\n",
    "    :param spark: spark session object\n",
    "    :param path: path of dir containing csv files\n",
    "    :type spark: SparkSession\n",
    "    :type path: str\n",
    "    :return: df\n",
    "    :rtype: pyspark.sql.dataframe\n",
    "    \"\"\"\n",
    "\n",
    "    # define csv schema\n",
    "    schema = StructType([\n",
    "        StructField('timestamp', StringType(), True),\n",
    "        StructField('duration', DoubleType(), True),\n",
    "        StructField('src_ip', StringType(), True),\n",
    "        StructField('dst_ip', StringType(), True),\n",
    "        StructField('src_port', LongType(), True),\n",
    "        StructField('dst_port', LongType(), True),\n",
    "        StructField('protocol', StringType(), True),\n",
    "        StructField('flags', StringType(), True),\n",
    "        StructField('forwarding_status', LongType(), True),\n",
    "        StructField('type_of_service', LongType(), True),\n",
    "        StructField('packet', LongType(), True),\n",
    "        StructField('num_of_bytes', LongType(), True),\n",
    "        StructField('label', StringType(), True),\n",
    "    ])\n",
    "\n",
    "    df = (\n",
    "        spark\n",
    "            .read\n",
    "            .schema(schema)\n",
    "            .csv(path)\n",
    "    )\n",
    "\n",
    "    # convert datetime column from string to unix_timestamp\n",
    "    df = (\n",
    "        df\n",
    "            .withColumn('timestamp', unix_timestamp(col('timestamp'), 'yyyy-MM-dd HH:mm:ss'))\n",
    "    )\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "def patch_time_windows(df: pyspark.sql.dataframe, window_seconds: int):\n",
    "    \"\"\"\n",
    "    Generate time window by\n",
    "    :param df: pyspark dataframe\n",
    "    :param window_seconds: window size in second\n",
    "    :type df: pyspark.sql.dataframe\n",
    "    :type window_seconds: int\n",
    "    :return: df\n",
    "    :rtype: pyspark.sql.dataframe\n",
    "    \"\"\"\n",
    "    time_window = from_unixtime(col('timestamp') - col('timestamp') % window_seconds)\n",
    "\n",
    "    df = (\n",
    "        df\n",
    "            .withColumn('time_window', time_window)\n",
    "    )\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "def init_local_spark():\n",
    "    # initialise local spark\n",
    "    os.environ['PYSPARK_PYTHON'] = sys.executable\n",
    "    os.environ['PYSPARK_DRIVER_PYTHON'] = sys.executable\n",
    "    memory_gb = psutil.virtual_memory().available // 1024 // 1024 // 1024\n",
    "    spark = (\n",
    "        SparkSession\n",
    "            .builder\n",
    "            .master('local[*]')\n",
    "            #.config('spark.driver.memory', f'{memory_gb}g')\n",
    "            .config('spark.driver.memory', f'2g')\n",
    "            .config('spark.driver.host', '127.0.0.1')\n",
    "            .getOrCreate()\n",
    "    )\n",
    "    return spark\n",
    "\n",
    "\n",
    "def normalise(x: float, min_val: float, max_val: float) -> float:\n",
    "    norm_x = (x - min_val) / (max_val - min_val)\n",
    "    if norm_x < 0:\n",
    "        norm_x = 0.0\n",
    "    elif norm_x > 1.0:\n",
    "        norm_x = 1.0\n",
    "\n",
    "    return norm_x\n",
    "\n",
    "\n",
    "def row_generator(x):\n",
    "    time_window, src_ip, feature, label = x\n",
    "    return {\n",
    "        'time_window': time_window,\n",
    "        'src_ip': src_ip,\n",
    "        'feature': np.expand_dims(np.array(feature, dtype=np.float32), axis=0),\n",
    "        'label': label,\n",
    "    }\n",
    "\n",
    "\n",
    "def change_df_schema(spark: SparkSession, schema: Unischema, df: pyspark.sql.DataFrame) -> pyspark.sql.DataFrame:\n",
    "    rows_rdd = (\n",
    "        df\n",
    "            .rdd\n",
    "            .map(row_generator)\n",
    "            .map(lambda x: dict_to_spark_row(schema, x))\n",
    "    )\n",
    "\n",
    "    df = spark.createDataFrame(\n",
    "        rows_rdd,\n",
    "        schema.as_spark_schema()\n",
    "    )\n",
    "\n",
    "    return df\n",
    "\n",
    "\n",
    "def save_parquet_for_petastorm_parquet(spark: SparkSession, df: pyspark.sql.DataFrame, output_path: str,\n",
    "                                       schema: Unischema):\n",
    "    #output_path = Path(output_path).absolute().as_uri()\n",
    "    output_path = 'file://' + str(Path(output_path).absolute())\n",
    "    with materialize_dataset(spark, output_path, schema, row_group_size_mb=256):\n",
    "        (\n",
    "            df\n",
    "                .write\n",
    "                .mode('overwrite')\n",
    "                .parquet(output_path)\n",
    "        )"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "QPwHz0YkI5iI"
   },
   "source": [
    "## Build Model Input"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "1TFB3KlnI-DS"
   },
   "outputs": [],
   "source": [
    "import logging\n",
    "\n",
    "import click\n",
    "import pyspark\n",
    "import numpy as np\n",
    "from petastorm.codecs import CompressedNdarrayCodec, ScalarCodec\n",
    "from petastorm.unischema import Unischema, UnischemaField\n",
    "from pyspark.sql import SparkSession\n",
    "from pyspark.sql.functions import udf, col\n",
    "from pyspark.sql.types import StringType\n",
    "\n",
    "#from utils import init_local_spark, feature_min_max, normalise, change_df_schema, save_parquet_for_petastorm_parquet\n",
    "\n",
    "\n",
    "class FeatureComposer:\n",
    "    def __init__(self, spark: SparkSession, df: pyspark.sql.DataFrame):\n",
    "        self.spark = spark\n",
    "        self.df = df\n",
    "        self.feature_column = (\n",
    "            'mean_duration', 'mean_packet', 'mean_num_of_bytes', 'mean_packet_rate', 'mean_byte_rate', 'std_duration',\n",
    "            'std_packet', 'std_num_of_bytes', 'std_packet_rate', 'std_byte_rate', 'entropy_protocol', 'entropy_dst_ip',\n",
    "            'entropy_src_port', 'entropy_dst_port', 'entropy_flags', 'proportion_src_port', 'proportion_dst_port',\n",
    "        )\n",
    "\n",
    "        self.feature_compose_udf = udf(self.feature_compose, 'array<double>')\n",
    "\n",
    "    @staticmethod\n",
    "    def feature_compose(\n",
    "            mean_duration: float, mean_packet: float, mean_num_of_bytes: float, mean_packet_rate: float,\n",
    "            mean_byte_rate: float, std_duration: float, std_packet: float, std_num_of_bytes: float,\n",
    "            std_packet_rate: float, std_byte_rate: float, entropy_protocol: float, entropy_dst_ip: float,\n",
    "            entropy_src_port: float, entropy_dst_port: float, entropy_flags: float, proportion_src_port: list,\n",
    "            proportion_dst_port: list\n",
    "    ) -> list:\n",
    "        \"\"\"\n",
    "        Compose the feature array\n",
    "        :param mean_duration: mean duration\n",
    "        :param mean_packet: mean packet\n",
    "        :param mean_num_of_bytes: mean number of bytes\n",
    "        :param mean_packet_rate: mean packet rate\n",
    "        :param mean_byte_rate: mean byte rate\n",
    "        :param std_duration: std duration\n",
    "        :param std_packet: std packet\n",
    "        :param std_num_of_bytes: std number of bytes\n",
    "        :param std_packet_rate: std packet rate\n",
    "        :param std_byte_rate: std byte rate\n",
    "        :param entropy_protocol: entropy of protocol\n",
    "        :param entropy_dst_ip: entropy of dest ip\n",
    "        :param entropy_src_port: entropy of src ip\n",
    "        :param entropy_dst_port: entropy of dest port\n",
    "        :param entropy_flags: entropy of flags\n",
    "        :param proportion_src_port: proportion of src common ports\n",
    "        :param proportion_dst_port: proportion of dest common port\n",
    "        :type mean_duration: float\n",
    "        :type mean_packet: float\n",
    "        :type mean_num_of_bytes: float\n",
    "        :type mean_packet_rate: float\n",
    "        :type mean_byte_rate: float\n",
    "        :type std_duration: float\n",
    "        :type std_packet: float\n",
    "        :type std_num_of_bytes: float\n",
    "        :type std_packet_rate: float\n",
    "        :type std_byte_rate: float\n",
    "        :type entropy_protocol: float\n",
    "        :type entropy_dst_ip: float\n",
    "        :type entropy_src_port: float\n",
    "        :type entropy_dst_port: float\n",
    "        :type entropy_flags: float\n",
    "        :type proportion_src_port: list\n",
    "        :type proportion_dst_port: list\n",
    "        :return: feature array\n",
    "        :rtype list\n",
    "        \"\"\"\n",
    "        # normalise\n",
    "        mean_duration = normalise(mean_duration, *feature_min_max.get('mean_duration'))\n",
    "        mean_packet = normalise(mean_packet, *feature_min_max.get('mean_packet'))\n",
    "        mean_num_of_bytes = normalise(mean_num_of_bytes, *feature_min_max.get('mean_num_of_bytes'))\n",
    "        mean_packet_rate = normalise(mean_packet_rate, *feature_min_max.get('mean_packet_rate'))\n",
    "        mean_byte_rate = normalise(mean_byte_rate, *feature_min_max.get('mean_byte_rate'))\n",
    "        std_duration = normalise(std_duration, *feature_min_max.get('std_duration'))\n",
    "        std_packet = normalise(std_packet, *feature_min_max.get('std_packet'))\n",
    "        std_num_of_bytes = normalise(std_num_of_bytes, *feature_min_max.get('std_num_of_bytes'))\n",
    "        std_packet_rate = normalise(std_packet_rate, *feature_min_max.get('std_packet_rate'))\n",
    "        std_byte_rate = normalise(std_byte_rate, *feature_min_max.get('std_byte_rate'))\n",
    "        entropy_protocol = normalise(entropy_protocol, *feature_min_max.get('entropy_protocol'))\n",
    "        entropy_dst_ip = normalise(entropy_dst_ip, *feature_min_max.get('entropy_dst_ip'))\n",
    "        entropy_src_port = normalise(entropy_src_port, *feature_min_max.get('entropy_src_port'))\n",
    "        entropy_dst_port = normalise(entropy_dst_port, *feature_min_max.get('entropy_dst_port'))\n",
    "        entropy_flags = normalise(entropy_flags, *feature_min_max.get('entropy_flags'))\n",
    "\n",
    "        feature_arr = [\n",
    "            mean_duration, mean_packet, mean_num_of_bytes, mean_packet_rate, mean_byte_rate, std_duration, std_packet,\n",
    "            std_num_of_bytes, std_packet_rate, std_byte_rate, entropy_protocol, entropy_dst_ip, entropy_src_port,\n",
    "            entropy_dst_port, entropy_flags,\n",
    "        ]\n",
    "\n",
    "        feature_arr.extend(proportion_src_port)\n",
    "        feature_arr.extend(proportion_dst_port)\n",
    "\n",
    "        return feature_arr\n",
    "\n",
    "    def transform(self, remove_malicious=True, remove_null_label=True) -> pyspark.sql.DataFrame:\n",
    "        df = (\n",
    "            self.df\n",
    "                # compose feature\n",
    "                .withColumn('features', self.feature_compose_udf(*self.feature_column))\n",
    "        )\n",
    "\n",
    "        if remove_null_label:\n",
    "            df = df.filter(col('label').isNotNull())\n",
    "\n",
    "        if remove_malicious:\n",
    "            df = df.filter(col('label') == 'background')\n",
    "\n",
    "        # select only time_window, src_ip, feature and label columns\n",
    "        df = df.select(\n",
    "            'time_window', 'src_ip', 'features', 'label',\n",
    "        )\n",
    "\n",
    "        return df"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2dCA2AdFq6AQ"
   },
   "outputs": [],
   "source": [
    "#temporary to generate complete train dataset (including attacks)\n",
    "test = 'feature/train.feature.parquet'\n",
    "target_test = 'model_input/train.model_input.parquet.complete'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "FSUc3dERq6OX"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "\n",
    "spark = init_local_spark()\n",
    "\n",
    "# petastorm schema\n",
    "schema = Unischema(\n",
    "    'data_schema', [\n",
    "        UnischemaField('time_window', np.str, (), ScalarCodec(StringType()), False),\n",
    "        UnischemaField('src_ip', np.str, (), ScalarCodec(StringType()), False),\n",
    "        UnischemaField('feature', np.float32, (1, 69), CompressedNdarrayCodec(), False),\n",
    "        UnischemaField('label', np.str, (), ScalarCodec(StringType()), True),\n",
    "    ]\n",
    ")\n",
    "\n",
    "# processing test\n",
    "test_feature_df = spark.read.parquet(test)\n",
    "test_input = FeatureComposer(spark, test_feature_df).transform(remove_malicious=False, remove_null_label=True)\n",
    "test_input = change_df_schema(spark, schema, test_input)\n",
    "save_parquet_for_petastorm_parquet(spark, test_input, target_test, schema)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {
    "id": "bi6y96_yU_B2"
   },
   "source": [
    "## Generate dataset"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "2bwDLj84kpKE"
   },
   "outputs": [],
   "source": [
    "#temporary change to load local data\n",
    "data_path = 'model_input/train.model_input.parquet.complete'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "tNrNS1UYkj1x"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "from petastorm import make_reader\n",
    "from petastorm.pytorch import DataLoader\n",
    "\n",
    "num_cores = 2#psutil.cpu_count(logical=True)\n",
    "\n",
    "reader = make_reader(\n",
    "    'file://' + str(Path(data_path).absolute()), reader_pool_type='process', workers_count=num_cores,\n",
    "    pyarrow_serialize=True, num_epochs=1\n",
    ")\n",
    "dataloader = DataLoader(reader, batch_size=300, shuffling_queue_capacity=4096)\n",
    "\n",
    "x_list = []\n",
    "label_list = []\n",
    "\n",
    "for data in dataloader:\n",
    "    x = np.squeeze(data['feature'].numpy())\n",
    "    label = data['label']\n",
    " \n",
    "    label_list.extend(label)\n",
    "    x_list.extend(x.tolist())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "Qg_EAdEZnfHj"
   },
   "outputs": [],
   "source": [
    "#dataset creation for training classifier (balancing)\n",
    "import pandas as pd\n",
    "\n",
    "train_df = pd.DataFrame(\n",
    "    {\n",
    "        'x': x_list,\n",
    "        'label': label_list\n",
    "    }\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "pOHOZTYZuyGo"
   },
   "outputs": [],
   "source": [
    "train_df['label'].value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "zWb9b3Dcoc8J"
   },
   "outputs": [],
   "source": [
    "%%time\n",
    "train_df_background = train_df[train_df['label'] == \"background\"].sample(n=10000, random_state=42, ignore_index=True)\n",
    "train_df_background1 = train_df_background.sample(n=1000, random_state=42, ignore_index=True)\n",
    "train_df_background2 = train_df_background.sample(n=2000, random_state=42, ignore_index=True)\n",
    "train_df_background3 = train_df_background.sample(n=3000, random_state=42, ignore_index=True)\n",
    "train_df_background4 = train_df_background.sample(n=4000, random_state=42, ignore_index=True)\n",
    "train_df_background5 = train_df_background.sample(n=5000, random_state=42, ignore_index=True)\n",
    "train_df_background6 = train_df_background.sample(n=6000, random_state=42, ignore_index=True)\n",
    "train_df_dos = train_df[train_df['label'] == \"dos\"].sample(n=1000, random_state=42, ignore_index=True, replace=True)\n",
    "train_df_scan11 = train_df[train_df['label'] == \"scan11\"].sample(n=1000, random_state=42, ignore_index=True, replace=True)\n",
    "train_df_scan44 = train_df[train_df['label'] == \"scan44\"].sample(n=1000, random_state=42, ignore_index=True, replace=True)\n",
    "train_df_botnet = train_df[train_df['label'] == \"nerisbotnet\"].sample(n=1000, random_state=42, ignore_index=True, replace=True)\n",
    "train_df_spam = train_df[train_df['label'] == \"anomaly-spam\"].sample(n=1000, random_state=42, ignore_index=True, replace=True)\n",
    "train_df_blacklist = train_df[train_df['label'] == \"blacklist\"].sample(n=1000, random_state=42, ignore_index=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "NQTxWWR4yWnj"
   },
   "outputs": [],
   "source": [
    "from sklearn.utils import shuffle\n",
    "\n",
    "train_df_all = pd.concat([train_df_background1, train_df_dos, train_df_scan11, train_df_scan44, train_df_botnet, train_df_spam, train_df_blacklist])\n",
    "train_df_all = shuffle(train_df_all).reset_index().drop(['index'], axis=1)\n",
    "\n",
    "train_df_all_binary_balanced = pd.concat([train_df_background6, train_df_dos, train_df_scan11, train_df_scan44, train_df_botnet, train_df_spam, train_df_blacklist])\n",
    "train_df_all_binary_balanced = shuffle(train_df_all_binary_balanced).reset_index().drop(['index'], axis=1)\n",
    "\n",
    "train_df_without_blacklist = pd.concat([train_df_background1, train_df_dos, train_df_scan11, train_df_scan44, train_df_botnet, train_df_spam])\n",
    "train_df_without_blacklist = shuffle(train_df_without_blacklist).reset_index().drop(['index'], axis=1)\n",
    "\n",
    "train_df_without_spam = pd.concat([train_df_background1, train_df_dos, train_df_scan11, train_df_scan44, train_df_botnet, train_df_blacklist])\n",
    "train_df_without_spam = shuffle(train_df_without_spam).reset_index().drop(['index'], axis=1)\n",
    "\n",
    "train_df_without_botnet = pd.concat([train_df_background1, train_df_dos, train_df_scan11, train_df_scan44, train_df_spam, train_df_blacklist])\n",
    "train_df_without_botnet = shuffle(train_df_without_botnet).reset_index().drop(['index'], axis=1)\n",
    "\n",
    "train_df_without_scan44 = pd.concat([train_df_background1, train_df_dos, train_df_scan11, train_df_botnet, train_df_spam, train_df_blacklist])\n",
    "train_df_without_scan44 = shuffle(train_df_without_scan44).reset_index().drop(['index'], axis=1)\n",
    "\n",
    "train_df_without_scan11 = pd.concat([train_df_background1, train_df_dos, train_df_scan44, train_df_botnet, train_df_spam, train_df_blacklist])\n",
    "train_df_without_scan11 = shuffle(train_df_without_scan11).reset_index().drop(['index'], axis=1)\n",
    "\n",
    "train_df_without_dos = pd.concat([train_df_background1, train_df_scan11, train_df_scan44, train_df_botnet, train_df_spam, train_df_blacklist])\n",
    "train_df_without_dos = shuffle(train_df_without_dos).reset_index().drop(['index'], axis=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "id": "E_kFJi-rZkMY"
   },
   "outputs": [],
   "source": [
    "#save generated balanced datasets\n",
    "train_df_all.to_csv('model_input/train_df_all.csv')\n",
    "train_df_all_binary_balanced.to_csv('model_input/train_df_all_binary_balanced.csv')\n",
    "train_df_without_blacklist.to_csv('model_input/train_df_without_blacklist.csv')\n",
    "train_df_without_spam.to_csv('model_input/train_df_without_spam.csv')\n",
    "train_df_without_botnet.to_csv('model_input/train_df_without_botnet.csv')\n",
    "train_df_without_scan44.to_csv('model_input/train_df_without_scan44.csv')\n",
    "train_df_without_scan11.to_csv('model_input/train_df_without_scan11.csv')\n",
    "train_df_without_dos.to_csv('model_input/train_df_without_dos.csv')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Train Classifier"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import pandas as pd\n",
    "from ast import literal_eval\n",
    "from sklearn.utils import shuffle\n",
    "from sklearn.ensemble import RandomForestClassifier\n",
    "from sklearn import metrics"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "binary = True\n",
    "#use binary_balanced for binary classifier\n",
    "if not binary:\n",
    "  train_df_all = pd.read_csv('model_input/train_df_all.csv')\n",
    "else:\n",
    "  train_df_all = pd.read_csv('model_input/train_df_all_binary_balanced.csv')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#multiclass (each class 1000 samples), we can just not use some class (no need to rebalance)\n",
    "if not binary:\n",
    "  #selected_df = train_df_all\n",
    "  selected_df = train_df_all[train_df_all['label'] != \"blacklist\"]\n",
    "  #selected_df = train_df_all[(train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"dos\")]\n",
    "  #selected_df = train_df_all[(train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"anomaly-spam\")]\n",
    "  #selected_df = train_df_all[(train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"nerisbotnet\")]\n",
    "  #selected_df = train_df_all[(train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"scan11\")]\n",
    "  #selected_df = train_df_all[(train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"scan44\")]\n",
    "  #selected_df = train_df_all[(train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"scan11\") & (train_df_all['label'] != \"scan44\")]\n",
    "  selected_df.x = selected_df.x.apply(literal_eval).tolist()\n",
    "  X = pd.DataFrame(selected_df['x'].to_list())\n",
    "  y = selected_df.label\n",
    "  y.value_counts()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#binary, balanced background vs attacks in ratio of 1:1, adjust number of samples for background according to how many attack classes are omitted (1000 for each)\n",
    "if binary:\n",
    "  selected_df_background = train_df_all[train_df_all['label'] == \"background\"].sample(n=5000, random_state=42, ignore_index=True)\n",
    "  #selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\")]\n",
    "  selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\") & (train_df_all['label'] != \"blacklist\")]\n",
    "  #selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\") & (train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"dos\")]\n",
    "  #selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\") & (train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"anomaly-spam\")]\n",
    "  #selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\") & (train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"nerisbotnet\")]\n",
    "  #selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\") & (train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"scan11\")]\n",
    "  #selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\") & (train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"scan44\")]\n",
    "  #selected_df_anomalies = train_df_all[(train_df_all['label'] != \"background\") & (train_df_all['label'] != \"blacklist\") & (train_df_all['label'] != \"scan11\") & (train_df_all['label'] != \"scan44\")]\n",
    "  selected_df = pd.concat([selected_df_background, selected_df_anomalies])\n",
    "  selected_df = shuffle(selected_df).reset_index().drop(['index'], axis=1)\n",
    "  selected_df.x = selected_df.x.apply(literal_eval).tolist()\n",
    "  X = pd.DataFrame(selected_df['x'].to_list())\n",
    "  y = selected_df.label\n",
    "  print(y.value_counts())\n",
    "\n",
    "  #change to train binary classifier\n",
    "  y[y != 'background'] = '1'\n",
    "  y[y == 'background'] = '0'\n",
    "  print(y.value_counts())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "#train on first 5000 samples and test on the rest, then retrain on the whole train dataset\n",
    "\n",
    "RFC = RandomForestClassifier(random_state=42, n_jobs=-1)\n",
    "RFC.fit(X[:5000], y[:5000])\n",
    "\n",
    "preds = RFC.predict(X[5000:])\n",
    "print(metrics.classification_report(y[5000:], preds, digits=4))\n",
    "\n",
    "RFC.fit(X, y)\n",
    "preds = RFC.predict(X)\n",
    "print(metrics.classification_report(y, preds, digits=4))\n",
    "\n",
    "classifier_model = RFC"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "#temporarily try another classifier\n",
    "from sklearn.tree import DecisionTreeClassifier\n",
    "\n",
    "DTC = DecisionTreeClassifier(random_state=42)\n",
    "DTC.fit(X[:5000], y[:5000])\n",
    "preds = DTC.predict(X[5000:])\n",
    "print(metrics.classification_report(y[5000:], preds, digits=4))\n",
    "#classifier_model = DTC"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Run Experiment\n",
    "Generate prediction results for the classifier trained on the selected data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_data = pd.read_feather('results_ad_test.feather.with_mse')\n",
    "labels = test_data['labels']\n",
    "mse = test_data['mse']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%time\n",
    "predictions = classifier_model.predict(test_data['features'].to_list())\n",
    "results = pd.DataFrame(\n",
    "    {\n",
    "        'labels': labels,\n",
    "        'mse': mse,\n",
    "        'predictions': predictions\n",
    "    }\n",
    ")\n",
    "#rename the resulting file according to experiment (which classes where used / not used for training, which data were used for testing)\n",
    "results.to_feather('results_bb_without_blacklist_test.feather')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#binarize labels\n",
    "results.labels[results.labels != 'background'] = '1'\n",
    "results.labels[results.labels == 'background'] = '0'"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#rough evaluation\n",
    "print(metrics.classification_report(results.labels, results.predictions, digits=4))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "len(test_data)/16.8"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "colab": {
   "provenance": []
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.7.7"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 1
}