Knowledge Graph Threat Detection with GNNs

Build a GNN-powered threat detector using a Neo4j knowledge graph.

Introduction

30 Second Summary

Cyberattacks don't happen in isolation. Attackers reuse infrastructure, target clusters of connected victims, and follow patterns that span dozens of relationships. Traditional security tools analyze each network event as an isolated row in a spreadsheet, completely blind to the web of connections that reveals the real threat.

In this project, you will build a threat detection system that models real cyber attack data as a knowledge graph and uses a Graph Neural Network to predict which network entities are malicious based on their connections.

What You'll Build

You feed raw cyber threat data into a pipeline that constructs a knowledge graph of who-talks-to-whom and what-attacks-what, then a GraphSAGE model learns which patterns of connectivity signal danger and prints its predictions alongside accuracy metrics in your notebook.

By the end of this project, you'll have:

  • A live knowledge graph in Neo4j where you can visually trace attack paths and communication patterns between senders, receivers, and threat types.
  • A trained GraphSAGE model that classifies network entities as benign or malicious, with precision, recall, and F1 scores printed in your notebook.
  • A side-by-side comparison showing how your GNN outperforms a traditional Random Forest by leveraging graph connectivity that flat tabular data cannot capture.
  • Secret Mission: A temporal link prediction model that forecasts which entities are likely to be involved in future attacks.

Are there any prerequisites?

This project assumes basic Python familiarity and some prior exposure to machine learning concepts like training, features, and classification. You don't need to be an expert. The project explains each new concept (knowledge graphs, GNNs, Cypher) as it appears.

Before We Start

Before we dive in, let's lock in what you're about to build and why it matters.

Traditional machine learning treats each network event as an isolated row in a spreadsheet. In this project, you'll take a different approach: modelling cyber threats as a connected graph, then training a Graph Neural Network to detect attacks by understanding the relationships between entities.

Set Up Google Colab and Install Dependencies

Training a Graph Neural Network on thousands of nodes and edges requires serious compute power. Google Colab gives you free access to a GPU in the browser, which is exactly what you need to train your GraphSAGE model later in this project.

In this step, you will create a Colab notebook, enable GPU acceleration, and install the libraries your project depends on. This includes PyTorch Geometric for graph neural networks, the Neo4j Python driver for connecting to your knowledge graph database, and kagglehub for downloading the dataset.

In this step, get ready to:
  • Create a Google Colab notebook with a T4 GPU runtime.
  • Install PyTorch Geometric, the Neo4j driver, and kagglehub.
  • Verify all core libraries import without errors.
Create a new Colab notebook with GPU

Google Colab provides a free cloud-based Python notebook environment. You need a GPU runtime because GNN training involves heavy matrix operations across the entire graph structure.

  • Go to colab.research.google.com and sign in with your Google account.
  • Click New notebook to create a fresh notebook.
  • Click Runtime in the top menu bar.
  • Select Change runtime type.
  • In the Hardware accelerator dropdown, select T4 GPU.
  • Click Save.

Why a T4 GPU?

The T4 is a free GPU option in Colab that dramatically speeds up tensor operations. Training a GNN on CPU could take 10x longer. The T4 handles the matrix multiplications in PyTorch Geometric efficiently, keeping each training run under a minute.

Don't see T4 GPU as an option?

GPU availability depends on Colab's current hardware pool. If T4 is not listed, select any available GPU option. The project works with any Colab GPU type.

If no GPU options appear at all, try disconnecting and reconnecting your runtime, or wait a few minutes and try again. Help me enable GPU in Google Colab

Install PyTorch Geometric and project dependencies

PyTorch Geometric requires GPU extension libraries that match your exact PyTorch and CUDA versions. The code below detects those versions automatically so the right wheels are installed.

  • In your new notebook, click into the first code cell and paste the following installation code:
import torch

# Detect PyTorch and CUDA versions for compatible GPU extensions
TORCH = torch.__version__.split('+')[0]
CUDA = 'cu' + torch.version.cuda.replace('.', '')
print(f"Detected PyTorch {TORCH} with CUDA {CUDA}")

# Install PyTorch Geometric core library
!pip install torch_geometric

# Install GPU-accelerated extensions matched to your runtime
!pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html

# Install Neo4j Python driver and Kaggle dataset access
!pip install neo4j kagglehub

What does this code do?

  • torch.__version__ and torch.version.cuda detect which PyTorch and CUDA versions Colab is currently running.
  • The -f flag points pip at a custom wheel repository on data.pyg.org that hosts pre-built binaries for each PyTorch/CUDA combination.
  • torch_geometric is the core Graph Neural Network library. The extension packages (torch_scatter, torch_sparse, etc.) provide GPU-accelerated sparse operations that GraphSAGE depends on.
  • neo4j is the official Python driver for connecting to your Neo4j AuraDB knowledge graph.
  • kagglehub lets you download datasets from Kaggle programmatically without manual file uploads.
  • Run the cell by pressing Shift+Enter.
  • Wait for the installation to complete. You will see pip output scrolling as packages download and install. The first line printed should show your detected PyTorch and CUDA versions.

Seeing installation errors?

If you see a version mismatch error, your Colab runtime may have updated its PyTorch version since the cell was written. Try restarting the runtime (Runtime > Restart runtime) and running the cell again. The version detection code adapts automatically.

If torch_sparse fails to install, the project still works for basic GNN training. You can proceed and revisit the extension later. Help me fix PyTorch Geometric installation errors in Colab

Verify all imports work

Before you move on to loading data, confirm that every library your project needs is importable. A single failed import here will cause confusing errors in later steps.

  • Click the + Code button to add a new code cell below the installation cell.
  • Paste the following verification code:
import torch
import torch_geometric
import neo4j
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt

