graph-test2.py

import time
from datetime import timedelta
from neo4j import GraphDatabase

# Neo4j configuration
URI = "bolt://localhost:7687"
USERNAME = "neo4j"
PASSWORD = "password"  # Change this to your actual Neo4j password

class CloudSecurityGraphSimulator:
    def __init__(self, uri, user, password):
        self.driver = GraphDatabase.driver(uri, auth=(user, password))

    def close(self):
        self.driver.close()

    def clear_database(self):
        with self.driver.session() as session:
            session.run("MATCH (n) DETACH DELETE n")
            print("Database cleared.")

    def create_container(self, container_id, risk_score, has_public_connection, services):
        """
        Create a container node with given properties.
        """
        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)
            print(f"Created container {container_id} with riskScore {risk_score}")
            return result.single()[0]

    def create_edge(self, from_id, to_id, label="COMMUNICATES_WITH"):
        """
        Create an edge between two container nodes.
        """
        timestamp = int(time.time() * 1000)
        query = """
            MATCH (src:Container {id: $from_id}), (dst:Container {id: $to_id})
            CREATE (src)-[r:%s {createdAt: $timestamp}]->(dst)
            RETURN r
        """ % label
        with self.driver.session() as session:
            result = session.run(query, from_id=from_id, to_id=to_id, timestamp=timestamp)
            print(f"Created edge {label} from {from_id} to {to_id}")
            return result.single()[0]

    def update_risk(self, container_id, new_risk):
        """
        Update the riskScore and lastEventTimestamp for a container node.
        """
        timestamp = int(time.time() * 1000)
        query = """
            MATCH (c:Container {id: $id})
            SET c.riskScore = $newRisk, c.lastEventTimestamp = $timestamp,
                c.label = c.id + ' (Risk: ' + toString($newRisk) + ')'
            RETURN c
        """
        with self.driver.session() as session:
            result = session.run(query, id=container_id, newRisk=new_risk, timestamp=timestamp)
            print(f"Updated {container_id} risk to {new_risk}")
            return result.single()[0]

    def remove_edge(self, edge_id):
        """
        Remove an edge by its ID.
        """
        with self.driver.session() as session:
            session.run("MATCH ()-[r]->() WHERE id(r) = $edge_id DELETE r", edge_id=edge_id)
            print(f"Removed edge with id {edge_id}")

    def detection_query(self, time_window_minutes=10, risk_threshold=70, anomaly_threshold=70):
        """
        Run a Cypher query to detect communication paths where:
          - The source container has a riskScore above risk_threshold.
          - The destination container is publicly accessible and its riskScore is above anomaly_threshold.
          - Both nodes had events within the defined time window.
        A composite risk score is calculated dynamically.
        """
        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 src.hasPublicConnection = true
              AND dst.riskScore >= aThreshold
            RETURN src.id AS Source, 
                   dst.id AS Destination,
                   ((src.riskScore * 0.7) + (dst.riskScore * 0.3)) AS compositeRisk
            ORDER BY compositeRisk 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)
            detections = [record.data() for record in result]
            print("Detection query results:")
            for d in detections:
                print(d)
            return detections
    
    def make_public(self, container_id):
        """
        Update a container node to have a public connection.
        """
        timestamp = int(time.time() * 1000)
        query = """
            MATCH (c:Container {id: $id})
            SET c.hasPublicConnection = true, c.lastEventTimestamp = $timestamp
            RETURN c
        """
        with self.driver.session() as session:
            result = session.run(query, id=container_id, timestamp=timestamp)
            print(f"Updated {container_id} to have a public connection")
            return result.single()[0]   

