launcher_low_level.py

This shows how to connect to reven launcher in python, using launcher services from reven low level api.

You will find some helpers to:

• Open/close a project.
• Generate a scenario from a binary.
• List available users, projects and virtual machines.

And a sample use case using them.

#!/usr/bin/env python2

# launcher_low_level.py -f /tmp/crash.html -- ~/samples/lynx32.bin -dump crash.html

import reven_api # to connect reven launcher
import reven # to give a high level handler on a reven server after opening a project

from os.path import basename
from os import system
from time import sleep

### App ########################################################################

import argparse

# default vnc password to use if not set in VM's config
default_vnc_password = 'passw0rd'
# default vnc port to use if not set in VM's config
default_vnc_port = 5900

def cli_args():
    """
    Parse command line arguments.
    """
    parser = argparse.ArgumentParser()
    # parser.add_argument('-v', '--verbose', action='store_true')
    parser.add_argument('--vm', default='debian_jessie_auto', help="VM name to select for scenario recording")
    parser.add_argument('--host', default='localhost', help="Reven launcher's hostname")
    parser.add_argument('--port', type=int, default=8080, help="Reven launcher's port")
    parser.add_argument('--user', default='reven', help="Reven username")
    parser.add_argument('--project', default='sample_project', help="Reven project name")
    parser.add_argument('--kill', action='store_true', help="close running project or abort previous scenario recording if any")
    parser.add_argument('-f', '--file', action='append', metavar='FILE', dest='files', default=[], help="add input file for scenario recording")
    parser.add_argument('binary', nargs='?', help="binary to analyse (will be uploaded as input file)")
    parser.add_argument('args', nargs='*', help="arguments to launch binary with")

    return parser.parse_args()

### Launcher ###################################################################

def connect_to_launcher(host, port):
    """
    Connect to Reven's launcher.
    On error, print and exit.
    @param host: host to connect to.
    @param port: port to connect to.
    @return: Handle to launcher (reven_api.launcher_connection).
    """
    try:
        lc = reven_api.launcher_connection(host, port)
        # add a hostname attribute to lc for future vnc connections
        lc.hostname = host

        return lc
    except RuntimeError as e:
        print("Could not connect to launcher at %s:%s (%s)" % (host, port, e))
        exit()

def vm_config_by_name(launcher_connection, vm_name):
    """
    Get a VM config by name.
    @param launcher_connection: launcher handle (reven_api.launcher_connection).
    @param vm_name: Name of the config to lookup.
    @return: VM config (reven_api.vm_info) or None if not found.
    """
    for vm_cfg in launcher_connection.list_vms():
        if vm_cfg.name == vm_name:
            return vm_cfg

def print_available_vms(launcher_connection):
    """
    Print available VM configurations.
    @param launcher_connection: launcher handle (reven_api.launcher_connection).
    """
    print "Available vm configs:"
    for vm_cfg in launcher_connection.list_vms():
        print "\t%s (%s)" % (vm_cfg.name, vm_cfg.display)

def user_exists(launcher_connection, username):
    """
    @param launcher_connection: launcher handle (reven_api.launcher_connection).
    @param username: Username to check for (str)
    @return: True if a user with username exists.
    """
    return username in launcher_connection.list_users()

def project_exists(launcher_connection, project):
    """
    @param launcher_connection: launcher handle (reven_api.launcher_connection).
    @param project: Project handle to check for (reven_api.project_id)
    @return: True if project exists.
    """
    return project in launcher_connection.list_projects(project.user)

def check_project_not_running(launcher_connection, project, kill=False):
    """
    If a server is running for the given project, kill it or exit.
    @param launcher_connection: launcher handle (reven_api.launcher_connection).
    @param project: Project handle to check for (reven_api.project_id)
    @param kill: If True will try to kill project's server if running.
    """
    try:
        project_port = launcher_connection.project_details(project).reven_server.port
    except RuntimeError:
        # launcher_connection.project_details may legitimately raise if project is not fully set up yet.
        return

    if project_port > 0:
        if kill:
            print "Stopping running project on port %s" % project_port
            launcher_connection.server_kill(project_port)
        else:
            print "Project %s already running on port %s" % (project, project_port)
            exit()

