Upload 12 files

codeql_kernel/__init__.py

@@ -0,0 +1,3 @@
+from .kernel import CodeQLKernel
+from .codeql import CLIClient, QueryClient
+from ._version import __version__

codeql_kernel/__main__.py

@@ -0,0 +1,4 @@
+from ipykernel.kernelapp import IPKernelApp
+from . import CodeQLKernel
+
+IPKernelApp.launch_instance(kernel_class=CodeQLKernel)

codeql_kernel/_version.py

@@ -0,0 +1 @@
+__version__ = '0.0.1'

codeql_kernel/codeql.py

@@ -0,0 +1,220 @@
+import logging
+import os
+import tempfile
+import time
+from subprocess import PIPE, Popen
+from typing import Optional, Tuple
+
+from .jsonrpc import RPC as JSONRPC
+from .rawrpc import RPC as RawRPC
+
+
+class CLIClient:
+    """
+    Represents a JSONRPC client to connect to CodeQL CLI Server
+    """
+
+    def __init__(self):
+        self.cache = {"ram": []}
+        self.conn = RawRPC(
+            [
+                "codeql",
+                "execute",
+                "cli-server",
+                "--logdir",
+                "/tmp/codeql_kernel_cliserver",
+            ]
+        )
+
+    def stop(self):
+        self.conn.stop()
+
+    def resolve_ram(self) -> Tuple[Optional[str], Optional[list]]:
+        if self.cache.get("ram"):
+            return (None, self.cache.get("ram"))
+        else:
+            cmd = ["resolve", "ram", "--format=json"]
+            (err, result) = self.conn.request(cmd)
+            if err:
+                return (err, None)
+            self.cache["ram"] = [x for x in result if x.startswith("-J")]
+            return (None, self.cache.get("ram"))
+
+    def resolve_metadata(self, query) -> Tuple[Optional[str], dict]:
+        cmd = ["resolve", "metadata", "--format=json", query]
+        return self.conn.request(cmd)
+
+    def resolve_database(self, db_path) -> Tuple[Optional[str], dict]:
+        cmd = ["resolve", "database", "--format=json", db_path]
+        return self.conn.request(cmd)
+
+    def resolve_library_path(self, query) -> Tuple[Optional[str], Optional[dict]]:
+        cmd = ["resolve", "library-path", "--format=json", "--query", query]
+        return self.conn.request(cmd)
+
+    def bqrs_info(self, bqrs_path) -> Tuple[Optional[str], dict]:
+        cmd = ["bqrs", "info", "--format=json", bqrs_path]
+        return self.conn.request(cmd)
+
+    def bqrs_decode(self, bqrs_path) -> Tuple[Optional[str], Optional[str]]:
+        (err, ram_opts) = self.resolve_ram()
+        if err or not ram_opts:
+            return (f"Error resolving ram options {err}", None)
+        results_path = tempfile.NamedTemporaryFile(delete=False)
+        cmd = [
+            "bqrs",
+            "decode",
+            "--format=csv",
+            f"-o={results_path.name}",
+            "--entities=string,url",
+            bqrs_path,
+        ]
+        cmd.extend(ram_opts)
+        (err, _) = self.conn.request(cmd)
+        if err:
+            return (f"Error decoding bqrs file {err}", None)
+        if os.path.exists(results_path.name):
+            with open(results_path.name, "r") as f:
+                data = f.read()
+                # return json.loads(data)
+                return (None, data)
+        else:
+            return ("Error decoding results", None)
+
+
+class QueryClient:
+    """
+    Represents a JSONRPC client to connect to CodeQL Query Server
+    """
+
+    def __init__(self, on_progress=None, on_result=None):
+        self._cli_client: CLIClient = CLIClient()
+        cmd = ["codeql", "execute", "query-server2", "--threads=0", "--evaluator-log-level", "5"]
+        # debug
+        # cmd.extend(["--debug", "--tuple-counting", "-v", "--log-to-stderr"])
+        # --save-cache --max-disk-cache XX
+        (err, ram_opts) = self._cli_client.resolve_ram()
+        if err or not ram_opts:
+            return (f"Error resolving ram options {err}", None)
+        cmd.extend(ram_opts)
+        self._proc = Popen(cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+        handlers = {}
+        if on_progress:
+            handlers["ql/progressUpdated"] = on_progress
+        self._conn = JSONRPC(
+            handlers=handlers, stdout=self._proc.stdin, stdin=self._proc.stdout
+        )
+        self._progress_id = 0
+        self._evaluate_id = 0
+        self._db_metadata = {}
+        # TODO: wait for query server to be ready
+        time.sleep(2)
+
+    def stop(self):
+        if self._proc.stdin:
+            self._proc.stdin.close()
+        if self._proc.stdout:
+            self._proc.stdout.close()
+        self._proc.terminate()
+        self._proc.wait()
+        if self._cli_client:
+            self._cli_client.stop()
+
+    def next_progress_id(self) -> int:
+        self._progress_id += 1
+        return self._progress_id
+
+    def next_evaluate_id(self) -> int:
+        self._evaluate_id += 1
+        return self._evaluate_id
+
+    def register_database(self, db_path) -> Optional[str]:
+        """
+        Register a database with the query server
+        """
+        if not db_path.endswith("/"):
+            db_path = db_path + "/"
+        if not os.path.isdir(db_path):
+            return f"Database path {db_path} is not a directory"
+
+        (err, db_metadata) = self._cli_client.resolve_database(db_path)
+        if err:
+            return "Failed to resolve database metadata"
+
+        # TODO: implement on-the-fly query patching
+
+        params = {
+            "body": {
+                "databases": [db_path],
+                "progressId": self.next_progress_id(),
+            }
+        }
+        (err, _) = self._conn.request("evaluation/registerDatabases", args=params)
+
+        if err:
+            return err
+
+        self._db_metadata = db_metadata
+        self._db_metadata["path"] = db_path
+
+        return None
+
+    def run_query(
+        self, query_path, quick_eval={}
+    ) -> Tuple[Optional[str], Optional[str]]:
+        logging.info(f"Running query {query_path}")
+        bqrs_path = tempfile.NamedTemporaryFile(suffix=".bqrs").name
+        target = {"query": {"xx": ""}}
+        if bool(quick_eval):
+            target = {
+                "quickEval": {
+                    "quickEvalPos": {
+                        "fileName": query_path,
+                        "line": quick_eval.get("startLine"),
+                        "column": quick_eval.get("startColumn"),
+                        "endLine": quick_eval.get("endLine"),
+                        "endColumn": quick_eval.get("endColumn"),
+                    }
+                }
+            }
+
+        run_queries_params = {
+            "body": {
+                "db": self._db_metadata["path"],
+                # TODO: get additional packs from ENV, command, config, etc.
+                "additionalPacks": ["/Users/pwntester/src/github.com/github/codeql"],
+                "externalInputs": [],
+                "singletonExternalInputs": [],  # opts.templateValues or {},
+                "outputPath": bqrs_path,
+                "queryPath": query_path,
+                "target": target,
+            },
+            "progressId": self.next_progress_id(),
+        }
+
+        (err, resp) = self._conn.request(
+            "evaluation/runQuery", args=run_queries_params
+        )
+
+        if resp and resp["resultType"] != 0:
+            return (resp["message"], None)
+
+        if err:
+            return (str(err), None)
+
+        if os.path.exists(bqrs_path):
+            (err, bqrs_info) = self._cli_client.bqrs_info(bqrs_path)
+            if err:
+                return (err, "")
+            if not bqrs_info or not bqrs_info["result-sets"]:
+                return ("Failed to get bqrs info", "")
+            count = bqrs_info["result-sets"][0]["rows"]
+            for result_set in bqrs_info["result-sets"]:
+                if result_set["name"] == "#select":
+                    count = result_set["rows"]
+            if count > 0:
+                return self._cli_client.bqrs_decode(bqrs_path)
+            else:
+                return (None, "No results")
+        else:
+            return (f"Failed to find results file at {bqrs_path}", "")

codeql_kernel/jsonrpc.py

@@ -0,0 +1,664 @@
+# coding: utf-8
+
+"""
+Minimal python RPC implementation in a single file based on the JSON-RPC 2.0 specs from
+http://www.jsonrpc.org/specification.
+"""
+
+
+__author__ = "Marcel Rieger"
+__email__ = "[email protected]"
+__copyright__ = "Copyright 2016-2021, Marcel Rieger"
+__credits__ = ["Marcel Rieger"]
+__contact__ = "https://github.com/riga/jsonrpyc"
+__license__ = "BSD-3-Clause"
+__status__ = "Development"
+__version__ = "1.1.1"
+__all__ = ["RPC"]
+
+
+import io
+import json
+import logging
+import sys
+import threading
+import time
+from queue import Queue
+
+
+class Spec(object):
+    """
+    This class wraps methods that create JSON-RPC 2.0 compatible string representations of
+    request, response and error objects. All methods are class members, so you might never want to
+    create an instance of this class, but rather use the methods directly:
+
+    .. code-block:: python
+
+        Spec.request("my_method", 18)  # the id is optional
+        # => '{"jsonrpc":"2.0","method":"my_method","id": 18}'
+
+        Spec.response(18, "some_result")
+        # => '{"jsonrpc":"2.0","id":18,"result":"some_result"}'
+
+        Spec.error(18, -32603)
+        # => '{"jsonrpc":"2.0","id":18,"error":{"code":-32603,"message":"Internal error"}}'
+    """
+
+    @classmethod
+    def check_id(cls, id, allow_empty=False):
+        """
+        Value check for *id* entries. When *allow_empty* is *True*, *id* is allowed to be *None*.
+        Raises a *TypeError* when *id* is neither an integer nor a string.
+        """
+        if (id is not None or not allow_empty) and not isinstance(id, (int, str)):
+            raise TypeError(
+                "id must be an integer or string, got {} ({})".format(id, type(id))
+            )
+
+    @classmethod
+    def check_method(cls, method):
+        """
+        Value check for *method* entries. Raises a *TypeError* when *method* is not a string.
+        """
+        if not isinstance(method, str):
+            raise TypeError(
+                "method must be a string, got {} ({})".format(method, type(method))
+            )
+
+    @classmethod
+    def check_code(cls, code):
+        """
+        Value check for *code* entries. Raises a *TypeError* when *code* is not an integer, or a
+        *KeyError* when there is no :py:class:`RPCError` subclass registered for that *code*.
+        """
+        if not isinstance(code, int):
+            raise TypeError("code must be an integer, got {} ({})".format(id, type(id)))
+
+        if not get_error(code):
+            raise ValueError("unknown code, got {} ({})".format(code, type(code)))
+
+    @classmethod
+    def request(cls, method, id=None, params=None):
+        """
+        Creates the string representation of a request that calls *method* with optional *params*
+        which are encoded by ``json.dumps``. When *id* is *None*, the request is considered a
+        notification.
+        """
+        try:
+            cls.check_method(method)
+            cls.check_id(id, allow_empty=True)
+        except Exception as e:
+            raise RPCInvalidRequest(str(e))
+
+        # start building the request string
+        req = '{{"jsonrpc":"2.0","method":"{}"'.format(method)
+
+        # add the id when given
+        if id is not None:
+            # encode string ids
+            if isinstance(id, str):
+                id = json.dumps(id)
+            req += ',"id":{}'.format(id)
+
+        # add parameters when given
+        if params is not None:
+            try:
+                req += ',"params":{}'.format(json.dumps(params))
+            except Exception as e:
+                raise RPCParseError(str(e))
+
+        # end the request string
+        req += "}"
+
+        return req
+
+    @classmethod
+    def response(cls, id, result):
+        """
+        Creates the string representation of a respone that was triggered by a request with *id*.
+        A *result* is required, even if it is *None*.
+        """
+        try:
+            cls.check_id(id)
+        except Exception as e:
+            raise RPCInvalidRequest(str(e))
+
+        # encode string ids
+        if isinstance(id, str):
+            id = json.dumps(id)
+
+        # build the response string
+        try:
+            res = '{{"jsonrpc":"2.0","id":{},"result":{}}}'.format(
+                id, json.dumps(result)
+            )
+        except Exception as e:
+            raise RPCParseError(str(e))
+
+        return res
+
+    @classmethod
+    def error(cls, id, code, data=None):
+        """
+        Creates the string representation of an error that occured while processing a request with
+        *id*. *code* must lead to a registered :py:class:`RPCError`. *data* might contain
+        additional, detailed error information and is encoded by ``json.dumps`` when set.
+        """
+        try:
+            cls.check_id(id)
+            cls.check_code(code)
+        except Exception as e:
+            raise RPCInvalidRequest(str(e))
+
+        # build the inner error data
+        message = get_error(code).title
+        err_data = '{{"code":{},"message":"{}"'.format(code, message)
+
+        # insert data when given
+        if data is not None:
+            try:
+                err_data += ',"data":{}}}'.format(json.dumps(data))
+            except Exception as e:
+                raise RPCParseError(str(e))
+        else:
+            err_data += "}"
+
+        # encode string ids
+        if isinstance(id, str):
+            id = json.dumps(id)
+
+        # start building the error string
+        err = '{{"jsonrpc":"2.0","id":{},"error":{}}}'.format(id, err_data)
+
+        return err
+
+
+class RPC(object):
+    """
+    The main class of *jsonrpyc*. Instances of this class wrap an input stream *stdin* and an output
+    stream *stdout* in order to communicate with other services. A service is not even forced to be
+    written in Python as long as it strictly implements the JSON-RPC 2.0 specification. RPC
+    instances may wrap a *target* object. By means of a :py:class:`Watchdog` instance, incoming
+    requests are routed to methods of this object whose result might be sent back as a response.
+    The watchdog instance is created but not started yet, when *watch* is not *True*.
+    Example implementation:
+
+    *server.py*
+
+    .. code-block:: python
+
+        import jsonrpyc
+
+        class MyTarget(object):
+
+            def greet(self, name):
+                return f"Hi, {name}!"
+
+       jsonrpc.RPC(MyTarget())
+
+    *client.py*
+
+    .. code-block:: python
+
+        import jsonrpyc
+        from subprocess import Popen, PIPE
+
+        p = Popen(["python", "server.py"], stdin=PIPE, stdout=PIPE)
+        rpc = jsonrpyc.RPC(stdout=p.stdin, stdin=p.stdout)
+
+        # non-blocking remote procedure call with callback and js-like signature
+        def cb(err, res=None):
+            if err:
+                throw err
+            print(f"callback got: {res}")
+
+        rpc("greet", args=("John",), callback=cb)
+
+        # cb is called asynchronously which prints
+        # => "callback got: Hi, John!"
+
+        # blocking remote procedure call with 0.1s polling
+        print(rpc("greet", args=("John",), block=0.1))
+        # => "Hi, John!"
+
+        # shutdown the process
+        p.stdin.close()
+        p.stdout.close()
+        p.terminate()
+        p.wait()
+
+    .. py:attribute:: target
+
+       The wrapped target object. Might be *None* when no object is wrapped, e.g. for the *client*
+       RPC instance.
+
+    .. py:attribute:: stdin
+
+       The input stream, re-opened with ``"rb"``.
+
+    .. py:attribute:: stdout
+
+       The output stream, re-opened with ``"wb"``.
+
+    .. py:attribute:: watch
+
+       The :py:class:`Watchdog` instance that optionally watches *stdin* and dispatches incoming
+       requests.
+    """
+
+    EMPTY_RESULT = object()
+
+    def __init__(self, handlers=None, stdin=None, stdout=None, watch=True, **kwargs):
+        super(RPC, self).__init__()
+
+        self.handlers = handlers
+
+        # open streams
+        stdin = sys.stdin if stdin is None else stdin
+        stdout = sys.stdout if stdout is None else stdout
+        self.stdin = io.open(stdin.fileno(), "rb")
+        self.stdout = io.open(stdout.fileno(), "wb")
+
+        # other attributes
+        self._i = 0
+        self._callbacks = {}
+        self._results = {}
+
+        # create and optional start the watchdog
+        kwargs["start"] = watch
+        # kwargs.setdefault("daemon", handlers is None)
+        self.watchdog = Watchdog(self, **kwargs)
+
+    def __del__(self):
+        watchdog = getattr(self, "watchdog", None)
+        if watchdog:
+            watchdog.stop()
+
+    def __call__(self, *args, **kwargs):
+        """
+        Shorthand for :py:meth:`call`.
+        """
+        return self.call(*args, **kwargs)
+
+    def request(
+        self,
+        cmd,
+        args=(),
+    ):
+        """
+        Sends a request to the remote service and waits for the response
+        """
+        q = Queue()
+
+        def cb(err, resp):
+            q.put((err, resp))
+
+        self.call(cmd, args=args, callback=cb)
+        return q.get()
+
+    def call(self, method, args=(), kwargs=None, callback=None, block=0):
+        """
+        Performs an actual remote procedure call by writing a request representation (a string) to
+        the output stream. The remote RPC instance uses *method* to route to the actual method to
+        call with *args* and *kwargs*. When *callback* is set, it will be called with the result of
+        the remote call. When *block* is larger than *0*, the calling thread is blocked until the
+        result is received. In this case, *block* will be the poll interval, emulating synchronuous
+        return value behavior. When both *callback* is *None* and *block* is *0* or smaller, the
+        request is considered a notification and the remote RPC instance will not send a response.
+        """
+        # default kwargs
+        if kwargs is None:
+            kwargs = {}
+
+        # check if the call is a notification
+        is_notification = callback is None and block <= 0
+
+        # create a new id for requests expecting a response
+        id = None
+        if not is_notification:
+            self._i += 1
+            id = self._i
+
+        # register the callback
+        if callback is not None:
+            self._callbacks[id] = callback
+
+        # store an empty result for the meantime
+        if block > 0:
+            self._results[id] = self.EMPTY_RESULT
+
+        # create the request
+        req = Spec.request(method, id=id, params=args)
+        self._write(req)
+
+        # blocking return value behavior
+        if block > 0:
+            while True:
+                if self._results[id] != self.EMPTY_RESULT:
+                    result = self._results[id]
+                    del self._results[id]
+                    if isinstance(result, Exception):
+                        raise result
+                    else:
+                        return result
+                time.sleep(block)
+
+    def _handle(self, line):
+        """
+        Handles an incoming *line* and dispatches the parsed object to the request, response, or
+        error handlers.
+        """
+        obj = json.loads(line)
+
+        # dispatch to the correct handler
+        if "method" in obj:
+            # request
+            self._handle_request(obj)
+        elif "error" not in obj:
+            # response
+            self._handle_response(obj)
+        else:
+            # error
+            self._handle_error(obj)
+
+    def _handle_request(self, req):
+        """
+        Handles an incoming request *req*. When it containes an id, a response or error is sent
+        back.
+        """
+        logging.debug(f"Handling request to {req['method']}")
+        try:
+            method = self._route(req["method"])
+            result = method(req["params"])
+            if "id" in req:
+                res = Spec.response(req["id"], result)
+                self._write(res)
+        except Exception as e:
+            if "id" in req:
+                if isinstance(e, RPCError):
+                    err = Spec.error(req["id"], e.code, e.data)
+                else:
+                    err = Spec.error(req["id"], -32603, str(e))
+                self._write(err)
+
+    def _handle_response(self, res):
+        """
+        Handles an incoming successful response *res*. Blocking calls are resolved and registered
+        callbacks are invoked with the first error argument being set to *None*.
+        """
+        logging.debug(f"Handling response for {res['id']}: {res}")
+        # set the result
+        if res["id"] in self._results:
+            self._results[res["id"]] = res["result"]
+
+        # lookup and invoke the callback
+        if res["id"] in self._callbacks:
+            callback = self._callbacks[res["id"]]
+            del self._callbacks[res["id"]]
+            callback(None, res["result"])
+
+    def _handle_error(self, res):
+        """
+        Handles an incoming failed response *res*. Blocking calls throw an exception and
+        registered callbacks are invoked with an exception and the second result argument set to
+        *None*.
+        """
+        logging.debug(f"Handling error {res}")
+        # extract the error and create an actual error instance to raise
+        err = res["error"]
+        error = get_error(err["code"])(err.get("data", err["message"]))
+
+        # set the error
+        if res["id"] in self._results:
+            self._results[res["id"]] = error
+
+        # lookup and invoke the callback
+        if res["id"] in self._callbacks:
+            callback = self._callbacks[res["id"]]
+            del self._callbacks[res["id"]]
+            callback(error, None)
+
+    def _route(self, method):
+        if method in self.handlers.keys():
+            return self.handlers[method]
+        else:
+            raise RPCMethodNotFound(data=method)
+
+    def _write(self, s):
+        """
+        Writes a string *s* to the output stream.
+        """
+        msg = f"Content-Length: {len(s)}\n\n{s}"
+        logging.debug("SENT: \n" + str(msg) + "\n\n")
+        self.stdout.write(bytearray(msg, "utf-8"))
+        self.stdout.flush()
+
+
+class Watchdog(threading.Thread):
+    """
+    This class represents a thread that watches the input stream of an :py:class:`RPC` instance for
+    incoming content and dispatches requests to it.
+
+    .. py:attribute:: rpc
+
+       The :py:class:`RPC` instance.
+
+    .. py:attribute:: name
+
+       The thread's name.
+
+    .. py:attribute:: interval
+
+       The polling interval of the run loop.
+
+    .. py:attribute:: daemon
+
+       The thread's daemon flag.
+    """
+
+    def __init__(self, rpc, name="watchdog", interval=0.1, daemon=True, start=True):
+        super(Watchdog, self).__init__()
+
+        # store attributes
+        self.rpc = rpc
+        self.name = name
+        self.interval = interval
+        self.daemon = daemon
+
+        # register a stop event
+        self._stop = threading.Event()
+
+        if start:
+            self.start()
+
+    def start(self):
+        """
+        Starts with thread's activity.
+        """
+        super(Watchdog, self).start()
+
+    def stop(self):
+        """
+        Stops with thread's activity.
+        """
+        self._stop.set()
+
+    def run(self):
+        # reset the stop event
+        self._stop.clear()
+
+        # stop here when stdin is not set or closed
+        if not self.rpc.stdin or self.rpc.stdin.closed:
+            return
+
+        # read new incoming lines
+        last_pos = 0
+        while not self._stop.is_set():
+            lines = None
+
+            # stop when stdin is closed
+            if self.rpc.stdin.closed:
+                break
+
+            # read from stdin depending on whether it is a tty or not
+            if self.rpc.stdin.isatty():
+                cur_pos = self.rpc.stdin.tell()
+                if cur_pos != last_pos:
+                    self.rpc.stdin.seek(last_pos)
+                    lines = self.rpc.stdin.readlines()
+                    last_pos = self.rpc.stdin.tell()
+                    self.rpc.stdin.seek(cur_pos)
+            else:
+                try:
+                    header = self.rpc.stdin.readline()
+                    header = header.decode("utf-8").strip()
+                    if header.startswith("Content-Length:"):
+                        length = 2 + int(header[len("Content-Length:") :])
+                        lines = [self.rpc.stdin.read(length)]
+                except Exception:
+                    # prevent residual race conditions occurring when stdin is closed externally
+                    pass
+
+            # handle new lines if any
+            if lines:
+                for line in lines:
+                    line = line.decode("utf-8").strip()
+                    if line:
+                        self.rpc._handle(line)
+            else:
+                self._stop.wait(self.interval)
+
+
+class RPCError(Exception):
+
+    """
+    Base class for RPC errors.
+
+    .. py:attribute:: message
+
+       The message of this error, i.e., ``"<title> (<code>)[, data: <data>]"``.
+
+    .. py:attribute:: data
+
+       Additional data of this error. Setting the data attribute will also change the message
+       attribute.
+    """
+
+    def __init__(self, data=None):
+        # build the error message
+        message = "{} ({})".format(self.title, self.code)
+        if data is not None:
+            message += ", data: {}".format(data)
+        self.message = message
+
+        super(RPCError, self).__init__(message)
+
+        self.data = data
+
+    def __str__(self):
+        return self.message
+
+
+error_map_distinct = {}
+error_map_range = {}
+
+
+def is_range(code):
+    return (
+        isinstance(code, tuple)
+        and len(code) == 2
+        and all(isinstance(i, int) for i in code)
+        and code[0] < code[1]
+    )
+
+
+def register_error(cls):
+    """
+    Decorator that registers a new RPC error derived from :py:class:`RPCError`. The purpose of
+    error registration is to have a mapping of error codes/code ranges to error classes for faster
+    lookups during error creation.
+
+    .. code-block:: python
+
+       @register_error
+       class MyCustomRPCError(RPCError):
+           code = ...
+           title = "My custom error"
+    """
+    # it would be much cleaner to add a meta class to RPCError as a registry for codes
+    # but in CPython 2 exceptions aren't types, so simply provide a registry mechanism here
+    if not issubclass(cls, RPCError):
+        raise TypeError("'{}' is not a subclass of RPCError".format(cls))
+
+    code = cls.code
+
+    if isinstance(code, int):
+        error_map = error_map_distinct
+    elif is_range(code):
+        error_map = error_map_range
+    else:
+        raise TypeError("invalid RPC error code {}".format(code))
+
+    if code in error_map:
+        raise AttributeError("duplicate RPC error code {}".format(code))
+
+    error_map[code] = cls
+
+    return cls
+
+
+def get_error(code):
+    """
+    Returns the RPC error class that was previously registered to *code*. *None* is returned when no
+    class could be found.
+    """
+    if code in error_map_distinct:
+        return error_map_distinct[code]
+
+    for (lower, upper), cls in error_map_range.items():
+        if lower <= code <= upper:
+            return cls
+
+    return None
+
+
+@register_error
+class RPCParseError(RPCError):
+
+    code = -32700
+    title = "Parse error"
+
+
+@register_error
+class RPCInvalidRequest(RPCError):
+
+    code = -32600
+    title = "Invalid Request"
+
+
+@register_error
+class RPCMethodNotFound(RPCError):
+
+    code = -32601
+    title = "Method not found"
+
+
+@register_error
+class RPCInvalidParams(RPCError):
+
+    code = -32602
+    title = "Invalid params"
+
+
+@register_error
+class RPCInternalError(RPCError):
+
+    code = -32603
+    title = "Internal error"
+
+
+@register_error
+class RPCServerError(RPCError):
+
+    code = (-32099, -32000)
+    title = "Server error"

codeql_kernel/kernel.json

@@ -0,0 +1,13 @@
+{
+  "argv": [
+    "python",
+    "-m",
+    "codeql_kernel",
+    "-f",
+    "{connection_file}"
+  ],
+  "display_name": "CodeQL",
+  "mimetype": "text/x-codeql",
+  "language": "codeql",
+  "name": "codeql"
+}

codeql_kernel/kernel.py

@@ -0,0 +1,189 @@
+import os
+from io import StringIO
+from tempfile import mkdtemp, mkstemp
+
+import pandas as pd
+import tree_sitter
+from IPython.display import HTML
+from metakernel import MetaKernel
+import subprocess
+
+from .codeql import QueryClient
+
+__version__ = "0.0.1"
+
+
+class CodeQLKernel(MetaKernel):
+    implementation = "CodeQL Kernel"
+    implementation_version = "1.0"
+    language = "ql"
+    language_version = "0.1"
+    banner = "CodeQL Kernel - Experimental"
+    language_info = {
+        "mimetype": "text/x-codeql",
+        "name": "codeql",
+        "file_extension": ".ql",
+        "help_links": MetaKernel.help_links,
+    }
+
+    def __init__(self, **kwargs):
+        # get absolute path of running script
+        here = os.path.dirname(os.path.abspath(__file__))
+        self.QL_LANGUAGE = tree_sitter.Language(
+            os.path.join(here, "tree-sitter-ql.so"), "ql"
+        )
+        self._select_query = self.QL_LANGUAGE.query(
+            "(moduleMember (select)) @select_statement"
+        )
+        self._predicate_query = self.QL_LANGUAGE.query(
+            """(moduleMember
+                (annotation name: (annotName) @aname (#eq? @aname "query")).
+                (classlessPredicate name: (predicateName) @pname)
+               ) @annotated_query """
+        )
+        self._parser = tree_sitter.Parser()
+        self._parser.set_language(self.QL_LANGUAGE)
+        self._context = ""
+
+        def on_progress(obj):
+            self.Display(obj["message"], clear_output=True)
+
+        def on_result(obj):
+            self.Display(
+                f"Query completed in {obj['evaluationTime']}!", clear_output=True
+            )
+
+        self._query_client: QueryClient = QueryClient(
+            on_progress=on_progress, on_result=on_result
+        )
+        MetaKernel.__init__(self, **kwargs)
+        print(kwargs)
+
+    def get_usage(self):
+        return "This is the CodeQL kernel."
+
+    def parse_cell(self, cell):
+        """
+        parse the cell code using tree-sitter
+
+        """
+        tree = self._parser.parse(bytes(cell, "utf8"))
+        select_statements = []
+        query_predicates = []
+        captures = self._select_query.captures(tree.root_node)
+        for capture in captures:
+            # capture[0] is the node, capture[1] is the capture name
+            if capture[1] == "select_statement":
+                start_point = capture[0].start_point
+                end_point = capture[0].end_point
+                select_statements.append((start_point, end_point))
+
+        captures = self._predicate_query.captures(tree.root_node)
+        for capture in captures:
+            # capture[0] is the node
+            # capture[1] is the capture name
+            if capture[1] == "annotated_query":
+                start_point = capture[0].start_point
+                end_point = capture[0].end_point
+                # extract the annotation name
+                # check if its a query predicate
+                for i, line in enumerate(cell.split("\n")):
+                    if i == start_point[0]:
+                        if (line[start_point[1]: start_point[1] + len("query")] == "query"):
+                            query_predicates.append((start_point, end_point))
+        return (select_statements, query_predicates)
+
+    def evaluate(self, code, quick_eval=None):
+        """
+        Evaluate the given code and return the result.
+        """
+        try:
+            if not self._query_client._db_metadata:
+                self.Error_display("No database registered! Use %set_database to register a database.")
+                return
+
+            # create a temporary directory to hold the query pack and the query
+            qlpack = "\n".join(
+                [
+                    "---",
+                    "library: false",
+                    "name: jupyter-kernel/temporary-qlpack",
+                    "version: 0.0.1",
+                    "dependencies:",
+                    "  codeql/{}-all: '*'",
+                    "",
+                ]
+            ).format(self._query_client._db_metadata["languages"][0])
+            tmp_dir = mkdtemp(dir="/tmp", prefix="codeql_kernel")
+            with open(os.path.join(tmp_dir, "qlpack.yml"), "w") as f:
+                f.write(qlpack)
+            subprocess.run("codeql pack install", cwd=tmp_dir, shell=True)
+            fd, query_path = mkstemp(suffix=".ql", dir=tmp_dir, text=True)
+            os.write(fd, bytearray(code, "utf-8"))
+            os.close(fd)
+            self.Display("Running query ...", clear_output=True)
+            (err, resp) = self._query_client.run_query(
+                query_path, quick_eval=quick_eval
+            )
+            if err:
+                self.clear_output(wait=True)
+                self.Error_display(
+                    "Error running query: {}".format(err)
+                )
+            else:
+                csv = StringIO(resp)
+                chunks = (chunk for chunk in pd.read_csv(csv, chunksize=5000))
+                df = pd.concat(chunks)
+                self.Display(HTML(df.to_html()), clear_output=True)
+
+        except Exception as e:
+            self.Error_display("Error running query: {}".format(e))
+
+    def do_execute_direct(self, code):
+        """
+        Execute the given code directly.
+        """
+        (select_statements, query_predicates) = self.parse_cell(code)
+        if len(query_predicates) == 1 and len(select_statements) == 0:
+            # we have exactly one query predicate:
+            # add cell to the context and evaluate the query predicate
+            offset = len(self._context.split("\n"))
+            self._context += code + "\n"
+            predicate = query_predicates[0][0]
+            pred_line = predicate[0]
+            pred_col = predicate[1]
+            cell_lines = code.split("\n")
+            words = cell_lines[pred_line].strip().split(" ")
+            position = {
+                "startLine": offset + pred_line,
+                "endLine": offset + pred_line,
+                "startColumn": pred_col + len(words[0]) + len(words[1]) + 3,
+                "endColumn": pred_col + len(words[0]) + len(words[1]) + 3,
+            }
+            self.Display("Evaluating predicate '" + words[2].split("(")[0] + "'", clear_output=True)
+            self.evaluate(self._context, quick_eval=position)
+        elif len(select_statements) == 1:
+            # we have exactly one select statement:
+            # add cell to the context and evaluate the whole context
+            self._context += code + "\n"
+            self.Display("Evaluating select statement ...", clear_output=True)
+            self.evaluate(self._context)
+        else:
+            self._context += code + "\n"
+
+    def repr(self, data):
+        return repr(data)
+
+    def do_shutdown(self, restart):
+        if self._query_client:
+            self._query_client.stop()
+        if restart:
+            self.Print("Restarting kernel...")
+            self.reload_magics()
+            self.restart_kernel()
+            self.Print("Done!")
+        super(CodeQLKernel, self).do_shutdown(restart)
+
+
+if __name__ == "__main__":
+    CodeQLKernel.run_as_main()

codeql_kernel/magics/set_database_magic.py

@@ -0,0 +1,22 @@
+from metakernel import Magic
+import os
+
+
+class SetDatabaseMagic(Magic):
+
+    def line_set_database(self, path):
+        if not os.path.exists(path):
+            self.kernel.Error_display("Databae path does not exist: {}".format(path))
+            return
+        self.kernel.Display("Registering database ...", clear_output=False)
+        err = self.kernel._query_client.register_database(path)
+        if err:
+            self.kernel.Error_display("Error registering database: {}".format(err))
+            return
+        else:
+            self.kernel.Display("Database registered!", clear_output=True)
+            return
+
+
+def register_magics(kernel):
+    kernel.register_magics(SetDatabaseMagic)

codeql_kernel/rawrpc.py

@@ -0,0 +1,131 @@
+import io
+import json
+import threading
+from queue import Queue
+from subprocess import PIPE, Popen
+
+
+class RPC(object):
+    def __init__(self, cmd):
+        self.cmd = cmd
+        self.start()
+        super(RPC, self).__init__()
+
+    def start(self):
+        self._proc: Popen = Popen(self.cmd, stdin=PIPE, stdout=PIPE, stderr=PIPE)
+
+        assert self._proc.stdin
+        assert self._proc.stdout
+        assert self._proc.stderr
+
+        self.stdin = io.open(self._proc.stdin.fileno(), "wb")
+        self.stdout = io.open(self._proc.stdout.fileno(), "rb")
+        self.stderr = io.open(self._proc.stderr.fileno(), "rb")
+
+        self._queue = Queue()
+
+        self.watchdog_stdout = Watchdog(self, name="stdout")
+        self.watchdog_stderr = Watchdog(self, name="stderr")
+
+    def stop(self):
+        if self.stdin and not self.stdin.closed:
+            self.stdin.close()
+        if self.stdout and not self.stdout.closed:
+            self.stdout.close()
+        if self.stderr and not self.stderr.closed:
+            self.stderr.close()
+        self.watchdog_stdout.stop()
+        self.watchdog_stderr.stop()
+        self._proc.terminate()
+        self._proc.wait()
+        self._proc.kill()
+
+    def restart(self):
+        self.stop()
+        self.start()
+
+    def request(self, cmd):
+        self._write(cmd)
+        return self._queue.get()
+
+    def _handle_stdout(self, resp):
+        try:
+            self._queue.put((None, json.loads(resp)))
+        except:
+            self._queue.put((None, resp))
+
+    def _handle_stderr(self, resp):
+        self._queue.put((resp, None))
+
+    def _write(self, s):
+        req = json.dumps(s)
+        req = req + "\0"
+        try:
+            self.stdin.write(bytearray(req, "utf-8"))
+            self.stdin.flush()
+        except:
+            pass
+
+
+class Watchdog(threading.Thread):
+    def __init__(self, rpc, name="watchdog", interval=0.1):
+        super(Watchdog, self).__init__()
+
+        if name == "stderr":
+            self.stream = rpc.stderr
+            self.handle = rpc._handle_stderr
+        elif name == "stdout":
+            self.stream = rpc.stdout
+            self.handle = rpc._handle_stdout
+
+        # store attributes
+        self.rpc = rpc
+        self.name = name
+        self.interval = interval
+        self.daemon = True
+
+        # register a stop event
+        self._stop = threading.Event()
+
+        self.start()
+
+    def start(self):
+        super(Watchdog, self).start()
+
+    def stop(self):
+        self._stop.set()
+
+    def run(self):
+        # reset the stop event
+        self._stop.clear()
+
+        # stop here when stream is not set or closed
+        if not self.stream or self.stream.closed:
+            return
+
+        # read new incoming lines
+        while not self._stop.is_set():
+            resp = None
+
+            # stop when stream is closed
+            if self.stream.closed:
+                break
+
+            try:
+                resp = ""
+                while True:
+                    c = self.stream.read(1).decode("utf-8")
+                    if c == "\x00" and self.name == "stdout":
+                        break
+                    elif c == "\n" and self.name == "stderr":
+                        break
+                    elif not c:
+                        # EOF
+                        break
+                    else:
+                        resp = resp + c
+            except IOError:
+                # prevent residual race conditions occurring when stream is closed externally
+                pass
+            self.handle(resp)
+            self._stop.wait(self.interval)