graph-test1.py

from neo4j import GraphDatabase
import time
from datetime import datetime, timedelta
import json

# Configuration for connecting to Neo4j
URI = "bolt://localhost:7687"
USERNAME = "neo4j"
PASSWORD = "password"  # Update this with your Neo4j password

class CloudSecurityGraph:
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))
    
    def close(self):
        self.driver.close()
    
    def clear_database(self):
        """Delete all nodes and relationships."""
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
    
    def create_container(self, container_id, risk_score, has_public_connection, services):
        """
        Create a container node with properties.
        lastEventTimestamp is set to the current time in milliseconds.
        """
        timestamp = int(time.time() * 1000)
        query = """
            CREATE (c:Container {
                id: $id,
                riskScore: $riskScore,
                hasPublicConnection: $hasPublicConnection,
                services: $services,
                lastEventTimestamp: $timestamp
            })
            RETURN c
        """
        with self.driver.session() as session:
            result = session.run(query, id=container_id, riskScore=risk_score, 
                                  hasPublicConnection=has_public_connection, 
                                  services=services, timestamp=timestamp)
            return result.single()[0]
    
    def create_communication(self, src_id, dst_id):
        """
        Create a COMMUNICATES_WITH relationship between two container nodes.
        """
        query = """
            MATCH (src:Container {id: $src_id}), (dst:Container {id: $dst_id})
            CREATE (src)-[r:COMMUNICATES_WITH {createdAt: $timestamp}]->(dst)
            RETURN r
        """
        timestamp = int(time.time() * 1000)
        with self.driver.session() as session:
            result = session.run(query, src_id=src_id, dst_id=dst_id, timestamp=timestamp)
            return result.single()[0]
    
    def update_risk_score(self, container_id, new_risk_score):
        """
        Update the riskScore and lastEventTimestamp of a container node.
        """
        timestamp = int(time.time() * 1000)
        query = """
            MATCH (c:Container {id: $id})
            SET c.riskScore = $newRiskScore, c.lastEventTimestamp = $timestamp
            RETURN c
        """
        with self.driver.session() as session:
            result = session.run(query, id=container_id, newRiskScore=new_risk_score, timestamp=timestamp)
            return result.single()[0]
    
    def dynamic_detection_query(self, time_window_minutes=10, risk_threshold=70, anomaly_threshold=80):
        """
        Run a detection query that finds communication paths where:
          - The source container (src) has riskScore above risk_threshold.
          - The destination container (dst) has a public connection and, for example, an implicit "anomalyScore"
            (we'll assume riskScore here for demonstration).
          - Both events happened within the last time_window_minutes.
          
        The composite risk score is computed as a weighted sum of the source and destination risk scores.
        """
        time_window_ms = int(timedelta(minutes=time_window_minutes).total_seconds() * 1000)
        current_ts = int(time.time() * 1000)
        
        query = """
            WITH $time_window_ms AS window, $current_ts AS now, $risk_threshold AS rThreshold, $anomaly_threshold AS aThreshold
            MATCH (src:Container)-[comm:COMMUNICATES_WITH]->(dst:Container)
            WHERE src.lastEventTimestamp >= (now - window)
              AND dst.lastEventTimestamp >= (now - window)
              AND src.riskScore >= rThreshold
              AND dst.hasPublicConnection = true
              AND dst.riskScore >= aThreshold
            RETURN src AS SourceContainer, 
                   dst AS DestinationContainer,
                   ((src.riskScore * 0.7) + (dst.riskScore * 0.3)) AS compositeRiskScore,
                   comm
            ORDER BY compositeRiskScore DESC
        """
        with self.driver.session() as session:
            result = session.run(query, time_window_ms=time_window_ms, current_ts=current_ts, 
                                 risk_threshold=risk_threshold, anomaly_threshold=anomaly_threshold)
            return [record.data() for record in result]

if __name__ == "__main__":
    # Initialize the graph instance and clear existing data
    graph = CloudSecurityGraph(URI, USERNAME, PASSWORD)
    graph.clear_database()

    # Create sample containers (nodes)
    # Container A: Misconfigured, high risk, no public connection
    container_a = graph.create_container("containerA", risk_score=80, has_public_connection=False, 
                                           services=["web", "database"])
    # Container B: Internal container with moderate risk
    container_b = graph.create_container("containerB", risk_score=65, has_public_connection=False, 
                                           services=["api", "cache"])
    # Container C: Exposed container with lower base risk but flagged as public
    container_c = graph.create_container("containerC", risk_score=75, has_public_connection=True, 
                                           services=["frontend"])
    
    # Create communication relationships
    # Scenario 1: Container A communicates with Container C (public exposure path)
    graph.create_communication("containerA", "containerC")
    # Scenario 2: Container B communicates with Container C (unexpected lateral movement)
    graph.create_communication("containerB", "containerC")
    
    # Run the dynamic detection query before any updates
    print("Initial detection results:")
    results = graph.dynamic_detection_query(time_window_minutes=10, risk_threshold=70, anomaly_threshold=70)
    for record in results:
        print(record)
    
    # Simulate a dynamic update: Container B's risk increases due to new vulnerability detection
    print("\nUpdating containerB risk score dynamically...")
    graph.update_risk_score("containerB", new_risk_score=85)
    
    # Run the detection query again after the update
    print("\nDetection results after dynamic update:")
    results = graph.dynamic_detection_query(time_window_minutes=10, risk_threshold=70, anomaly_threshold=70)
    for record in results:
        print(json.dumps(record, indent=4)) 
    
    # Close the connection
    graph.close()