~ben-hutchings/ensoft-sextant/upload-perf

# -----------------------------------------
# Sextant
# Copyright 2014, Ensoft Ltd.
# Author: Patrick Stevens, using work from Patrick Stevens and James Harkin
# -----------------------------------------
# API to interact with a Neo4J server: upload, query and delete programs in a DB

__all__ = ("Validator", "AddToDatabase", "FunctionQueryResult", "Function",
           "SextantConnection")

import re  # for validation of function/program names
import logging
from datetime import datetime
import os
import getpass
from collections import namedtuple

from neo4jrestclient.client import GraphDatabase
import neo4jrestclient.client as client

COMMON_CUTOFF = 10
# a function is deemed 'common' if it has more than this
# many connections


class Validator():
    """ Sanitises/checks strings, to prevent Cypher injection attacks"""

    @staticmethod
    def validate(input_):
        """
        Checks whether we can allow a string to be passed into a Cypher query.
        :param input_: the string we wish to validate
        :return: bool(the string is allowed)
        """
        regex = re.compile(r'^[A-Za-z0-9\-:\.\$_@\*\(\)%\+,]+$')
        return bool(regex.match(input_))

    @staticmethod
    def sanitise(input_):
        """
        Strips harmful characters from the given string.
        :param input_: string to sanitise
        :return: the sanitised string
        """
        return re.sub(r'[^\.\-_a-zA-Z0-9]+', '', input_)


class AddToDatabase():
    """Updates the database, adding functions/calls to a given program"""

    def __init__(self, program_name='', sextant_connection=None,
                 uploader='', uploader_id='', date=None):
        """
        Object which can be used to add functions and calls to a new program
        :param program_name: the name of the new program to be created
          (must already be validated against Validator)
        :param sextant_connection: the SextantConnection to use for connections
        :param uploader: string identifier of user who is uploading
        :param uploader_id: string Unix user-id of logged-in user
        :param date: string date of today
        """
        # program_name must be alphanumeric, to avoid injection attacks easily
        if not Validator.validate(program_name):
            return

        self.program_name = program_name
        self.parent_database_connection = sextant_connection

        self._funcs_tx = None  # transaction for uploading functions
        self._calls_tx = None  # transaction for uploading relationships

        self._calldict = {}

        if self.parent_database_connection:
            # we'll locally use db for short
            db = self.parent_database_connection._db

            parent_function = db.nodes.create(name=program_name,
                                              type='program',
                                              uploader=uploader,
                                              uploader_id=uploader_id,
                                              date=date)
            self._parent_id = parent_function.id

            self._funcs_tx = db.transaction(using_globals=False, for_query=True)
            self._calls_tx = db.transaction(using_globals=False, for_query=True)

    @staticmethod
    def _get_display_name(function_name):
        """
        Gets the name we will display to the user for this function name.

        For instance, if function_name were __libc_start_main@plt, we would
        return ("__libc_start_main", "plt_stub"). The returned function type is
        currently one of "plt_stub", "function_pointer" or "normal".

        :param function_name: the name straight from objdump of a function
        :return: ("display name", "function type")

        """

        if function_name[-4:] == "@plt":
            display_name = function_name[:-4]
            function_group = "plt_stub"
        elif function_name[:20] == "_._function_pointer_":
            display_name = function_name
            function_group = "function_pointer"
        else:
            display_name = function_name
            function_group = "normal"

        return display_name, function_group

    def add_function(self, function_name):
        """
        Adds a function to the program, ready to be sent to the remote database.
        If the function name is already in use, this method effectively does
          nothing and returns True.

        :param function_name: a string which must be alphanumeric
        :return: True if the request succeeded, False otherwise
        """
        if not Validator.validate(function_name):
            return False
        if function_name in self._calldict:
            return True

        display_name, function_group = self._get_display_name(function_name)

        query = ('START n = node({}) '
                 'CREATE (n)-[:subject]->(m:func {{type: "{}", name: "{}"}}) '
                 'RETURN m.name, id(m)')
        query = query.format(self._parent_id, function_group, display_name)

        self._funcs_tx.append(query)

        self._calldict[function_name] = set()

        return True

    def add_function_call(self, fn_calling, fn_called):
        """
        Adds a function call to the program, ready to be sent to the database.
        Effectively does nothing if there is already a function call between
          these two functions.
        Function names must be alphanumeric for easy security purposes;
          returns False if they fail validation.  :param fn_calling: the name of the calling-function as a string.
          It should already exist in the AddToDatabase; if it does not,
          this method will create a stub for it.
        :param fn_called: name of the function called by fn_calling.
          If it does not exist, we create a stub representation for it.
        :return: True if successful, False otherwise
        """
        if not all((Validator.validate(fn_calling),
                    Validator.validate(fn_called))):
            return False
    
        if not fn_called in self._calldict:
            self.add_function(fn_called)
        if not fn_calling in self._calldict:
            self.add_function(fn_calling)


        self._calldict[fn_calling].add(fn_called)

        return True

    def commit(self):
        """
        Call this when you are finished with the object.
        Changes are not synced to the remote database until this is called.
        """
 
        functions = (result[0] for result in self._funcs_tx.commit())  # send off the function names
        id_funcs = dict(functions)
        
        logging.info('Functions uploaded.')

        # so id_funcs is a dict with id_funcs['name'] == id
        for (caller, called) in self._calldict.items():
            if not called:
                pass
            else:
                # add all the connections for this caller in one query
                caller_id = id_funcs[self._get_display_name(caller)[0]]
                called_ids = (id_funcs[self._get_display_name(fn)[0]] for fn in called)

                query = (' MATCH n WHERE id(n) = {}'
                         ' UNWIND[{}] as called_id'
                         ' MATCH m WHERE id(m) = called_id'
                         ' CREATE (n)-[:calls]->(m)')

                query = query.format(caller_id, ','.join(str(i) for i in called_ids))

                self._calls_tx.append(query)

        self._calls_tx.commit()
        logging.info('Calls uploaded')