# Confirm all core libraries loaded without errors
print("All libraries loaded successfully")

What does this cell verify?

  • torch and torch_geometric confirm your GNN training stack is ready.
  • neo4j confirms you can connect to your knowledge graph database.
  • pandas, sklearn, and matplotlib are pre-installed in Colab but importing them here catches any rare environment issues early.

Before you run this cell, what do you expect to see if everything installed correctly?

  • Run the cell by pressing Shift+Enter.

You should see the message All libraries loaded successfully printed with no error tracebacks above it.

Seeing an ImportError?

If import torch_geometric fails, go back to the installation cell and re-run it. Sometimes Colab needs a runtime restart after installing compiled extensions.

Restart your runtime (Runtime > Restart runtime), then run both cells again in order (install first, then verify). Help me fix this ImportError in my Colab notebook

✔️ Awesome, I've got everything!

Great. Make sure your notebook is saved to Google Drive (it saves automatically, but you can press Ctrl+S to be safe).

ⓧ I'd like to double check the full code

Your notebook should have two code cells. Scroll up and confirm they match the following:

Cell 1 (Installation):

import torch

# Detect PyTorch and CUDA versions for compatible GPU extensions
TORCH = torch.__version__.split('+')[0]
CUDA = 'cu' + torch.version.cuda.replace('.', '')
print(f"Detected PyTorch {TORCH} with CUDA {CUDA}")

# Install PyTorch Geometric core library
!pip install torch_geometric

# Install GPU-accelerated extensions matched to your runtime
!pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-{TORCH}+{CUDA}.html

# Install Neo4j Python driver and Kaggle dataset access
!pip install neo4j kagglehub

Cell 2 (Verification):

import torch
import torch_geometric
import neo4j
import pandas as pd
from sklearn.model_selection import train_test_split
from sklearn.metrics import classification_report
import matplotlib.pyplot as plt

# Confirm all core libraries loaded without errors
print("All libraries loaded successfully")

Your environment is locked and loaded. Next up, you will download the cyber threat dataset and explore the entities and relationships that will become nodes and edges in your knowledge graph.

Load the Dataset and Explore the Threat Landscape

Your libraries are installed and your GPU is ready. Now you need the raw material: real cyber threat data that you can model as a graph.

Before you can build a knowledge graph, you need to understand what entities and relationships exist in the data. This step downloads a real-world cyber threat dataset, reveals the structure you will encode as nodes and edges, and prepares clean entity mappings for graph construction.

In this step, get ready to:
  • Download the Cyber Threat Dataset from Kaggle and load it into a pandas DataFrame.
  • Explore the label distribution, unique sender/receiver counts, and identify graph entities.
  • Clean the data and create integer ID mappings for all entities.
Download the dataset from Kaggle

The dataset you will use is the "Cyber Threat Dataset: Network, Text & Relation" by ramoliyafenil on Kaggle. It contains network traffic data, textual data, and entity relationships that map directly to a knowledge graph structure.

  • Download the dataset and load it into a DataFrame by running this code in a new cell:
import kagglehub
import pandas as pd
import os

# Download the cyber threat dataset from Kaggle
path = kagglehub.dataset_download('ramoliyafenil/text-based-cyber-threat-detection')

# Find the CSV file in the downloaded directory
csv_file = [f for f in os.listdir(path) if f.endswith('.csv')][0]
print(f"Dataset downloaded to: {path}")
print(f"CSV file found: {csv_file}")

# Load into a pandas DataFrame
df = pd.read_csv(os.path.join(path, csv_file))
print(f"Dataset shape: {df.shape}")

What does this code do?

The kagglehub.dataset_download function downloads the full dataset to a local cache directory and returns the path. You then scan that directory for the CSV file and load it into a DataFrame called df.

  • Run the cell. You should see the download path, CSV filename, and a shape tuple printed (rows, columns).

Seeing an authentication error?

Kaggle requires authentication to download datasets. In Google Colab, you can authenticate by running kagglehub.login() in a cell before the download, or by uploading your kaggle.json API token file. Make sure your Kaggle account is set up from Step 1.

Help me fix my Kaggle authentication error in Google Colab

Inspect the dataset structure

Now you need to see what columns exist and what types of threats the data contains. This tells you which entities become nodes and which connections become edges in your knowledge graph.

  • Inspect the columns and data types by running this in a new cell:
# Inspect column names, data types, and non-null counts
df.info()
  • Run the cell. You should see columns including sender_id, receiver_ids, label, start-offset, and end-offset listed with their data types.
  • Explore the threat type distribution by running this in a new cell:
# Show distribution of threat types (labels)
print("Label distribution:")
print(df['label'].value_counts())
print(f"\nUnique senders: {df['sender_id'].nunique()}")
print(f"Unique labels: {df['label'].nunique()}")

What do these numbers tell you?

The label column shows the type of cyber threat for each record. The unique sender count tells you how many distinct entities are initiating communications. These senders and receivers become nodes in your knowledge graph, and the communications become edges.

The label distribution also reveals class imbalance. If one threat type dominates, your GNN will need careful handling during training to avoid bias toward the majority class.

  • Run the cell. You should see a count of records per threat type and the unique entity counts printed.

Seeing a KeyError on a column name?

Column names may vary slightly depending on the dataset version. Run print(df.columns.tolist()) to see the exact column names in your download and adjust accordingly.

Help me fix a KeyError when accessing my DataFrame columns

Clean the data and create entity mappings

Raw data often has missing values and multi-valued fields. The receiver_ids column may contain multiple receivers per row, separated by delimiters. You need to parse these into individual relationships and assign each unique entity an integer ID for the graph.

  • Handle missing values and parse receiver IDs by running this in a new cell:
# Drop rows with missing critical fields
df = df.dropna(subset=['sender_id', 'receiver_ids', 'label'])
print(f"Shape after dropping missing values: {df.shape}")

