{
 "cells": [
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-24T13:28:00.794477Z",
     "start_time": "2024-02-24T13:28:00.782180Z"
    }
   },
   "cell_type": "code",
   "source": [
    "import os\n",
    "from typing import List, Callable\n",
    "import queue\n",
    "import joblib\n",
    "import threading\n",
    "import time\n",
    "import logging\n",
    "import json\n",
    "\n",
    "import networkx as nx\n",
    "\n",
    "from kafka import KafkaConsumer, KafkaProducer\n",
    "import redis\n",
    "\n",
    "from graph_based_intrusion_detection.utils import constants\n",
    "from graph_based_intrusion_detection.utils.kafka_utils import deserialize, serialize\n",
    "from graph_based_intrusion_detection.utils.logging import create_logger\n",
    "from graph_based_intrusion_detection.utils.state_merging import merge_states\n",
    "from graph_based_intrusion_detection.packet_processing import processing_functions\n",
    "from graph_based_intrusion_detection.packet_processing.processing_functions import process_packets\n",
    "from graph_based_intrusion_detection.graph_processing.data_extraction import create_dataset_for_node"
   ],
   "id": "48d278a5f5448c0e",
   "outputs": [],
   "execution_count": 16
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-23T14:28:05.510736Z",
     "start_time": "2024-02-23T14:28:05.507164Z"
    }
   },
   "cell_type": "code",
   "source": [
    "# define constants\n",
    "log_file_path = os.path.abspath(os.path.join(os.path.dirname(\".\"), \"logs\", \"main.log\"))\n",
    "print(f\"Log file path: {log_file_path}\")\n",
    "\n",
    "update_interval_in_ms = 1000\n",
    "\n"
   ],
   "id": "bbed502fad42ce8f",
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Log file path: /home/alex/projects/it-security-praktikum/devices/dev08-echo-dot-l4s3re/graph_based_intrusion_detection/notebooks/logs/main.log\n"
     ]
    }
   ],
   "execution_count": 11
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-23T14:13:07.915623Z",
     "start_time": "2024-02-23T14:13:07.912275Z"
    }
   },
   "cell_type": "code",
   "source": "# define function for creating logger for each thread\n",
   "id": "44cf10deef903383",
   "outputs": [],
   "execution_count": 8
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7759e6a6c4b91cd6",
   "metadata": {
    "collapsed": true
   },
   "outputs": [],
   "source": [
    "\n",
    "packet_queue = queue.Queue()\n",
    "state_queue = queue.Queue()\n",
    "connection_update_queue = queue.Queue()\n",
    "connection_update_done_queue = queue.Queue()\n",
    "\n",
    "current_state = dict()\n",
    "state_lock = threading.Lock()\n",
    "\n",
    "connection_registry = dict()\n",
    "connection_registry_lock = threading.Lock()\n"
   ]
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "def packet_worker(packet_processing_functions: list, worker_id: int):\n",
    "    \"\"\"\n",
    "    Worker thread, processes packets from the packet queue and puts the resulting state into the state queue\n",
    "    :param packet_processing_functions: function for processing packets\n",
    "    :param worker_id: id of the worker\n",
    "    :return: None\n",
    "    \"\"\"\n",
    "\n",
    "    logger = create_logger(f\"Packet-Worker {worker_id}\")\n",
    "    logger.info(f\"Worker {worker_id} started\")\n",
    "\n",
    "    while True:\n",
    "        # get current batch size based on approximate queue length\n",
    "        n_packets = packet_queue.qsize()\n",
    "        batch_size = n_packets // n_workers\n",
    "        logger.info(f\"Worker {worker_id} current batch size: {batch_size}\")\n",
    "\n",
    "        # fetch batch from queue\n",
    "        batch = list()\n",
    "        try:\n",
    "            for _ in range(batch_size):\n",
    "                batch.append(packet_queue.get_nowait())\n",
    "        except queue.Empty:\n",
    "            pass\n",
    "\n",
    "        logger.info(f\"Worker {worker_id} processing {len(batch)} packets\")\n",
    "        # process batch\n",
    "        if len(batch) > 0:\n",
    "            state = process_packets(batch, packet_processing_functions)\n",
    "            state_queue.put(state)\n",
    "        else:\n",
    "            time.sleep(0.1)\n",
    "\n",
    "        logger.info(f\"Worker {worker_id} done\")\n",
    "\n",
    "        # mark batch as done\n",
    "        for _ in batch:\n",
    "            packet_queue.task_done()"
   ],
   "id": "852ee0a61981f880"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "\n",
    "def merger(merger_function: Callable, worker_id: int):\n",
    "    \"\"\"\n",
    "    Merger thread, merges states from the state queue into one central state representing the state of the network\n",
    "    :param merger_function: function for merging states\n",
    "    :param worker_id: id of the worker\n",
    "    :return: None\n",
    "    \"\"\"\n",
    "    global current_state\n",
    "\n",
    "    logger = create_logger(f\"Merger {worker_id}\")\n",
    "    logger.info(f\"Merger {worker_id} started\")\n",
    "\n",
    "    while True:\n",
    "        # get state lock\n",
    "        state_lock.acquire()\n",
    "\n",
    "        # fetch state from queue\n",
    "        state = state_queue.get()\n",
    "\n",
    "        logger.info(f\"Merging state\")\n",
    "        states = [current_state, state]\n",
    "        # merge state\n",
    "        new_state = merger_function(states)\n",
    "\n",
    "        # update state\n",
    "        current_state = new_state\n",
    "\n",
    "        # mark state as done\n",
    "        state_queue.task_done()\n",
    "\n",
    "        # release state lock\n",
    "        state_lock.release()\n",
    "\n",
    "        logger.info(f\"State merged\")\n"
   ],
   "id": "initial_id"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "def dispatcher(packet: dict,\n",
    "               worker_id: int):\n",
    "    \"\"\"\n",
    "    Dispatcher thread, watches connection registry and dispatches connection to workers for analysis\n",
    "    :param packet: packet to dispatch\n",
    "    :param worker_id: id of the worker\n",
    "    :return: None\n",
    "    \"\"\"\n",
    "\n",
    "    logger = create_logger(f\"Dispatcher {worker_id}\")\n",
    "    logger.info(f\"Dispatcher {worker_id} started\")\n",
    "\n",
    "    while True:\n",
    "        # get connections, based on layer 3 graph -> TODO: extend functionality\n",
    "        # dirty read is not a problem here, modification will be picked up in the next iteration\n",
    "        base_graph = current_state[\"layer_3_graph\"]\n",
    "\n",
    "        # get connections, connections are edges in the graph\n",
    "        connections = list(base_graph.edges())\n",
    "\n",
    "        # get connection registry lock\n",
    "        connection_registry_lock.acquire()\n",
    "\n",
    "        # get connections that are not in the registry\n",
    "        new_connections = [connection for connection in connections if connection not in connection_registry]\n",
    "\n",
    "        # process connection updates since last update\n",
    "        while not connection_update_done_queue.empty():\n",
    "            # get connection from queue, wait at most update interval\n",
    "            try:\n",
    "                connection = connection_update_done_queue.get(timeout=update_interval_in_ms)\n",
    "            except queue.Empty:\n",
    "                break\n",
    "            # update connection in registry, connection is tuple with connection and last updated\n",
    "            connection_registry[connection][\"last_updated\"] = connection[1]\n",
    "\n",
    "        # add new connections to registry\n",
    "        for connection in new_connections:\n",
    "            # add connection to registry, set last updated to None so that it is processed\n",
    "            connection_registry[connection] = {\"last_updated\": None}\n",
    "\n",
    "        # add connections to work queue, if last update is longer ago than the update interval or None\n",
    "        for connection, connection_info in connection_registry.items():\n",
    "            if connection_info[\"last_updated\"] is None or time.time() - connection_info[\n",
    "                \"last_updated\"] > update_interval_in_ms:\n",
    "                connection_update_queue.put(connection)\n",
    "\n",
    "        # release connection registry lock\n",
    "        connection_registry_lock.release()\n"
   ],
   "id": "cca86765185e962c"
  },