class FunctionQueryResult:
    """A graph of function calls arising as the result of a Neo4J query."""

    def __init__(self, parent_db, program_name='', rest_output=None):
        self.program_name = program_name
        self._parent_db_connection = parent_db
        self.functions = self._rest_node_output_to_graph(rest_output)
        self._update_common_functions()

    def __eq__(self, other):
        # we make a dictionary so that we can perform easy comparison
        selfdict = {func.name: func for func in self.functions}
        otherdict = {func.name: func for func in other.functions}

        return self.program_name == other.program_name and selfdict == otherdict

    def _update_common_functions(self):
        """
        Loop over all functions: increment the called-by count of their callees.
        """
        for func in self.functions:
            for called in func.functions_i_call:
                called.number_calling_me += 1

    def _rest_node_output_to_graph(self, rest_output):
        """
        Convert the output of a REST API query into our internal representation.
        :param rest_output: output of the REST call as a Neo4j QuerySequence
        :return: iterable of <Function>s ready to initialise self.functions.
        """

        if rest_output is None or not rest_output.elements:
            return []

        # how we store this is: a dict
        #   with keys  'functionname'
        #   and values [the function object we will use,
        #               and a set of (function names this function calls),
        #               and numeric ID of this node in the Neo4J database]

        result = {}

        # initial pass for names of functions

        # if the following assertion failed, we've probably called db.query
        # to get it to not return client.Node objects, which is wrong.
        # we attempt to handle this a bit later; this should never arise, but
        # we can cope with it happening in some cases, like the test suite

        if type(rest_output.elements) is not list:
            logging.warning('Not a list: {}'.format(type(rest_output.elements)))

        for node_list in rest_output.elements:
            assert(isinstance(node_list, list))
            for node in node_list:
                if isinstance(node, client.Node):
                    name = node.properties['name']
                    node_id = node.id
                    node_type = node.properties['type']
                else:  # this is the handling we mentioned earlier;
                    # we are a dictionary instead of a list, as for some
                    # reason we've returned Raw rather than Node data.
                    # We should never reach this code, but just in case.
                    name = node['data']['name']
                    # hacky workaround to get the id
                    node_id = node['self'].split('/')[-1]
                    node_type = node['data']['type']

                result[name] = [Function(self.program_name,
                                         function_name=name,
                                         function_type=node_type),
                                set(),
                                node_id]

        # end initialisation of names-dictionary

        if self._parent_db_connection is not None:
            # This is the normal case, of extracting results from a server.
            # We leave the other case in because it is useful for unit testing.

            # We collect the name-name pairs of caller-callee, batched for speed
            new_tx = self._parent_db_connection.transaction(using_globals=False,
                                                            for_query=True)
            for index in result:
                q = ("START n=node({})"
                     "MATCH n-[calls:calls]->(m)"
                     "RETURN n.name, m.name").format(result[index][2])
                new_tx.append(q)

            logging.debug('exec')
            results = new_tx.execute()

            # results is a list of query results, each of those being a list of
            # calls.

            for call_list in results:
                if call_list:
                    # call_list has element 0 being an arbitrary call this
                    # function makes; element 0 of that call is the name of the
                    # function itself. Think {{'orig', 'b'}, {'orig', 'c'}}.
                    orig = call_list[0][0]
                    # result['orig'] is [<Function>, ('callee1','callee2')]
                    result[orig][1] |= set(list(zip(*call_list.elements))[1])
                    # recall: set union is denoted by |

        else:
            # we don't have a parent database connection.
            # This has probably arisen because we created this object from a
            # test suite, or something like that.
            for node in rest_output.elements:
                node_name = node[0].properties['name']
                result[node_name][1] |= {relationship.end.properties['name']
                                         for relationship in node[0].relationships.outgoing()}

        logging.debug('Relationships complete.')

        # named_function takes a function name and returns the Function object
        # with that name, or None if none exists.
        named_function = lambda name: result[name][0] if name in result else None

        for function, calls, node_id in result.values():
            what_i_call = [named_function(name)
                           for name in calls
                           if named_function(name) is not None]
            function.functions_i_call = what_i_call

        return [list_element[0]
                for list_element in result.values()
                if list_element[0]]

    def get_functions(self):
        """
        :return: a list of Function objects present in the query result
        """
        return self.functions

    def get_function(self, name):
        """
        Given a function name, returns the Function object which has that name.
        If no function with that name exists, returns None.
        """
        func_list = [func for func in self.functions if func.name == name]
        return None if len(func_list) == 0 else func_list[0]