if __name__ == "__main__":
    simulator = CloudSecurityGraphSimulator(URI, USERNAME, PASSWORD)
    simulator.clear_database()

    # --- Create Nodes for Three Scenarios ---
    # Scenario 1: Misconfigured Public-Facing Service
    # Container A: high risk, misconfigured, public
    containerA = simulator.create_container("containerA", risk_score=80, has_public_connection=True, services=["web"])
    # Container X: internal sensitive API
    containerX = simulator.create_container("containerX", risk_score=50, has_public_connection=False, services=["sensitive API"])

    # Scenario 2: Unexpected Lateral Movement
    # Container B: internal container, low risk initially
    containerB = simulator.create_container("containerB", risk_score=60, has_public_connection=False, services=["internal service"])
    # Container Y: normally isolated internal service
    containerY = simulator.create_container("containerY", risk_score=55, has_public_connection=False, services=["isolated service"])

    # Scenario 3: Cascade from a Compromised Container
    # Container C: compromised container with high risk
    containerC = simulator.create_container("containerC", risk_score=85, has_public_connection=False, services=["compromised container"])
    # Downstream containers
    containerD = simulator.create_container("containerD", risk_score=50, has_public_connection=False, services=["downstream service"])
    containerE = simulator.create_container("containerE", risk_score=45, has_public_connection=False, services=["downstream service"])

    # --- Create Initial Edges ---
    # Scenario 1: Container A -> Container X
    edge_A_X = simulator.create_edge("containerA", "containerX")
    # Scenario 3: Container C -> Container D and Container C -> Container E
    edge_C_D = simulator.create_edge("containerC", "containerD")
    edge_C_E = simulator.create_edge("containerC", "containerE")
    # (Scenario 2: initially, no edge between Container B and Container Y)

    # --- Simulate Dynamic Changes ---
    print("\n--- Simulation Start ---\n")

    # After 3 seconds, update Container B's risk and add edge to Container Y (unexpected lateral movement)
    time.sleep(3)
    simulator.update_risk("containerB", new_risk=80)
    edge_B_Y = simulator.create_edge("containerB", "containerY")

    # After 6 seconds, update Container A's risk (misconfiguration worsens)
    time.sleep(3)
    simulator.update_risk("containerA", new_risk=90)

    # After 9 seconds, simulate a remediation on Scenario 3 by removing the edge between Container C and Container D
    time.sleep(3)
    # First, we need to retrieve the edge ID for edge_C_D.
    # For simplicity, we'll re-run a query to get it (in production you would track edge IDs more robustly)
    query_edge = "MATCH ()-[r:COMMUNICATES_WITH]->() WHERE r.createdAt IS NOT NULL RETURN id(r) AS edgeId LIMIT 1"
    with simulator.driver.session() as session:
        record = session.run(query_edge).single()
        if record:
            edge_id = record["edgeId"]
            simulator.remove_edge(edge_id)

    # Run the detection query to see current risky paths (for Scenario 1)
    simulator.detection_query(time_window_minutes=10, risk_threshold=70, anomaly_threshold=70)


    # After 12 seconds, update Container B's risk back to normal and add edge to Container X (expected communication)
    time.sleep(3)
    simulator.update_risk("containerB", new_risk=60)
    simulator.create_edge("containerB", "containerX")
    simulator.detection_query(time_window_minutes=10, risk_threshold=70, anomaly_threshold=70)

    # After 12 seconds, make container A talk to container B
    time.sleep(3)
    simulator.create_edge("containerA", "containerB")
    simulator.detection_query(time_window_minutes=10, risk_threshold=70, anomaly_threshold=70)

    simulator.close()
    print("\n--- Simulation End ---\n")
    # simulator.close()



# Query to find the path that has the highest composite risk score
def highest_risk_path():
    cypher = """
        WITH $time_window_ms AS window, $current_ts AS now, $risk_threshold AS rThreshold, $anomaly_threshold AS aThreshold
        MATCH path = (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
        WITH path, src, dst, nodes(path) AS containers
        RETURN src.id AS Source, 
            dst.id AS Destination,
            reduce(s = 0, n IN containers | s + n.riskScore) AS compositeRisk
        ORDER BY compositeRisk DESC
        LIMIT 1
    """