329 lines
13 KiB
Plaintext
329 lines
13 KiB
Plaintext
{
|
|
"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
|
|
}
|