def set_common_cutoff(common_def):
    """
    Sets the number of incoming connections at which we deem a function 'common'
    Default is 10 (which is used if this method is never called).
    :param common_def: number of incoming connections
    """
    global COMMON_CUTOFF
    COMMON_CUTOFF = common_def


class Function(object):
    """Represents a function which might appear in a FunctionQueryResult."""

    def __eq__(self, other):
        funcs_i_call_list = {func.name for func in self.functions_i_call}
        funcs_other_calls_list = {func.name for func in other.functions_i_call}

        return (self.parent_program == other.parent_program
                and self.name == other.name
                and funcs_i_call_list == funcs_other_calls_list
                and self.attributes == other.attributes)

    @property
    def number_calling_me(self):
        return self._number_calling_me

    @number_calling_me.setter
    def number_calling_me(self, value):
        self._number_calling_me = value
        self.is_common = (self._number_calling_me > COMMON_CUTOFF)

    def __init__(self, program_name='', function_name='', function_type=''):
        self.parent_program = program_name
        self.attributes = []
        self.type = function_type
        self.functions_i_call = []
        self.name = function_name
        self.is_common = False
        self._number_calling_me = 0
        # care: _number_calling_me is not automatically updated, except by
        # any invocation of FunctionQueryResult._update_common_functions.


