Graph_Cursor.py

import socket
import struct
from utils import pack_string
from GraphQLSender import run_cypher_query

user_agent = "FPOSCustomClient/1.0"
PREAMBLE = b'\x60\x60\xb0\x17'
slot1 = 0x00000003  # Explicitly Bolt v3.0
slot2 = 0x00050003  # Bolt v3.0 through v3.5 (Range format)
slot3 = 0x00000002  # Explicitly Bolt v4.0 (Fallback)

class GraphCursor():
    """
    A cursor item used to query graph databases using the bolt protocol.
    """
    def __init__(self, uri: str="", port: int=7687, user: str="", password: str="", version: int=0x00000002):
        """
        :param uri: The URI of the Graph database, exclude protocol.
        :param port: The port via which to access the database (standard is 7687)
        :param user: The username used to access the database.
        :param password: The password for the provided user.
        :param version: An integer representing the protocol version of the Graph database to be used.
        """
        socket.getaddrinfo(uri, port, socket.AF_INET, socket.SOCK_STREAM)
        self.s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.s.connect((uri, port))
        ready = self.__initialize_connection(version, user, password)
        if not ready:
            raise Exception("Connection failed")

    def __initialize_connection(self, version, user, password) -> bool:
        """
        An internal method used to establish the open connection to the database.
        :param version: The user version provided in the init.
        :param user: The username used to access the database.
        :param password: The password for the provided user.
        :return: True if the connection was successful, False otherwise.
        """
        version_proposals = struct.pack('>IIII', version, slot1, slot2, slot3)
        self.s.sendall(PREAMBLE + version_proposals)
        # time.sleep(.5)
        selected_version = self.s.recv(4)

        auth_map = (
                bytes([0xA3]) +
                pack_string("scheme") +
                pack_string("basic") +
                pack_string("principal") +
                pack_string(user) +
                pack_string("credentials") +
                pack_string("password")
        )

        init_body = bytes([0xB2, 0x01]) + pack_string(user_agent) + auth_map

        chunk_header = struct.pack('>H', len(init_body))
        chunk_end = b'\x00\x00'

        full_auth_packet = chunk_header + init_body + chunk_end
        self.s.sendall(full_auth_packet)
        resp_header = self.s.recv(2)
        if resp_header:
            chunk_size = struct.unpack('>H', resp_header)[0]
            resp_body = self.s.recv(chunk_size)
            end_marker = self.s.recv(2)

            if resp_body[1] == 0x70:
               return True
            else:
                print("AUTH FAILED")
                return False
        else:
            print("No response...")
            return False

    def run(self, query) -> list[dict]:
        """
        Runs a basic query against the graph database.
        :param query: The query string to execute on the database.
        """
        return run_cypher_query(self.s, query)