
  {
   "metadata": {},
   "cell_type": "code",
   "outputs": [],
   "execution_count": null,
   "source": [
    "\n",
    "\n",
    "\n",
    "\n",
    "\n"
   ],
   "id": "b9b68e0a300b6b39"
  },
  {
   "metadata": {
    "ExecuteTime": {
     "end_time": "2024-02-23T14:00:52.816881Z",
     "start_time": "2024-02-23T14:00:52.644607Z"
    }
   },
   "cell_type": "code",
   "source": [
    "\n",
    "n_workers = 1\n",
    "packet_workers = [threading.Thread(target=packet_worker, args=(processing_functions, worker_id))\n",
    "                  for worker_id in range(n_workers)]\n",
    "for worker_thread in packet_workers:\n",
    "    worker_thread.daemon = True\n",
    "\n",
    "merger_thread = threading.Thread(target=merger, args=(merge_states,))\n",
    "merger_thread.daemon = True\n",
    "\n",
    "# start threads\n",
    "# for worker_thread in workers:\n",
    "#     worker_thread.start()\n",
    "# \n",
    "# merger_thread.start()\n"
   ],
   "id": "dde42074dc4deda9",
   "outputs": [
    {
     "ename": "NameError",
     "evalue": "name 'threading' is not defined",
     "output_type": "error",
     "traceback": [
      "\u001B[0;31m---------------------------------------------------------------------------\u001B[0m",
      "\u001B[0;31mNameError\u001B[0m                                 Traceback (most recent call last)",
      "Cell \u001B[0;32mIn[1], line 2\u001B[0m\n\u001B[1;32m      1\u001B[0m n_workers \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[0;32m----> 2\u001B[0m workers \u001B[38;5;241m=\u001B[39m \u001B[43m[\u001B[49m\u001B[43mthreading\u001B[49m\u001B[38;5;241;43m.\u001B[39;49m\u001B[43mThread\u001B[49m\u001B[43m(\u001B[49m\u001B[43mtarget\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43mworker\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43margs\u001B[49m\u001B[38;5;241;43m=\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mprocessing_functions\u001B[49m\u001B[43m,\u001B[49m\u001B[43m \u001B[49m\u001B[43mworker_id\u001B[49m\u001B[43m)\u001B[49m\u001B[43m)\u001B[49m\n\u001B[1;32m      3\u001B[0m \u001B[43m           \u001B[49m\u001B[38;5;28;43;01mfor\u001B[39;49;00m\u001B[43m \u001B[49m\u001B[43mworker_id\u001B[49m\u001B[43m \u001B[49m\u001B[38;5;129;43;01min\u001B[39;49;00m\u001B[43m \u001B[49m\u001B[38;5;28;43mrange\u001B[39;49m\u001B[43m(\u001B[49m\u001B[43mn_workers\u001B[49m\u001B[43m)\u001B[49m\u001B[43m]\u001B[49m\n\u001B[1;32m      4\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m worker_thread \u001B[38;5;129;01min\u001B[39;00m workers:\n\u001B[1;32m      5\u001B[0m     worker_thread\u001B[38;5;241m.\u001B[39mdaemon \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mTrue\u001B[39;00m\n",
      "Cell \u001B[0;32mIn[1], line 2\u001B[0m, in \u001B[0;36m<listcomp>\u001B[0;34m(.0)\u001B[0m\n\u001B[1;32m      1\u001B[0m n_workers \u001B[38;5;241m=\u001B[39m \u001B[38;5;241m1\u001B[39m\n\u001B[0;32m----> 2\u001B[0m workers \u001B[38;5;241m=\u001B[39m [\u001B[43mthreading\u001B[49m\u001B[38;5;241m.\u001B[39mThread(target\u001B[38;5;241m=\u001B[39mworker, args\u001B[38;5;241m=\u001B[39m(processing_functions, worker_id))\n\u001B[1;32m      3\u001B[0m            \u001B[38;5;28;01mfor\u001B[39;00m worker_id \u001B[38;5;129;01min\u001B[39;00m \u001B[38;5;28mrange\u001B[39m(n_workers)]\n\u001B[1;32m      4\u001B[0m \u001B[38;5;28;01mfor\u001B[39;00m worker_thread \u001B[38;5;129;01min\u001B[39;00m workers:\n\u001B[1;32m      5\u001B[0m     worker_thread\u001B[38;5;241m.\u001B[39mdaemon \u001B[38;5;241m=\u001B[39m \u001B[38;5;28;01mTrue\u001B[39;00m\n",
      "\u001B[0;31mNameError\u001B[0m: name 'threading' is not defined"
     ]
    }
   ],
   "execution_count": 1
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 2
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython2",
   "version": "2.7.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