# Parse receiver_ids into individual relationships
# Split comma-separated receiver IDs into lists
df['receiver_list'] = df['receiver_ids'].astype(str).str.split(',')

# Check how many individual relationships exist
total_edges = df['receiver_list'].apply(len).sum()
print(f"Total individual sender-receiver relationships: {total_edges}")

Why parse receiver_ids?

A single row might show one sender communicating with multiple receivers. In graph terms, that is multiple edges from one node. By splitting the receiver list, you create one edge per sender-receiver pair, which is what Neo4j and PyTorch Geometric expect.

  • Run the cell. You should see the cleaned shape and total relationship count printed.
  • Create the entity-to-ID mappings by running this in a new cell:
# Collect all unique entities
all_senders = df['sender_id'].unique().tolist()

# Flatten all individual receivers across all rows
all_receivers = list(set(
    r.strip() for receivers in df['receiver_list']
    for r in receivers
))

all_labels = df['label'].unique().tolist()

# Create integer ID mappings for each entity type
sender_to_id = {s: i for i, s in enumerate(all_senders)}
receiver_to_id = {r: i for i, r in enumerate(all_receivers)}
label_to_id = {l: i for i, l in enumerate(all_labels)}

print(f"Sender entities: {len(sender_to_id)}")
print(f"Receiver entities: {len(receiver_to_id)}")
print(f"Threat label types: {len(label_to_id)}")
print(f"\nLabel mapping: {label_to_id}")

What are entity mappings for?

Graph databases and GNN frameworks work with integer node IDs, not string identifiers. These dictionaries translate between human-readable entity names (like sender IDs and threat labels) and the integer indices that Neo4j and PyTorch Geometric consume. You will use these mappings in the next step when building the knowledge graph.

  • Run the cell. You should see the count of unique senders, receivers, and threat labels printed, along with the label-to-integer mapping dictionary.

Seeing duplicate receivers or unexpected counts?

Make sure you are stripping whitespace from parsed receiver IDs. The r.strip() call removes leading and trailing spaces that could create false duplicates like " node1" vs "node1".

Help me debug unexpected entity counts in my mapping dictionaries

Before you move on, what do you think you will see if you run df.info() and df['label'].value_counts() together now?

  • Verify the final state of your DataFrame by running this in a new cell:
# Final verification of the cleaned dataset
df.info()
print("\n")
print(df['label'].value_counts())

You should see the dataset columns (including your new receiver_list column), the data types, and the distribution of threat types printed. The non-null counts should be consistent across all rows since you dropped missing values earlier.

✔️ Awesome, I've got everything!

Great. You have a cleaned DataFrame, a parsed receiver list, and three entity mapping dictionaries ready for graph construction.

ⓧ I'd like to double check the full code

Here are the cells you should have in order. Confirm each one is present and runs without error:

  • Cell 1: Dataset download with kagglehub.dataset_download and CSV loading into df.
  • Cell 2: df.info() to inspect columns and types.
  • Cell 3: df['label'].value_counts() and unique entity counts.
  • Cell 4: Missing value handling and receiver_list parsing.
  • Cell 5: Entity mapping dictionaries (sender_to_id, receiver_to_id, label_to_id).
  • Cell 6: Final verification with df.info() and df['label'].value_counts().

Your data is clean, explored, and mapped. Next up, you will load these entities and relationships into a Neo4j knowledge graph where you can query and visualize the attack patterns.

Build the Knowledge Graph in Neo4j AuraDB

Your DataFrame holds clean, structured threat data with senders, receivers, and attack labels all mapped to unique IDs. But right now, every row sits in isolation.

A flat table can't answer questions like "which receivers are targeted by multiple attackers sharing the same threat type?" because those answers live in the connections between entities. By loading this data into Neo4j as a knowledge graph, you turn those invisible relationships into something you can query, visualize, and feed into a graph neural network.

In this step, get ready to:
  • Connect to your Neo4j AuraDB instance from Google Colab using the Python driver.
  • Create Cypher queries that define Sender, Receiver, and ThreatType nodes plus COMMUNICATES_WITH and CLASSIFIED_AS edges.
  • Batch-load the full dataset and verify the graph is populated with node and relationship counts.
Connect to Neo4j AuraDB from Colab

The neo4j Python driver lets you run Cypher queries against your AuraDB instance directly from a Colab cell. You installed the driver in Step 1, so now you just need your connection credentials.

  • In your Colab notebook, create a new code cell and establish the connection by pasting this code:
from neo4j import GraphDatabase

# Replace with your Neo4j AuraDB credentials from the Aura Console
NEO4J_URI = "neo4j+s://your-auradb-uri-here"
NEO4J_USERNAME = "neo4j"
NEO4J_PASSWORD = "your-password-here"

# Connect to the AuraDB instance
driver = GraphDatabase.driver(NEO4J_URI, auth=(NEO4J_USERNAME, NEO4J_PASSWORD))

# Verify the connection works
driver.verify_connectivity()
print("Connected to Neo4j AuraDB successfully")

What does this code do?

  • GraphDatabase.driver() creates a driver instance that manages connections to your AuraDB database.
  • The neo4j+s:// protocol prefix tells the driver to use an encrypted connection, which AuraDB requires.
  • verify_connectivity() immediately tests the connection and credentials before you try to load any data.
  • Replace your-auradb-uri-here with the connection URI from your Aura Console (it looks like xxxxxxxx.databases.neo4j.io).
  • Replace your-password-here with the password you saved when creating the instance.
  • Run the cell. You should see Connected to Neo4j AuraDB successfully printed below.