def check_project_not_generating(launcher_connection, project, kill=False):
    """
    If a scenario is being recorder for the project, abort it or exit.
    @param launcher_connection: launcher handle (reven_api.launcher_connection).
    @param project: Project handle to check for (reven_api.project_id)
    @param kill: If True will try to abort any current scenario generation.
    """
    try:
        project_scenario = launcher_connection.project_scenario(project)
    except RuntimeError:
        # launcher_connection.project_scenario may legitimately raise if project is not fully set up yet.
        return

    if project_scenario.is_generating:
        if kill:
            print "Aborting previous scenario generation"
            abort_recording(launcher_connection, project)
        else:
            print "Project %s is already recording a scenario" % (project)
            exit()

### Scenario ###################################################################

def upload_input_files(launcher_connection, project, files):
    """
    Upload input files to project for scenario recording.
    @param launcher_connection: launcher handle (reven_api.launcher_connection)
    @param project: Project to upload files to (reven_api.project_id)
    @param files: List of file paths to be uploaded (list(str))
    @return: yield uploaded filepath.
    """
    for filepath in files:
        launcher_connection.project_upload_file(project, filepath)
        yield filepath

class ScenarioConfigurator(object):
    """
    Helper class to configure scenario recording.
    """

    def __init__(self, vm_cfg, binary=None, args=[], dump_hint=None):
        """
        @param vm_cfg: VM config to use (reven_api.vm_info)
        @param binary: Optional binary to record, filepath must be accessible (str)
        @param args: Arguments to launch the binary with (list(str))
        @param dump_hint: Address or symbol in binary to start recording from (str).
        """
        self._sc = reven_api.scenario_generation_config()
        self._sc.generation.enable_instruction_tracing = False
        self._sc.generation.is_interactive = True

        self._set_vm(vm_cfg)
        if binary:
            self._set_binary(binary, args, dump_hint)

    def _set_vm(self, vm_cfg):
        """
        @param: vm_cfg: VM config to use (reven_api.vm_info)
        """
        self._sc.generation.vnc_password = vm_cfg.vnc_password or default_vnc_password
        self._sc.generation.vnc_port = vm_cfg.vnc_port or default_vnc_port

        self._sc.scenario.vm_config_name = vm_cfg.name
        self._sc.scenario.system_pdb_path = vm_cfg.pdb_path

    def _set_binary(self, binary, args=[], dump_hint=None):
        self._sc.generation.is_interactive = False

        self._sc.scenario.binary_name = basename(binary) if binary else ''
        self._sc.scenario.binary_arguments = ' '.join(args)
        self._sc.scenario.binary_dump_hint = dump_hint or ''
        self._sc.scenario.binary_dump_address = dump_hint or ''

    def validate(self):
        """
        @return: Scenario configuration (reven_api.scenario_generation_config)
        """
        return self._sc

def stop_recording(vnc_host, vnc_port, vnc_password):
    """
    Gracefully stop scenario recording using VNC to send magic key sequence 'F9', 'F7'.
    @param vnc_host: hostname to connect to.
    @param vnc_port: port to connect to.
    @param vnc_password: password to use.
    """
    from vncdotool import api as vnc

    vnc.connect('%s::%s' % (vnc_host, vnc_port), vnc_password).keyDown('f9').keyDown('f7')
    vnc.shutdown()


def abort_recording(launcher_connection, project):
    """
    Abort scenario recording for given project (will discard currently recorded data).
    @param launcher_connection: launcher handle (reven_api.launcher_connection)
    @param project: project to abort recording for (reven_api.project_id)
    """
    launcher_connection.project_abort_scenario_generation(project)

def watch_recording(launcher_connection, project):
    """
    @param launcher_connection: launcher handle (reven_api.launcher_connection)
    @param project: project to watch recording (reven_api.project_id)
    @return: yield recording progress (reven_api.scenario_generation_info)
    """
    recording = True
    while recording:
        progress = launcher_connection.project_scenario(project)
        recording = progress.is_generating

        yield progress

def wait_recording(launcher_connection, project, soft_timeout=10, hard_timeout=30):
    """
    Blocks until scenario is recorded or stopped.
    @param launcher_connection: launcher handle (reven_api.launcher_connection)
    @param project: project to wait recording (reven_api.project_id)
    @param soft_timeout: Timeout in seconds before trying to stop scenario recording
    @param hard_timeout: Timeout in seconds before trying to abort scenario recording (recorded data are discarded)
    """
    for i, progress in enumerate(watch_recording(launcher_connection, project)):
        if progress.log_chunk:
            print progress.log_chunk

        if i == soft_timeout:
            print "Soft timeout: stopping scenario recording"
            try:
                stop_recording(launcher_connection.hostname, vnc_port, vnc_password)
            except Exception as e:
                print "Could not gracefully stop recording through vnc (%s)" % e
        elif i == hard_timeout:
            print "Hard timeout: aborting scenario recording"
            abort_recording(launcher_connection, project)
        else:
            sleep(1)

    return progress.is_successful