class SextantConnection:
    """
    RESTful connection to a remote database.
    It can be used to create/delete/query programs.
    """

    ProgramWithMetadata = namedtuple('ProgramWithMetadata',
                                     ['uploader', 'uploader_id',
                                      'program_name', 'date', 
                                      'number_of_funcs'])

    def __init__(self, url):
        self.url = url
        self._db = GraphDatabase(url)

    def new_program(self, name_of_program):
        """
        Request that the remote database create a new program with the given name.
        This procedure will create a new program remotely; you can manipulate
          that program using the returned AddToDatabase object.
        The name can appear in the database already, but this is not recommended
          because then delete_program will not know which to delete. Check first
          using self.check_program_exists.
        The name specified must pass Validator.validate()ion; this is a measure
          to prevent Cypher injection attacks.
        :param name_of_program: string program name
        :return: AddToDatabase instance if successful
        """

        if not Validator.validate(name_of_program):
            raise ValueError(
                "{} is not a valid program name".format(name_of_program))
        
        uploader = getpass.getuser()
        uploader_id = os.getuid()

        return AddToDatabase(sextant_connection=self,
                             program_name=name_of_program,
                             uploader=uploader, uploader_id=uploader_id,
                             date=str(datetime.now()))

    def delete_program(self, name_of_program):
        """
        Request that the remote database delete a specified program.
        :param name_of_program: a string which must be alphanumeric only
        :return: bool(request succeeded)
        """
        if not Validator.validate(name_of_program):
            return False

        q = """MATCH (n) WHERE n.name= "{}" AND n.type="program"
        OPTIONAL MATCH (n)-[r]-(b) OPTIONAL MATCH (b)-[rel]-()
        DELETE  b,rel DELETE n, r""".format(name_of_program)

        self._db.query(q)

        return True

    def _execute_query(self, prog_name='', query=''):
        """
        Executes a Cypher query against the remote database.
        Note that this returns a FunctionQueryResult, so is unsuitable for any
          other expected outputs (such as lists of names). For those instances,
          it is better to run self._parent_database_connection_object.query
          explicitly.
        Intended only to be used for non-updating queries
          (such as "get functions" rather than "create").
        :param prog_name: name of the program the result object will reflect
        :param query: verbatim query we wish the server to execute
        :return: a FunctionQueryResult corresponding to the server's output
        """
        rest_output = self._db.query(query, returns=client.Node)

        return FunctionQueryResult(parent_db=self._db,
                                   program_name=prog_name,
                                   rest_output=rest_output)

    def get_program_names(self):
        """
        Execute query to retrieve a list of all programs in the database.
        Any name in this list can be used verbatim in any SextantConnection
          method which requires a program-name input.
        :return: a list of function-name strings.
        """
        q = """MATCH (n) WHERE n.type = "program" RETURN n.name"""
        program_names = self._db.query(q, returns=str).elements

        result = [el[0] for el in program_names]

        return set(result)

    def programs_with_metadata(self):
        """
        Returns a set of namedtuples which represent the current database.
        
        The namedtuples have .uploader, .uploader_id, .program_name, .date,
        .number_of_funcs.
        :return: set of namedtuples
       
        """
        
        q = ("MATCH (base) WHERE base.type = 'program' "
             "MATCH (base)-[:subject]->(n)"
             "RETURN base.uploader, base.uploader_id, base.name, base.date, count(n)")
        result = self._db.query(q)
        return {self.ProgramWithMetadata(*res) for res in result}

    def check_program_exists(self, program_name):
        """
        Execute query to check whether a program with the given name exists.
        Returns False if the program_name fails validation against Validator.
        :return: bool(the program exists in the database).
        """

        if not Validator.validate(program_name):
            return False

        q = ("MATCH (base) WHERE base.name = '{}' AND base.type = 'program' "
             "RETURN count(base)").format(program_name)

        result = self._db.query(q, returns=int)
        return result.elements[0][0] > 0

    def check_function_exists(self, program_name, function_name):
        """
        Execute query to check whether a function with the given name exists.
        We only check for functions which are children of a program with the
          given program_name.
        :param program_name: string name of the program within which to check
        :param function_name: string name of the function to check for existence
        :return: bool(names validate correctly, and function exists in program)
        """
        if not self.check_program_exists(program_name):
            return False

        if not Validator.validate(program_name):
            return False

        q = ("MATCH (base) WHERE base.name = '{}' AND base.type = 'program'"
             "MATCH (base)-[r:subject]->(m) WHERE m.name = '{}'"
             "RETURN count(m)").format(program_name, function_name)

        result = self._db.query(q, returns=int)
        return result.elements[0][0] > 0

    def get_function_names(self, program_name):
        """
        Execute query to retrieve a list of all functions in the program.
        Any of the output names can be used verbatim in any SextantConnection
          method which requires a function-name input.
        :param program_name: name of the program whose functions to retrieve
        :return: None if program_name doesn't exist in the remote database,
          a set of function-name strings otherwise.
        """

        if not self.check_program_exists(program_name):
            return None

        q = ("MATCH (base) WHERE base.name = '{}' AND base.type = 'program' "
             "MATCH (base)-[r:subject]->(m) "
             "RETURN  m.name").format(program_name)
        return {func[0] for func in self._db.query(q)}

    def get_all_functions_called(self, program_name, function_calling):
        """
        Execute query to find all functions called by a function (indirectly).
        If the given function is not present in the program, returns None;
          likewise if the program_name does not exist.
        :param program_name: a string name of the program we wish to query under
        :param function_calling: string name of a function whose children to find
        :return: FunctionQueryResult, maximal subgraph rooted at function_calling
        """

        if not self.check_program_exists(program_name):
            return None

        if not self.check_function_exists(program_name, function_calling):
            return None

        # @@@ type in query - does it matter?
        q = """MATCH (base) WHERE base.name = '{}' ANd base.type = 'program'
            MATCH (base)-[:subject]->(m) WHERE m.name='{}'
            MATCH (m)-[:calls*]->(n)
            RETURN distinct n, m""".format(program_name, function_calling)

        return self._execute_query(program_name, q)

    def get_all_functions_calling(self, program_name, function_called):
        """
        Execute query to find all functions which call a function (indirectly).
        If the given function is not present in the program, returns None;
          likewise if the program_name does not exist.
        :param program_name: a string name of the program we wish to query
        :param function_called: string name of a function whose parents to find
        :return: FunctionQueryResult, maximal connected subgraph with leaf function_called
        """

        if not self.check_program_exists(program_name):
            return None

        if not self.check_function_exists(program_name, function_called):
            return None

        q = """MATCH (base) WHERE base.name = '{}' AND base.type = 'program'
            MATCH (base)-[r:subject]->(m) WHERE m.name='{}'
            MATCH (n)-[:calls*]->(m) WHERE n.name <> '{}'
            RETURN distinct n , m"""
        q = q.format(program_name, function_called, program_name)

        return self._execute_query(program_name, q)

    def get_call_paths(self, program_name, function_calling, function_called):
        """
        Execute query to find all possible routes between two specific nodes.
        If the given functions are not present in the program, returns None;
          ditto if the program_name does not exist.
        :param program_name: string program name
        :param function_calling: string
        :param function_called: string
        :return: FunctionQueryResult, the union of all subgraphs reachable by
          adding a source at function_calling and a sink at function_called.
        """

        if not self.check_program_exists(program_name):
            return None

        if not self.check_function_exists(program_name, function_called):
            return None

        if not self.check_function_exists(program_name, function_calling):
            return None

        q = r"""MATCH (pr) WHERE pr.name = '{}' AND pr.type = 'program'
                MATCH p=(start {{name: "{}" }})-[:calls*]->(end {{name:"{}"}})
                  WHERE (pr)-[:subject]->(start)
                WITH DISTINCT nodes(p) AS result
                UNWIND result AS answer
                RETURN answer"""
        q = q.format(program_name, function_calling, function_called)

        return self._execute_query(program_name, q)

    def get_whole_program(self, program_name):
        """Execute query to find the entire program with a given name.
        If the program is not present in the remote database, returns None.
        :param: program_name: a string name of the program we wish to return.
        :return: a FunctionQueryResult consisting of the program graph.
        """

        if not self.check_program_exists(program_name):
            return None

        query = """MATCH (base) WHERE base.name = '{}' AND base.type = 'program'
                MATCH (base)-[subject:subject]->(m)
                RETURN DISTINCT (m)""".format(program_name)

        return self._execute_query(program_name, query)

    def get_shortest_path_between_functions(self, program_name, func1, func2):
        """
        Execute query to get a single, shortest, path between two functions.
        :param program_name: string name of the program we wish to search under
        :param func1: the name of the originating function of our shortest path
        :param func2: the name of the function at which to terminate the path
        :return: FunctionQueryResult shortest path between func1 and func2.
        """
        if not self.check_program_exists(program_name):
            return None

        if not self.check_function_exists(program_name, func1):
            return None

        if not self.check_function_exists(program_name, func2):
            return None

        q = """MATCH (func1 {{ name:"{}" }}),(func2 {{ name:"{}" }}),
            p = shortestPath((func1)-[:calls*]->(func2))
            UNWIND nodes(p) AS ans
            RETURN ans""".format(func1, func2)

        return self._execute_query(program_name, q)