Seeing a connection error?

  • Double-check that your URI starts with neo4j+s:// and not bolt:// or https://. AuraDB requires the neo4j+s scheme.
  • Make sure your AuraDB instance is in the "Running" state in the Aura Console. Free instances auto-pause after 3 days of inactivity.
  • Confirm the password matches exactly what was generated at instance creation time. Passwords are case-sensitive.

Why can't I connect to my Neo4j AuraDB instance from Colab?

Create Cypher queries that build the knowledge graph schema

Your knowledge graph needs three types of nodes (Sender, Receiver, ThreatType) and two types of relationships (COMMUNICATES_WITH, CLASSIFIED_AS). You will batch-create all the nodes first, then wire up the edges.

  • Create a new code cell and batch-load all three node types by running this code:
# Prepare node data from entity mappings
sender_batch = [{"sender_id": name} for name in sender_to_id.keys()]
receiver_batch = [{"receiver_id": name} for name in receiver_to_id.keys()]
threat_batch = [{"name": name} for name in label_to_id.keys()]

# Batch-create all node types using UNWIND
with driver.session() as session:
    session.execute_write(
        lambda tx: tx.run("""
            UNWIND $batch AS item
            MERGE (s:Sender {sender_id: item.sender_id})
        """, batch=sender_batch)
    )
    session.execute_write(
        lambda tx: tx.run("""
            UNWIND $batch AS item
            MERGE (r:Receiver {receiver_id: item.receiver_id})
        """, batch=receiver_batch)
    )
    session.execute_write(
        lambda tx: tx.run("""
            UNWIND $batch AS item
            MERGE (t:ThreatType {name: item.name})
        """, batch=threat_batch)
    )

print(f"Created {len(sender_batch)} Sender nodes")
print(f"Created {len(receiver_batch)} Receiver nodes")
print(f"Created {len(threat_batch)} ThreatType nodes")

What does this code do?

  • UNWIND takes a list of items and processes each one individually inside the query. This is far faster than running a separate query per node.
  • MERGE creates a node only if one with the same property doesn't already exist. This makes the query safe to re-run without creating duplicates.
  • Each node type gets a label (Sender, Receiver, ThreatType) and a unique identifying property so you can match them when creating edges.
  • Run the cell. You should see three lines printing the count of each node type created.

Now you need relationships connecting these nodes. In your knowledge graph, senders communicate with receivers and senders are classified under threat types.

  • Create a new code cell and batch-load both relationship types by running this code:
# Prepare edge data from the DataFrame
comm_edges = df[["sender_id", "receiver_id"]].drop_duplicates().to_dict("records")
class_edges = df[["sender_id", "label"]].drop_duplicates().to_dict("records")

# Batch-create relationships using UNWIND
with driver.session() as session:
    session.execute_write(
        lambda tx: tx.run("""
            UNWIND $batch AS item
            MATCH (s:Sender {sender_id: item.sender_id})
            MATCH (r:Receiver {receiver_id: item.receiver_id})
            MERGE (s)-[:COMMUNICATES_WITH]->(r)
        """, batch=comm_edges)
    )
    session.execute_write(
        lambda tx: tx.run("""
            UNWIND $batch AS item
            MATCH (s:Sender {sender_id: item.sender_id})
            MATCH (t:ThreatType {name: item.label})
            MERGE (s)-[:CLASSIFIED_AS]->(t)
        """, batch=class_edges)
    )

print(f"Created {len(comm_edges)} COMMUNICATES_WITH relationships")
print(f"Created {len(class_edges)} CLASSIFIED_AS relationships")

What does this code do?

  • MATCH finds existing nodes by their identifying property so the relationship connects the correct pair.
  • COMMUNICATES_WITH edges connect a Sender to each Receiver it contacted. This encodes the network communication pattern.
  • CLASSIFIED_AS edges connect a Sender to the ThreatType label associated with its activity. This encodes the threat categorization.
  • drop_duplicates() ensures you only create one edge per unique pair, keeping the graph clean.
  • Run the cell. You should see two lines printing the count of each relationship type created.

Seeing zero relationships created?

  • Check that the column names in df[["sender_id", "receiver_id"]] match your actual DataFrame columns exactly. Run df.columns.tolist() to confirm.
  • Make sure you ran the node creation cell first. Edges require matching nodes to exist.

Why are my Neo4j relationship counts showing zero?

Verify the knowledge graph is populated

The nodes and edges are loaded. Now you need to confirm the graph structure matches what you expect from the dataset.

  • Before you run the verification query, predict: how many ThreatType nodes do you expect based on the number of unique labels you saw in Step 2?
  • Create a new code cell and run these verification queries:
# Verify the knowledge graph is populated
with driver.session() as session:
    # Count nodes by type
    result = session.run("""
        MATCH (n)
        RETURN labels(n)[0] AS type, count(n) AS count
    """)
    print("Node counts:")
    for record in result:
        print(f"  {record['type']}: {record['count']}")

    # Count relationships by type
    result = session.run("""
        MATCH ()-[r]->()
        RETURN type(r) AS type, count(r) AS count
    """)
    print("\nRelationship counts:")
    for record in result:
        print(f"  {record['type']}: {record['count']}")

    # Sample paths showing the graph structure
    result = session.run("""
        MATCH (s:Sender)-[:COMMUNICATES_WITH]->(r:Receiver)
        MATCH (s)-[:CLASSIFIED_AS]->(t:ThreatType)
        RETURN s.sender_id AS sender, r.receiver_id AS receiver,
               t.name AS threat
        LIMIT 5
    """)
    print("\nSample paths:")
    for record in result:
        print(f"  {record['sender']} -> {record['receiver']} [{record['threat']}]")