### Reven ######################################################################

def open_project(launcher_connection, project):
    """
    Open project (launch a server)
    @param launcher_connection: launcher handle (reven_api.launcher_connection)
    @param project: project to open (reven_api.project_id)
    @return: handle to project (reven.Project)
    """
    # actually open project
    port = launcher_connection.server_launch(project)
    if not port:
        print "Error server not started for project %s" % project

    print "Server started on port %s" % port

    # reven server cannnot accept connections immediatly after beeing launched
    # try to reconnect if first attempts failed
    client = None
    for _ in range(10):
        try:
            client = reven.Project(launcher_connection.hostname, port)
            print "Connected!"
            break
        except RuntimeError as e:
            print "Could not connect (%s)" % e
            sleep(1)

    return client


def wait_execution(reven_client):
    """
    Block until execution is completed.
    @param reven_client: opened project handle (reven.Project)
    """
    while True:
        progress = reven_client.execution_status()

        print "# %s/%s - %s/%s - %s" % (
                progress.tsc, progress.last_tsc,
                progress.point_index, progress.last_point_index,
                progress.status)

        if not progress.is_busy:
            break

        sleep(1)


if __name__ == '__main__':
    # Arguments ############################################################

    args = cli_args()

    launcher_hostname = args.host
    launcher_port = args.port
    sample_project = reven_api.project_id(args.user, args.project)
    sample_binary = args.binary
    sample_binary_arguments = args.args
    sample_input_files = args.files
    sample_vm = args.vm
    sample_overkill = args.kill

    sample_is_interactive = not sample_binary

    # Connect ##############################################################

    print "Connecting to %s:%s" % (launcher_hostname, launcher_port)

    lc = connect_to_launcher(launcher_hostname, launcher_port)

    # Check Vm #############################################################

    vm_cfg = vm_config_by_name(lc, sample_vm)
    if not vm_cfg:
        print "Vm config %s not found." % (sample_vm)
        print_available_vms(lc)
        exit()

    print "Using vm config %s: %s" % (sample_vm, vm_cfg.display)

    # Check project ########################################################

    if not user_exists(lc, sample_project.user):
        print "User %s does not exist (will be created)" % sample_project.user

    if project_exists(lc, sample_project):
        print "Project %s already exists, checking project status" % (sample_project)
        check_project_not_running(lc, sample_project, kill=sample_overkill)
        check_project_not_generating(lc, sample_project, kill=sample_overkill)
    else:
        print "Project %s does not exists, creating..." % (sample_project)
        lc.project_create(sample_project)

    print "Saving to project %s" % (sample_project)

    # Scenario setup #######################################################

    # Input files

    if sample_binary and not sample_binary in sample_input_files:
        print "Adding binary %s to files being uploaded" % (sample_binary)
        sample_input_files.append(sample_binary)

    print "Uploading files: %s" % (sample_input_files)

    for f in upload_input_files(lc, sample_project, sample_input_files):
        print "\tUploaded %s" % f

    # Recording

    print "Recording: %s" % ("<interactive>" if not sample_binary else "%s with %s" % (sample_binary, sample_binary_arguments))

    print "Configuring scenario recording"

    scenario_config = ScenarioConfigurator(vm_cfg=vm_cfg, binary=sample_binary, args=sample_binary_arguments, dump_hint=None).validate()

    print "Start scenario recording"

    vnc_port = scenario_config.generation.vnc_port
    vnc_password = scenario_config.generation.vnc_password

    print "VNC conf %s::%s -p %s" % (lc.hostname, vnc_port, vnc_password)

    # Scenario recording ###################################################

    lc.project_generate_scenario(sample_project, scenario_config)

    is_scenario_successful = wait_recording(lc, sample_project)

    if is_scenario_successful:
        print "Scenario recorded!"
    else:
        print "Error while recording scenario"
        exit()

    # Launch Reven #########################################################

    client = open_project(lc, sample_project)
    if not client:
        print "Error while connecting to project"
        exit()
    port = client.port

    # Execution ############################################################

    print "Launching execution"

    client.start_execution()
    wait_execution(client)

    # Launch axion #########################################################

    axion_launch_command = "axion --host %s --port %s" % (lc.hostname, port)
    print axion_launch_command
    system(axion_launch_command)