What does this code do?

  • The first query counts all nodes grouped by their label. You should see entries for Sender, Receiver, and ThreatType.
  • The second query counts all relationships grouped by type. You should see COMMUNICATES_WITH and CLASSIFIED_AS.
  • The third query walks a full path through the graph: from a Sender, through a COMMUNICATES_WITH edge to a Receiver, and through a CLASSIFIED_AS edge to a ThreatType. This proves the graph is connected end-to-end.
  • Run the cell. You should see node counts for Sender, Receiver, and ThreatType, relationship counts for COMMUNICATES_WITH and CLASSIFIED_AS, and five sample paths showing real attack patterns in your graph.

The sample paths show the core structure of your knowledge graph in action. Each path represents a real threat event: a sender communicated with a receiver and was classified under a specific threat type.

Node counts don't match your expectations?

  • Compare the printed Sender count against len(sender_to_id) from Step 2. They should match exactly.
  • If you see more nodes than expected, you may have run the creation cell twice with slightly different data. Clear the database with MATCH (n) DETACH DELETE n and re-run the node and edge creation cells.

Why do my Neo4j node counts not match my entity mapping dictionaries?

Your knowledge graph is live and queryable. Every sender, receiver, and threat type is connected through real relationships extracted from your dataset. Next up, you will convert this graph structure into a format that PyTorch Geometric can train on.

Extract Graph Features and Build the PyG Dataset

Your knowledge graph is alive in Neo4j with senders, receivers, and threat types connected by real relationships. Now you need to get that graph structure into a format that a neural network can learn from.

PyTorch Geometric expects graph data as tensors: an edge list, a node feature matrix, and labels. In this step, you'll query Neo4j to extract the full adjacency structure, engineer features from each node's connectivity, and package everything into a single Data object ready for training.

In this step, get ready to:
  • Query Neo4j for all edges and build a node feature matrix from degree, type encoding, and communication frequency.
  • Convert the extracted data into a PyTorch Geometric Data object with edge_index, x, and y tensors.
  • Create stratified train/validation/test masks that handle class imbalance.
Query Neo4j and map edges to a unified node index

Your knowledge graph has three types of nodes stored in separate Neo4j label groups. For PyTorch Geometric, you need every node in a single contiguous index space so edges can be represented as integer pairs.

  • Create the unified node index and query all edges from Neo4j by running this code in a new Colab cell:
import numpy as np
from collections import Counter

# Define unified node index offsets
num_senders = len(sender_to_id)
num_receivers = len(receiver_to_id)
num_threats = len(label_to_id)
total_nodes = num_senders + num_receivers + num_threats

# Query edges from Neo4j and map to unified indices
with driver.session() as session:
    # COMMUNICATES_WITH edges (sender -> receiver)
    comm_result = session.run("""
        MATCH (s:Sender)-[:COMMUNICATES_WITH]->(r:Receiver)
        RETURN s.sender_id AS src, r.receiver_id AS tgt
    """)
    comm_edges = [(sender_to_id[rec["src"]],
                   num_senders + receiver_to_id[rec["tgt"]])
                  for rec in comm_result]

    # CLASSIFIED_AS edges (sender -> threat type)
    class_result = session.run("""
        MATCH (s:Sender)-[:CLASSIFIED_AS]->(t:ThreatType)
        RETURN s.sender_id AS src, t.name AS tgt
    """)
    class_edges = [(sender_to_id[rec["src"]],
                    num_senders + num_receivers + label_to_id[rec["tgt"]])
                   for rec in class_result]

all_edges = comm_edges + class_edges
print(f"Total nodes: {total_nodes} (Senders: {num_senders}, Receivers: {num_receivers}, ThreatTypes: {num_threats})")
print(f"Total edges: {len(all_edges)} (COMMUNICATES_WITH: {len(comm_edges)}, CLASSIFIED_AS: {len(class_edges)})")

What does this code do?

The unified index places all node types into one number line: senders occupy indices 0 through num_senders - 1, receivers start at num_senders, and threat types start at num_senders + num_receivers.

The two Cypher queries pull every relationship from Neo4j and map them into integer pairs using the mapping dictionaries from the previous step. This gives you a flat edge list where every source and target is a simple integer.

  • Run the cell. You should see the total node count and edge counts for both relationship types printed.

Seeing a KeyError on sender or receiver IDs?

This means a node in Neo4j has an ID that doesn't exist in your mapping dictionaries. Confirm that sender_to_id and receiver_to_id were built from the same cleaned DataFrame you used to load the graph.

Re-run the cells in Step 3 that create the entity mappings, then try this cell again. Why am I getting a KeyError when querying Neo4j edges?

Build the node feature matrix

A GNN needs numeric features for each node. You'll engineer three features that capture each node's role in the network: its degree (how many connections it has), its type encoding (sender, receiver, or threat), and its normalized communication frequency.

  • Build the feature matrix by running this code in a new cell:
# Build node feature matrix: degree, type encoding, communication frequency
degree = np.zeros(total_nodes)
for src, tgt in all_edges:
    degree[src] += 1
    degree[tgt] += 1

# Type encoding: 0=sender, 1=receiver, 2=threat_type
node_type = np.zeros(total_nodes)
node_type[num_senders:num_senders + num_receivers] = 1
node_type[num_senders + num_receivers:] = 2

# Communication frequency (normalized degree within each type)
comm_freq = np.zeros(total_nodes)
sender_max = degree[:num_senders].max() + 1e-8
receiver_max = degree[num_senders:num_senders + num_receivers].max() + 1e-8
comm_freq[:num_senders] = degree[:num_senders] / sender_max
comm_freq[num_senders:num_senders + num_receivers] = (
    degree[num_senders:num_senders + num_receivers] / receiver_max
)

# Stack into feature matrix
features = np.stack([degree, node_type, comm_freq], axis=1)
print(f"Feature matrix shape: {features.shape}")
print(f"Features per node: [degree, type_encoding, comm_frequency]")

Why these three features?

Degree tells the GNN how connected a node is. Highly connected senders might be scanners or botnet controllers.

Type encoding distinguishes node roles so the model can learn different patterns for senders versus receivers versus threat categories.

Communication frequency normalizes the degree within each type, preventing senders (which naturally have more edges) from dominating receivers in the feature space.

  • Run the cell. You should see Feature matrix shape: followed by (total_nodes, 3) confirming three features per node.

Feature matrix shape shows (0, 3)?

This means total_nodes is zero. Check that the previous cell ran successfully and that num_senders, num_receivers, and num_threats all have non-zero values.

Help me debug why my total_nodes is zero when building the feature matrix.

Create the PyTorch Geometric Data object

Now you'll assign labels and convert everything into the tensor format that PyTorch Geometric expects. Each node gets a binary label: malicious (1) if it participates in attack activity, or benign (0) otherwise.

  • Create node labels and build the PyG Data object by running this code in a new cell:
import torch
from torch_geometric.data import Data

# Create binary labels: malicious (1) vs benign (0)
node_labels = np.zeros(total_nodes, dtype=np.int64)

# ThreatType nodes representing attacks are malicious
for label, idx in label_to_id.items():
    if label.lower() != 'benign':
        node_labels[num_senders + num_receivers + idx] = 1

# Senders connected to malicious threats are malicious
for src, tgt in class_edges:
    if node_labels[tgt] == 1:
        node_labels[src] = 1

# Receivers connected to malicious senders are malicious
for src, tgt in comm_edges:
    if node_labels[src] == 1:
        node_labels[tgt] = 1

# Convert to PyG tensors
edge_index = torch.tensor(all_edges, dtype=torch.long).t().contiguous()
# Make graph undirected by adding reverse edges
edge_index = torch.cat([edge_index, edge_index.flip(0)], dim=1)
x = torch.tensor(features, dtype=torch.float)
y = torch.tensor(node_labels, dtype=torch.long)

# Build the PyG Data object
data = Data(x=x, edge_index=edge_index, y=y)
print(data)
print(f"\nLabel distribution: {Counter(node_labels.tolist())}")

What does this code do?

The labelling logic propagates maliciousness through the graph. First, threat type nodes that represent actual attacks get labelled malicious. Then any sender connected to those threats inherits the label. Finally, receivers connected to malicious senders get marked too.

The edge_index.flip(0) call adds reverse edges, making the graph undirected. This is important because GraphSAGE aggregates information from neighbors in both directions.

  • Run the cell. You should see a Data(x=[...], edge_index=[2, ...], y=[...]) summary plus the count of malicious vs benign nodes.

Seeing all zeros in the label distribution?

If every node is labelled 0, none of your threat types matched the malicious condition. Print label_to_id.keys() to check the exact label names in your dataset. The comparison uses label.lower() != 'benign', so confirm whether your dataset uses that exact string for safe traffic.

Why are all my node labels zero when I expected some malicious nodes?

Create stratified train/validation/test masks

Class imbalance is common in cybersecurity data. Most network traffic is benign, so a random split could leave your validation or test set with too few malicious examples. Stratified splitting preserves the ratio of malicious to benign nodes in every subset.

  • Create the masks by running this code in a new cell:
from sklearn.model_selection import train_test_split

# Stratified split: 60% train, 20% validation, 20% test
indices = np.arange(total_nodes)
train_idx, temp_idx = train_test_split(
    indices, test_size=0.4, stratify=node_labels, random_state=42
)
val_idx, test_idx = train_test_split(
    temp_idx, test_size=0.5, stratify=node_labels[temp_idx], random_state=42
)

# Create boolean masks on the Data object
data.train_mask = torch.zeros(total_nodes, dtype=torch.bool)
data.val_mask = torch.zeros(total_nodes, dtype=torch.bool)
data.test_mask = torch.zeros(total_nodes, dtype=torch.bool)
data.train_mask[train_idx] = True
data.val_mask[val_idx] = True
data.test_mask[test_idx] = True

print(f"Train nodes: {data.train_mask.sum().item()}")
print(f"Validation nodes: {data.val_mask.sum().item()}")
print(f"Test nodes: {data.test_mask.sum().item()}")

Why stratified splitting matters here

The stratify parameter ensures each split (train, validation, test) contains the same proportion of malicious nodes as the full dataset. Without it, a small test set might contain zero malicious examples, making your evaluation metrics meaningless.

The two-stage split first separates 60% for training, then divides the remaining 40% equally between validation and test.

  • Run the cell. You should see the node counts for each split printed.

Getting a ValueError about stratify?

If the error says the least populated class has too few members, your label propagation produced very few malicious nodes. Check the label distribution from the previous cell. You need at least 3 malicious nodes (one per split) for stratification to work.

Help me fix a ValueError when using stratified train_test_split on my node labels.

Before you verify the final state, what do you expect the Data object to show now that masks have been added?

  • Verify the complete Data object by running this in a new cell:
# Final verification: print the complete PyG Data object
print("PyG Data object ready for training:")
print(data)

You should see the Data object summary including x, edge_index, y, train_mask, val_mask, and test_mask all with correct dimensions matching your node count.

✔️ I see the Data object with all tensors and masks

Your graph is now in the exact format PyTorch Geometric needs for training. The edge structure, node features, labels, and data splits are all packaged in a single object.

ⓧ Some fields are missing from the Data printout

If train_mask, val_mask, or test_mask are missing, re-run the mask creation cell. Colab cells must be run in order for variables to persist between cells.

If edge_index or x are missing, scroll up and re-run the Data object creation cell. Then re-run the mask cell since it modifies the same data variable.

Your entire knowledge graph is now a trainable PyG dataset with features, labels, and balanced splits. Next up, you'll feed this into a GraphSAGE model that learns to detect threats by aggregating information from each node's neighborhood.

Train a GraphSAGE Model and Detect Threats

Your knowledge graph is now encoded as tensors that PyTorch Geometric can work with. The next question is: can a neural network actually learn which nodes are threats by looking at their connections?

This is the ML heart of the project. GraphSAGE learns by aggregating information from each node's neighborhood. A node connected to many known-malicious senders gets a different embedding than one surrounded by benign connections. That structural signal is what gives graph-based ML its edge over traditional approaches.

In this step, get ready to:
  • Define a 2-layer GraphSAGE model with a threat classification head.
  • Train the model on GPU for 200 epochs and track validation accuracy.
  • Compare GraphSAGE against a Random Forest baseline to prove graph structure helps.
Define the GraphSAGE model architecture

GraphSAGE works by sampling and aggregating features from a node's neighbors. With two SAGEConv layers, the model sees 2 hops outward from each node.

  • In a new code cell in your Colab notebook, define the GraphSAGE model class by pasting the following code:
import torch
import torch.nn.functional as F
from torch_geometric.nn import SAGEConv

class GraphSAGEThreatDetector(torch.nn.Module):
    def __init__(self, in_channels, hidden_channels, out_channels):
        super().__init__()
        # First SAGEConv: aggregates 1-hop neighborhood
        self.conv1 = SAGEConv(in_channels, hidden_channels)
        # Second SAGEConv: aggregates 2-hop neighborhood
        self.conv2 = SAGEConv(hidden_channels, hidden_channels)
        # Classification head maps embeddings to class probabilities
        self.classifier = torch.nn.Linear(hidden_channels, out_channels)

    def forward(self, x, edge_index):
        # Layer 1: aggregate neighbors, activate, dropout
        x = self.conv1(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        # Layer 2: aggregate 2-hop neighbors
        x = self.conv2(x, edge_index)
        x = F.relu(x)
        x = F.dropout(x, p=0.5, training=self.training)
        return self.classifier(x)

    def get_embeddings(self, x, edge_index):
        """Extract node embeddings before classification head."""
        x = F.relu(self.conv1(x, edge_index))
        x = F.relu(self.conv2(x, edge_index))
        return x

What does this model do?

Each SAGEConv layer gathers feature information from neighboring nodes and combines it. After two layers, every node's representation encodes patterns from 2 hops away in the graph.

The classifier layer is a simple linear head that maps those learned embeddings to threat class probabilities. The get_embeddings method lets you extract the raw node representations for later use.

  • Run the cell. You should see no output, which means the class is defined and ready.

Seeing an import error?

If from torch_geometric.nn import SAGEConv fails, your PyTorch Geometric installation from Step 1 may not have completed properly. Re-run the installation cell and restart your runtime (Runtime > Restart runtime), then re-run all cells from the top.

Need more help? Why is SAGEConv not importing in my Colab notebook?

Train the model on GPU

With the architecture defined, you can now move everything to the GPU and train. The training loop computes loss only on training-masked nodes, leaving validation and test nodes untouched.

  • In a new code cell, paste the training loop that initializes the model and trains for 200 epochs:
# Move data to GPU and initialize model
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
data = data.to(device)
num_features = data.x.shape[1]
num_classes = len(torch.unique(data.y))
model = GraphSAGEThreatDetector(num_features, 64, num_classes).to(device)
optimizer = torch.optim.Adam(model.parameters(), lr=0.01, weight_decay=5e-4)
criterion = torch.nn.CrossEntropyLoss()

def train():
    model.train()
    optimizer.zero_grad()
    out = model(data.x, data.edge_index)
    loss = criterion(out[data.train_mask], data.y[data.train_mask])
    loss.backward()
    optimizer.step()
    return loss.item()

def evaluate(mask):
    model.eval()
    with torch.no_grad():
        pred = model(data.x, data.edge_index)[mask].argmax(dim=1)
    return (pred == data.y[mask]).sum().item() / mask.sum().item()

# Train for 200 epochs, print every 20
for epoch in range(1, 201):
    loss = train()
    if epoch % 20 == 0:
        val_acc = evaluate(data.val_mask)
        print(f'Epoch {epoch:03d}, Loss: {loss:.4f}, Val Acc: {val_acc:.4f}')

What is happening during training?

Each epoch, the train() function pushes node features through the two SAGEConv layers, computes cross-entropy loss on nodes in the training mask, and updates the model weights via backpropagation.

The evaluate() function checks accuracy on the validation mask every 20 epochs so you can watch the model improve. Only training-masked nodes contribute to the gradient. Validation and test nodes stay hidden from the optimizer.

  • Run the cell. You should see 10 lines of output (one every 20 epochs) showing the loss decreasing and validation accuracy increasing.

You should see validation accuracy climbing over the training run, ending well above random chance.

Loss not decreasing?

If loss stays flat or increases, confirm your runtime is set to T4 GPU (Runtime > Change runtime type). Also verify that data.train_mask contains True values by running data.train_mask.sum() in a new cell.

Still stuck? My GraphSAGE loss is not decreasing during training. How do I debug this?

Evaluate and compare against a baseline

A trained model's accuracy numbers mean nothing without a baseline to compare against. You will train a Random Forest on the exact same node features but without any graph structure.

  • In a new code cell, evaluate the trained GraphSAGE model on the test set by pasting this code:
from sklearn.metrics import classification_report

# Get GraphSAGE predictions on test set
model.eval()
with torch.no_grad():
    out = model(data.x, data.edge_index)
    gnn_preds = out[data.test_mask].argmax(dim=1).cpu().numpy()
    true_labels = data.y[data.test_mask].cpu().numpy()

# Store report for comparison
gnn_report = classification_report(true_labels, gnn_preds)
print("=== GraphSAGE Classification Report ===")
print(gnn_report)

What does this report show?

The classification report shows precision, recall, and F1-score for each threat class. Precision measures how many predicted threats were real. Recall measures how many real threats were caught. F1 balances the two.

  • Run the cell and confirm you see a table of per-class metrics.

Before you run the baseline comparison, think about this: the Random Forest gets the same node features (degree, frequency, threat encoding) but has no visibility into who connects to whom. Do you think graph connections will make a measurable difference in F1?

  • In a new code cell, train the Random Forest baseline and print a side-by-side comparison:
from sklearn.ensemble import RandomForestClassifier

# Train Random Forest on flat node features (no graph structure)
X_train = data.x[data.train_mask].cpu().numpy()
y_train = data.y[data.train_mask].cpu().numpy()
X_test = data.x[data.test_mask].cpu().numpy()
y_test = data.y[data.test_mask].cpu().numpy()

rf = RandomForestClassifier(n_estimators=100, random_state=42)
rf.fit(X_train, y_train)
rf_preds = rf.predict(X_test)

# Print baseline report
rf_report = classification_report(y_test, rf_preds)
print("=== Random Forest Baseline (No Graph Structure) ===")
print(rf_report)

# Side-by-side comparison
print("\n" + "="*50)
print("COMPARISON: Graph Structure vs Flat Features")
print("="*50)
print("\nGraphSAGE (uses graph connections):")
print(gnn_report)
print("Random Forest (flat features only):")
print(rf_report)

Why does graph structure help?

The Random Forest sees each node in isolation. It knows a node has high degree and certain frequency patterns, but it cannot see that the node's neighbors are all flagged as malicious.

GraphSAGE aggregates neighbor information, so a benign-looking node connected to a cluster of known attackers gets a very different embedding. This is the structural signal that flat tabular ML misses entirely.

  • Run the cell. Compare the F1-scores across both reports. You should see GraphSAGE achieving a measurably higher F1 on at least one class compared to the Random Forest.

Random Forest performing the same or better?

If the RF matches or beats GraphSAGE, check that your edge_index tensor was built correctly in the previous step. A graph with very few edges gives GraphSAGE nothing extra to work with. Run data.edge_index.shape and confirm the second dimension is much larger than the number of nodes.

Also verify your training ran for the full 200 epochs. Early stopping or a crashed cell can leave the GNN undertrained.

Need help interpreting? Why is my Random Forest baseline outperforming GraphSAGE on this cybersecurity dataset?

✔️ Awesome, I've got everything!

Great. Make sure you have saved your notebook (Ctrl+S or Cmd+S).

ⓧ I'd like to double check the full code

Scroll up through this step and confirm you have four code cells in order:

  • Cell 1: GraphSAGEThreatDetector class definition with conv1, conv2, classifier, forward, and get_embeddings methods.
  • Cell 2: Model initialization, train() and evaluate() helpers, and the 200-epoch training loop printing every 20 epochs.
  • Cell 3: GraphSAGE evaluation with classification_report printed.
  • Cell 4: Random Forest baseline training and the side-by-side comparison print.

You just proved that graph-based ML captures attack patterns that flat features miss. Your trained GraphSAGE model and its node embeddings are ready. In the Secret Mission, you will use those embeddings to predict future attack connections.

Secret mission

Temporal Link Prediction for Future Threat Forecasting

Your GNN can classify current threats. But can it predict the future? In this secret mission, you'll add temporal edge features and build a link prediction model that forecasts which entities are likely to be involved in future attacks. Instead of asking "is this entity dangerous now?", you'll ask "which new attack connections will form next?"

Clean Up Your Resources

Clean Up Your Resources

This project uses entirely free resources, so there are no ongoing costs. Decide whether to keep your resources available, let them auto-pause, or delete them entirely.

Resources you used:

  • Neo4j AuraDB Free instance (your knowledge graph database).
  • Google Colab notebook (saved automatically to your Google Drive).

Keep everything running

No action needed. Choose this if you plan to continue exploring the knowledge graph or extending the GNN model.

  • Your Neo4j AuraDB Free instance will auto-pause after 3 days of inactivity (no write queries). Your data stays safe until you resume it from the Aura Console.
  • Your Colab notebook is saved in your Google Drive under the Colab Notebooks folder. You can reopen it any time.
  • If you leave the AuraDB Free instance paused for more than 30 days without resuming it, Neo4j will automatically delete it and all associated data.

Pause - I'll come back to this later

Neo4j AuraDB Free auto-pauses on its own after 3 days of inactivity. There is nothing you need to do manually.

  • When you are ready to resume, open the Neo4j Aura Console at console.neo4j.io and click the play/resume button on your paused instance card.
  • Remember that a paused Free instance is automatically deleted after 30 days if not resumed.
  • Your Colab notebook remains in Google Drive with no expiry.

Delete - I don't want to use this again

Remove all project resources permanently. This cannot be undone.

  • Delete your Neo4j AuraDB instance:
  • Delete your Colab notebook from Google Drive:

Nice Work!

Nice Work!

You just built a full threat detection pipeline that models cyber attacks as a knowledge graph and uses a Graph Neural Network to predict threats.

You've learned how to:

  • Build a cybersecurity knowledge graph in Neo4j AuraDB from raw threat intelligence data and query attack patterns using Cypher.
  • Train a GraphSAGE Graph Neural Network with PyTorch Geometric that classifies network entities as malicious or benign based on their graph neighborhood.
  • Demonstrate that graph-based ML outperforms traditional flat-feature ML (Random Forest) on cyber threat data by comparing precision, recall, and F1 scores side by side.
  • Secret Mission: Build a temporal link prediction model that forecasts which entities are likely to be involved in future attacks.

Ready to quiz yourself?