cb-test

#!/usr/bin/python

"""
CB Testing tool

Copyright (C) 2014,2015 - Brian Caswell <bmc@lungetech.com>

Permission is hereby granted, free of charge, to any person obtaining a copy
of this software and associated documentation files (the "Software"), to deal
in the Software without restriction, including without limitation the rights
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
copies of the Software, and to permit persons to whom the Software is
furnished to do so, subject to the following conditions:

The above copyright notice and this permission notice shall be included in
all copies or substantial portions of the Software.

THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
THE SOFTWARE.

This tool allows verification of POV and POLLs with a CGC challenge binary
using 'cb-replay', 'tcpdump', and 'cb-server'.
"""

import time
import re
import uuid
import argparse
import platform
import glob
import logging
import os
import random
import resource
import signal
import socket
import subprocess
import sys
import thread
import threading
import Queue
import ansi_x931_aes128


class TimeoutException(Exception):
    """ TimeoutException - A class used to catch just timeouts created by
    Timeout() """
    pass


class Timeout(object):
    """ Timeout - A class to use within 'with' for timing out a block via
    exceptions and alarm."""

    def __init__(self, seconds):
        """ Instantiate a Timeout() instance, to be used in a 'with' context.

        Arguments:
            seconds: time until timeout

        Returns:
            None

        Raises:
            AssertionError: if a POV action is not in the pre-defined methods
        """

        assert isinstance(seconds, int)
        assert seconds > 0
        self.seconds = seconds

    @staticmethod
    def cb_handle_timeout(signum, frame):
        """ signal handler callback method, called when SIGALRM is raised

        Arguments:
            signum: signal number that caused the signal handler to be called
            frame:  frame that caused the signal

        Returns:
            None

        Raises:
            TimeoutException: Always raises this exception
        """

        raise TimeoutException("timed out")

    def __enter__(self):
        """ Context guard for handling timeouts as a 'with'

        Arguments:
            None

        Returns:
            None

        Raises:
            None
        """

        signal.signal(signal.SIGALRM, self.cb_handle_timeout)
        signal.alarm(self.seconds)

    def __exit__(self, exit_type, exit_value, traceback):
        """ Disable the alarm upon exiting the 'with' block

        Arguments:
            exit_type: the type of exit code being generated
            exit_value: the value of the exit method
            traceback: a stack trace

        Raises:
            None
        """

        signal.alarm(0)


class Background(object):
    """ Run an external command in a background thread

    Usage:
        a = Background(['echo', 'hello'])
        a.wait()

    Attributes:
        cmd: The command that should be run (as a list)
        process: The subprocess handle
        threads:  A list of threads being maintained, one for stderr, one for
            stdout.
    """
    def __init__(self, cmd, command_name=None):
        logging.warning('launching %s', ' '.join(cmd))
        if command_name is None:
            self.cmd = cmd[0]
        else:
            self.cmd = command_name
        self.process = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                        stderr=subprocess.PIPE)
        self.threads = []
        self.log_queue = Queue.Queue()
        self.log_handle(self.process.stdout, False)
        self.log_handle(self.process.stderr, True)

    def log_handle(self, filehandle, should_repr):
        """ Create a thread to log all of the output from the provided file
        handle.

        Arguments:
            filehandle: File handle to monitor
            should_repr: Should the output be passed through 'repr' before
                writing to the log.  (Useful for untrusted input)

        Returns:
            None

        Raises:
            None
        """
        def log_background(log, should_repr, queue):
            """ Thread callback to log the output from the background process

            Arguments:
                log: File handle to monitor
                should_repr: should the output be passed through repr before
                    writing to the log.  (Useful for untrusted input)
            """
            try:
                for line in iter(log.readline, ''):
                    queue.put(line)
                    if should_repr:
                        logging.warning('%s: %s', self.cmd, repr(line[:-1]))
                    else:
                        logging.warning('%s: %s', self.cmd, line[:-1])
                log.close()
                queue.put(None)
            except KeyboardInterrupt:
                thread.interrupt_main()

        my_thread = threading.Thread(target=log_background,
                                     args=(filehandle, should_repr, self.log_queue))
        my_thread.daemon = True
        my_thread.start()
        self.threads.append(my_thread)

    def terminate(self):
        """ Terminate the running process and associated threads.

        Arguments:
            None

        Returns:
            None

        Raises:
            None
        """
        logging.debug('terminating %s', self.cmd)
        try:
            self.process.terminate()
        except OSError:
            pass
        for my_thread in self.threads:
            my_thread.join()

    def wait(self):
        """ Wait for the process to exit.

        Arguments:
            None

        Returns:
            Return code from the process

        Raises:
            None
        """
        logging.debug('waiting for %s to terminate (pid: %s)', self.cmd, self.process.pid)
        wval = self.process.wait()
        result = []

        waiting = 2
        while waiting > 0:
            item = self.log_queue.get()
            if item is None:
                waiting -= 1
                continue
            result.append(item)
            self.log_queue.task_done()

        logging.debug('process returned %s', repr(wval))
        return wval, ''.join(result)


class Runner(object):
    """ Run an external command in a background thread

    Usage:
        a = Runner(port, cb_list, xml_list, pcap, wrapper, directory,
                   should_core, failure_ok, should_debug, timeout, log_fh,
                   cb_seed, cb_seed_skip, max_send, concurrent,
                   negotiate_seed, pov_seed, cb_no_attach, cb_seed_munge,
                   cb_env)
        a.run()

    Attributes:
        port: TCP port the CB should run on.
        cb_list: Set of CBs being tested
        xml_list: List of XML files to test the CB with
        pcap: Path the pcap file should be written to
        wrapper: The wrapper executable to run the CBs through
        directory: The directory containing the CBs
        should_core: Is the CB expected to core with the provided XML
        failure_ok: Should failures stop the testing process
        should_debug: Should debugging be enabled on child processes
        timeout: Timeout for child processes
        log_fh: The file handle for the log
        ids: Path to the IDS ruleset
        server_cb: Node the CB should run on (cb-server)
        server_pov: Node the XML should be run on (cb-replay)
        server_ids: Node the IDS should run on (cb-ids)
        ip_address: IP Address of the CB server
        cb_seed: PRNG for CBs
        cb_seed_skip: Amount to advance the PRNG for CBs
        max_send: Maximum data to send per request from the poll
        concurrent: Number of polls/povs to send concurrently
        negotiate_seed: Should the CB seed be negotiated from cb-replay
        pov_seed: the PRNG seed for POVs
        cb_no_attach: Should the CB not be attached within cb-server
        cb_seed_munge: Should the poll munge the seed
        cb_env: An enviornment variable to pass through to the CB in cb-server
    """
    pov_signals = [signal.SIGSEGV, signal.SIGILL, signal.SIGBUS]

    def __init__(self, port, cb_list, xml_list, pcap, wrapper, directory,
                 should_core, failure_ok, should_debug, timeout, log_fh,
                 cb_seed, cb_seed_skip, max_send, concurrent, negotiate_seed,
                 pov_seed, cb_no_attach, cb_seed_munge, cb_env):
        self.port = port
        self.cb_list = cb_list
        self.cb_no_attach = cb_no_attach
        self.cb_env = cb_env
        self.xml_list = xml_list
        self.pcap = pcap
        self.wrapper = wrapper
        self.concurrent = concurrent
        self.directory = directory
        self.should_core = should_core
        self.should_debug = should_debug
        self.failure_ok = failure_ok
        self.timeout = timeout
        self.processes = []
        self.log_fh = log_fh
        self.cb_seed = cb_seed
        self.cb_seed_skip = cb_seed_skip
        self.cb_seed_munge = cb_seed_munge
        self.ids_rules = None
        self.max_send = max_send
        self.negotiate_seed = negotiate_seed
        self.pov_seed = pov_seed

        self._remote = False
        self._remote_uploaded = []

        self.server_cb = None
        self.server_pov = None
        self.server_ids = None
        self._test_id = None
        self._tmp_dir = None

        self.ip_address = self.random_ip()
        if not isinstance(self.port, int) or self.port == 0:
            self.port = self.random_port()

        resource.setrlimit(resource.RLIMIT_CORE, (resource.RLIM_INFINITY,
                                                  resource.RLIM_INFINITY))

    def remote_cmd(self, node, cmd):
        """ Run the command on a remote host via ssh

        Arguments:
            node: host that the command should be run
            cmd: command that should be run

        Returns:
            Return the stdout from the command

        Raises:
            AssertionError: if the node is not in the set of known servers
            subprocess.error: on command excution failure
        """
        assert node in [self.server_cb, self.server_pov, self.server_ids]

        return subprocess.check_output(['ssh', node] + cmd)

    def remote_upload(self, node, files):
        """ Upload a set of files to the tempdir on the remote node.

        NOTE: This method maintains a history of what we've uploaded, and only
        upload it once.

        Arguments:
            host: host that the command should be run
            files: Files to be uploaded

        Returns:
            None

        Raises:
            subprocess.error: on command excution failure
        """
        assert node in [self.server_cb, self.server_pov, self.server_ids]
        assert isinstance(files, list)

        files = list(set(files) - set(self._remote_uploaded))

        if len(files):
            subprocess.check_output(['scp'] + files +
                                    ['%s:%s' % (node, self._tmp_dir)])

            self._remote_uploaded += files

    def set_ids_rules(self, ids_rules):
        """ Set IDS rules to use during teting

        Arguments:
            ids_rules: Rules to use

        Returns:
            None

        Raises:
            None
        """

        self.ids_rules = ids_rules

    def enable_remote(self, server_cb, server_pov, server_ids):
        """ Setup remote cb testing

        Arguments:
            server_cb: The node to run the cb-server
            server_pov: The node to run the POVs
            server_ids: The node to run the IDS

        Returns:
            None

        Raises:
            subprocess.error: on command excution failure
        """
        self.server_cb = server_cb
        self.server_pov = server_pov
        self.server_ids = server_ids
        self._remote = True
        self._test_id = str(uuid.uuid4())
        self._tmp_dir = '/tmp/%s' % self._test_id
#        if self.ids_rules is None:
#            self.ip_address = self.server_cb
#        else:
        self.ip_address = self.server_ids

        for node in [server_cb, server_pov, server_ids]:
            self.remote_cmd(node, ['mkdir', '-p', self._tmp_dir])

    def background(self, cmd, cmd_name=None):
        """ Run a command in the background, and verify it started

        Arguments:
            cmd: Command to be run
            cmd_name: Name of the command to show up in the logs

        Returns:
            None

        Raises:
            Exception: if the process didn't get created effectively
        """
        process = Background(cmd, cmd_name)
        self.processes.append(process)
        if process.process.poll():
            raise Exception('background process failed: %s' % (' '.join(cmd)))
        return process

    def launch(self, cmd):
        """ Run a command in the foreground, logging the stderr and stdout

        Arguments:
            cmd: Command to be run

        Returns:
            None

        Raises:
            None
        """
        logging.warning('launching %s', ' '.join(cmd))
        process = subprocess.Popen(cmd, stdout=subprocess.PIPE,
                                   stderr=subprocess.PIPE)
        stdout, stderr = process.communicate()
        if len(stderr):
            for line in stderr.split('\n'):
                logging.error('%s (stderr): %s', cmd[0], repr(line))
        self.log_fh.flush()
        self.log_fh.write(stdout)
        self.log_fh.flush()
        return process.returncode, stdout

    def cleanup(self):
        """ Cleanup the current enviornment.

        Terminates all background processes, and removes files uploaded to
        remote nodes.

        Arguments:
            None

        Returns:
            None

        Raises:
            None
        """
        for process in self.processes:
            process.terminate()

        if self._remote:
            for node in [self.server_cb, self.server_pov, self.server_ids]:
                try:
                    self.remote_cmd(node, ['rm', '-rf', self._tmp_dir])
                except subprocess.CalledProcessError:
                    pass

    @staticmethod
    def random_ip():
        """ Generate a random IP for local testing.

        Arguments:
            None

        Returns:
            A random IP in the 127/8 range

        Raises:
            None
        """
        return "127.%d.%d.%d" % (random.randint(1, 254),
                                 random.randint(1, 254),
                                 random.randint(1, 254))

    def verify_cb(self):
        """ Verify each CB with cgcef_verify

        Arguments:
            None

        Returns:
            0 if all of the CBs are valid DECREE executables
            -1 if any of the CBs is not a valid DECREE executable

        Raises:
            None
        """
        for cb_name in self.cb_list:
            cb_path = os.path.join(self.directory, cb_name)
            try:
                subprocess.check_output(['cgcef_verify', cb_path])
            except subprocess.CalledProcessError as err:
                logging.error('not ok - CB did not verify: %s', str(err), extra={'raw':True})
                return -1
        return 0

    @staticmethod
    def random_port():
        """ Generate a random port for testing.

        Arguments:
            None

        Returns:
            A random port between 2000 and 2999

        Raises:
            None
        """
        return random.randint(2000, 2999)

    @staticmethod
    def log_packages():
        """ Logs all of the installed CGC packages

        Arguments:
            None

        Returns:
            None

        Raises:
            None
        """
        if 'debian' in platform.dist()[0]:
            import apt
            cache = apt.Cache()
            cgc_packages = []
            for package_name in cache.keys():
                if 'cgc' not in package_name and 'cb-testing' not in package_name:
                    continue

                package = cache[package_name]

                status = '%s: (installed: %s)' % (package.name,
                                                  repr(package.is_installed))
                for version in package.versions:
                    status += ' version: %s' % version.version

                cgc_packages.append(status)

            cgc_packages.sort()
            for package in cgc_packages:
                logging.warning('package: %s', package)

    @staticmethod
    def signal_name(sig_id):
        """ Translate a signal number to its name

        Arguments:
            sig_id: signal number

        Returns:
            The signal name, or 'UNKNOWN' if it is an undefined signal number

        Raises:
            None
        """
        for name, value in signal.__dict__.iteritems():
            if sig_id == value:
                return name
        return 'UNKNOWN'

    def start_server(self, connection_count):
        """ Start cb-server in the background

        Arguments:
            connection_count: The maximum number of connections that cb-server
                should handle before exiting

        Returns:
            A handle to the Background() instance of cb-server

        Raises:
            None
        """
        assert connection_count > 0

        server_cmd = ['cb-server', '--insecure', '-p', "%d" % self.port,
                     '-m', "%d" % connection_count]

        if self._remote:
            cb_paths = [os.path.join(self.directory, x) for x in self.cb_list]
            self.remote_upload(self.server_cb, cb_paths)

            server_cmd = ['ssh', self.server_cb] + server_cmd
            server_cmd += ['-d', self._tmp_dir]
        else:
            server_cmd += ['-d', self.directory]

        if self.negotiate_seed:
            server_cmd += ['--negotiate']

        if self.cb_no_attach:
            server_cmd += ['--no_attach']

        if self.cb_seed is not None and self.negotiate_seed is None:
            server_cmd += ['-s', self.cb_seed]
        
        if self.cb_seed_skip is not None:
            server_cmd += ['-S', '%d' % self.cb_seed_skip]

        if self.cb_env is not None:
            server_cmd += ['-e', self.cb_env]

        if self.timeout > 0:
            server_cmd += ['-t', '%d' % self.timeout]
            
        if self.max_send is not None and self.max_send > 0:
            server_cmd += ['-M', '%d' % self.max_send]

        if self.should_debug:
            server_cmd += ['--debug']
        else:
            server_cmd += ['-c', '0']

        if self.wrapper:
            server_cmd += ['-w', self.wrapper]

        server_cmd += self.cb_list

        return self.background(server_cmd, 'cb-server')

    def start_ids(self, connection_count):
        """ Start cb-proxy in the background

        Arguments:
            None

        Returns:
            A handle to the Background() instance of cb-proxy

        Raises:
            None
        """
        assert connection_count > 0

        # the IDS is only used for remote commands
        proxy_cmd = ['sleep', '100']

        # if self._remote and self.ids_rules:
        if self._remote:
            if self.ids_rules:
                self.remote_upload(self.server_ids, [self.ids_rules])
                base_filename = os.path.basename(self.ids_rules)
                remote_path = os.path.join(self._tmp_dir, base_filename)

            proxy_cmd = ['ssh', self.server_ids, 'cb-proxy',
                        # '--listen_host', self.server_ids,
                         '--max_connections', '%d' % connection_count,
                         '--listen_port', '%d' % self.port,
                         '--host', self.server_cb,
                         '--port', '%d' % self.port]

            if self.ids_rules:
                proxy_cmd += ['--rules', remote_path]

            if self.pcap:
                proxy_cmd += ['--pcap_host', socket.gethostname()]

            if self.should_debug:
                proxy_cmd += ['--debug']
        
            if self.negotiate_seed:
                proxy_cmd += ['--negotiate']

        return self.background(proxy_cmd, 'cb-proxy')

    def start_pcap(self):
        """ Start pcap logging in the background

        Arguments:
            None

        Returns:
            If pcap logging is enabled, a handle to the Background() instance
                of the pcap logger
            Otherwise, None

        Raises:
            None
        """
        # XXX add 'remote' support
        if self.pcap:
            if self._remote:
                pcap = self.background(['cb-packet-log', '--pcap_file',
                                        self.pcap])
            else:
                pcap = self.background(['/usr/sbin/tcpdump', '-i', 'lo', '-w',
                                        self.pcap, '-s', '0', '-B', '65536',
                                        'port', '%d' % self.port, 'and',
                                        'host', self.ip_address], 'tcpdump')
            while not pcap.log_queue.qsize():
                time.sleep(0.1)
        return None

    def start_replay(self, xml):
        """ Start cb-replay

        Arguments:
            None

        Returns:
            The handle to the running process

        Raises:
            None
        """

        replay_bin = 'cb-replay'
        if xml[0].endswith('.pov'):
            replay_bin = 'cb-replay-pov'

        replay_cmd = [replay_bin, '--host', self.ip_address, '--port', '%d' %
                      self.port]

        if self.timeout > 0:
            replay_cmd += ['--timeout', '%d' % self.timeout]

        if self.negotiate_seed:
            replay_cmd += ['--negotiate']
            if self.cb_seed is not None:
                replay_cmd += ['--cb_seed', self.cb_seed]

        if self.pov_seed and replay_bin == 'cb-replay-pov':
            replay_cmd += ['--pov_seed', self.pov_seed]

        if self.cb_seed_munge and replay_bin != 'cb-replay-pov':
            replay_cmd += ['--munge_seed']

        if self.should_debug:
            replay_cmd += ['--debug']

        if replay_bin == 'cb-replay':
            if self.failure_ok:
                replay_cmd += ['--failure_ok']
    
            if self.concurrent:
                replay_cmd += ['--concurrent', '%d' % self.concurrent]
           
        if self.max_send is not None and self.max_send > 0:
            replay_cmd += ['--max_send', '%d' % self.max_send]

        if self._remote:
            self.remote_upload(self.server_pov, xml)
            xml = [os.path.join(self._tmp_dir, os.path.basename(x)) for x in
                   xml]
            replay_cmd = ['ssh', self.server_pov] + replay_cmd

        replay_cmd += xml

        return self.launch(replay_cmd)

    def use_server_connections(self, count):
        """ Connect to the cb-server a specifed number of times

        If the replay tool fails, the server might run forever waiting to
        process input. Give it the connections it's expecting and hope it exits
        normally.

        Arguments:
            count: number of times to connect to the cb-server

        Returns:
            None

        Raises:
            None
        """
        for _ in range(0, count):
            try:
                sock = socket.create_connection((self.ip_address,
                                                 self.port))
            except socket.error:
                continue
            sock.close()

    def _check_result_cqe(self, sig, ret):
        if sig == signal.SIGALRM:
            if self.failure_ok:
                logging.warning('ok - process timed out', extra={'raw':True})
                ret = 0
            else:
                logging.error('not ok - process timed out', extra={'raw':True})
                ret = -1
        elif sig in Runner.pov_signals:
            if self.should_core:
                logging.warning('ok - process cored as expected: (signal %d: %s)',
                              sig, self.signal_name(sig), extra={'raw':True})
                ret = 0
            else:
                logging.error('not ok - process cored: (signal %d: %s)',
                              sig, self.signal_name(sig), extra={'raw':True})
                ret = -1
        elif self.should_core:
            logging.error('not ok - core expected, no cores identified', extra={'raw':True})
            ret = -1

        return ret

    def _check_result_type1(self, sig, replay_stdout, server_stdout):
        if sig == 0:
            if self.should_core:
                logging.error('not ok - POV type 1 should have cored, and did not.', extra={'raw':True})
                return -1
            else:            
                logging.warning('ok - POV type 1 did not core, as expected', extra={'raw':True})
                return 0
        elif not self.should_core:
            logging.error('not ok - POV type 1 expected to not core, but did. (signal %d: %s)', sig, self.signal_name(sig), extra={'raw':True})
            return -1

        match = re.match(r"(?ms).*\n# type 1 masks: ([a-fA-F0-9]+) ([a-fA-F0-9]+)", replay_stdout)
        if not match:
            logging.error('not ok - POV type 1 negotiation failed', extra={'raw':True})
            return -1
        ip_mask = int(match.group(1), 16)
        reg_mask = int(match.group(2), 16)
            
        match = re.match(r"(?ms).*\n# type 1 pov: ([a-fA-F0-9]+) ([a-fA-F0-9]+) (\d)", replay_stdout)
        if not match:
            logging.error('not ok - POV type 1 negotiation failed', extra={'raw':True})
            return -1

        ip_value = int(match.group(1), 16)
        reg_value = int(match.group(2), 16)
        reg_num = int(match.group(3))
        
        logging.warning('POV type 1 negotiated masks: %08x %08x %d', ip_mask, reg_mask, reg_num)
        # logging.warning('process cored as expected (signal %d: %s)', sig, self.signal_name(sig))

        register_state_re = r"(?ms).*\nregister states"
        register_state_re2 = r"register states - eax: ([a-fA-F0-9]+) ecx: ([a-fA-F0-9]+) edx: ([a-fA-F0-9]+) ebx: ([a-fA-F0-9]+) esp: ([a-fA-F0-9]+) ebp: ([a-fA-F0-9]+) esi: ([a-fA-F0-9]+) edi: ([a-fA-F0-9]+) eip: ([a-fA-F0-9]+)"

        match = re.match(register_state_re, server_stdout)
        if not match:
            logging.error('not ok - cb-server did not log register states', extra={'raw':True})
            return -1

        lines = server_stdout.split('\n')

        expected_eip = ip_value & ip_mask
        expected_reg = reg_value & reg_mask

        results = []
        for line in lines:
            match = re.match(register_state_re2, line)
            if not match:
                continue

            registers = []
            for i in range(8):
                registers.append(int(match.group(1+i), 16))
            eip = int(match.group(9), 16)

            result_eip = eip & ip_mask
            result_reg = registers[reg_num] & reg_mask

            if (result_eip == expected_eip and result_reg == expected_reg):
                logging.warning('ok - TYPE 1 POV', extra={'raw':True})
                return 0

            results.append([result_eip, result_reg])

        for values in results:
            logging.error('not ok - TYPE 1 POV FAILED.  expected eip: %08x actual eip: %08x expected reg: %08x actual reg: %08x', expected_eip, values[0], expected_reg, values[1], extra={'raw': True})
        
        return -1

    def _check_result_type2(self, replay_stdout, server_stdout):
        seed = self.cb_seed
        if not seed:
            match = re.match(r"(?ms).*\nseed: ([a-fA-F0-9]+)", server_stdout)
            if not match:
                logging.error('not ok - POV type 2 server did not print seed', extra={'raw':True})
                return -1

            seed = match.group(1)
        
        prng = ansi_x931_aes128.PRNG(seed.decode('hex'))
        flag_page = prng.get(0x1000)

        match = re.match(r"(?ms).*\n# secret value: ([a-fA-F0-9]+)", replay_stdout)
        if not match:
            if self.should_core:
                logging.error('not ok - POV type 2 did not return a secret value', extra={'raw':True})
                return -1
            else:
                logging.warning('ok - POV type 2 did not return a secret value as expected', extra={'raw':True})
                return 0

        value = match.group(1).decode('hex')
        if value not in flag_page:
            if not self.should_core:
                logging.warning('ok - POV type 2 secret value is not part of the secret page (as expected)', extra={'raw':True})
                return 0
            else:
                logging.error('not ok - TYPE 2 POV FAILED.  %s not in %s', value.encode('hex'), flag_page.encode('hex'), extra={'raw':True})
                return -1

        if not self.should_core:
            logging.error('not ok - TYPE 2 unexpectedly passed.  %s in %s', value.encode('hex'), flag_page.encode('hex'), extra={'raw':True})
            return -1

        logging.warning('ok - TYPE 2 POV', extra={'raw': True})
        return 0

    def _check_result_cfe(self, sig, replay_stdout, server_stdout):
        pov_type = None
        match = re.match(r"(?ms).*\n# negotiation type: (\d)\n", replay_stdout)
        if not match:
            logging.error("not ok - pov did not negotiate", extra={'raw':True})
            return -1

        pov_type = int(match.group(1))
        if pov_type not in [1, 2]:
            logging.error("not ok - pov negotiation is not 1 or 2: %d\n" % pov_type, extra={'raw':True})
            return -1

        if pov_type == 1:
            return self._check_result_type1(sig, replay_stdout, server_stdout)

        if pov_type == 2:
            return self._check_result_type2(replay_stdout, server_stdout)

    def check_result(self, replay_result, server_result, xml_list):
        """ Check the results of cb-server and cb-replay

        Arguments:
            replay_result: the output of cb-replay
            sig: the return code of cb-server

        Returns:
            Returns 0 if the test 'passed'
            Returns -1 if the test 'failed'

        Raises:
            None
        """
        ret, replay_stdout = replay_result
        sig, server_stdout = server_result

        magic_value = None

        if not xml_list[0].endswith('.pov'):
            return self._check_result_cqe(sig, ret)
        else:
            return self._check_result_cfe(sig, replay_stdout, server_stdout)

    def run(self):
        """ Runs the test

        Arguments:
            None

        Returns:
            Returns 0 if the tests all pass
            Returns non-0 if any of the tests failed

        Raises:
            None
        """
        self.log_packages()

        if self.verify_cb() != 0:
            return -1

        xml_sets = []
        if self.failure_ok:
            for xml in self.xml_list:
                xml_sets.append([xml])
        else:
            cqe_xml = []
            cfe_xml = []
            for xml in self.xml_list:
                if xml.endswith('.xml'):
                    cqe_xml.append(xml)
                else:
                    cfe_xml.append(xml)

            for i in [cqe_xml, cfe_xml]:
                if len(i):
                    xml_sets.append(i)

        pcap = self.start_pcap()
        for xml_list in xml_sets:
            server = self.start_server(len(xml_list))
            ids = self.start_ids(len(xml_list))
            replay_result = self.start_replay(xml_list)
            
            if server.process.poll() is None:
                self.use_server_connections(len(xml_list))

            server_result = None, ''

            # wait a maximum of 30 seconds after our replay and socket cleanup
            # finishes before terminating cb-server if it hasn't returned yet.
            try:
                with Timeout(30):
                    server_result = server.wait()
            except TimeoutException:
                server.terminate()

            ids.terminate()

            result = self.check_result(replay_result, server_result, xml_list)
            if result != 0:
                return result

        if pcap is not None:
            pcap.terminate()

        return 0

class CB_Formatter(logging.Formatter):

    def format(self, record):
        s = super(CB_Formatter, self).format(record)
        result = []

        for line in s.split('\n'):
            if 'raw' in dir(record) and record.raw == True:
                if line.startswith('# '):
                    line = line[2:]
            else:
                if line[:2] != '# ':
                    line = '# ' + line
            result.append(line)
        return '\n'.join(result)

def main():
    """ Test a challenge binary using cb-server and cb-replay, with either
    local or remote support.

    Arguments:
        None

    Returns:
        Returns 0 on test success
        Returns non-0 if any of the tests failed

    Raises:
        None
        """

    formatter = argparse.ArgumentDefaultsHelpFormatter
    parser = argparse.ArgumentParser(description='Send CGC Polls and POVs',
                                     formatter_class=formatter)

    required = parser.add_argument_group(title='required')
    required.add_argument('--cb', required=True, type=str, nargs='+',
                          help='Challenge Binaries to run')
    required.add_argument('--directory', required=True, type=str,
                          help='Directory containing challenge binaries')

    remote = parser.add_argument_group(title='Remote Testing')
    remote.add_argument('--enable_remote', required=False, action='store_true',
                        default=False, help='Perform tests remotely')
    remote.add_argument('--remote_nodes', required=False, nargs=3, type=str,
                        default=['cb', 'pov', 'ids'],
                        metavar=('<cb-server>', '<pov-server>',
                                 '<ids-server>'),
                        help='Hosts to run each of the components')

    parser.add_argument('--debug', required=False, action='store_true',
                        default=False, help='Enable debugging')
    parser.add_argument('--log', required=False, type=str, help='Log filename')
    parser.add_argument('--pcap', required=False, type=str,
                        help='PCAP output file')
    parser.add_argument('--concurrent', required=False, type=int,
                        help='Number of Polls/POVs to throw concurrently')
    parser.add_argument('--cb_seed', required=False, type=str,
                        help='CB PRNG Seed')
    parser.add_argument('--cb_seed_skip', required=False, type=int,
                        help='Advance the CB PRNG Seed before launch')
    parser.add_argument('--cb_seed_munge', required=False, action='store_true',
                        default=False, help='Bit invert the seed for testing')
    parser.add_argument('--port', required=False, type=int,
                        help='PORT to run the CB')
    parser.add_argument('--pov_seed', required=False, type=str,
                        help='POV PRNG Seed')
    parser.add_argument('--timeout', required=False, type=int, default=15,
                        help='Maximum duration for each Poll or POV')
    parser.add_argument('--should_core', required=False, action='store_true',
                        default=False, help='This test should cause a core')
    parser.add_argument('--wrapper', required=False, type=str,
                        help='Executable to wrap each CB for instrumentation')
    parser.add_argument('--failure_ok', required=False, action='store_true',
                        default=False, help='Failures for this test are '
                        'accepted')
    parser.add_argument('--max_send', required=False, type=int,
                        help='Maximum data to send in each request from the '
                        'poll')
    parser.add_argument('--ids_rules', required=False, type=str,
                        help='IDS Rules')
    parser.add_argument('--negotiate_seed', required=False, action='store_true',
                        default=False, help='Negotate the CB seed from cb-replay')
    parser.add_argument('--cb_no_attach', required=False, action='store_true',
                        default=False, help='Do not attach to the CB')
    parser.add_argument('--cb_env', required=False, type=str,
                        help='CB Enviornment Variable')

    exgroup = parser.add_argument_group(title='XML files')
    group = exgroup.add_mutually_exclusive_group(required=True)
    group.add_argument('--xml', type=str, nargs='+', help='POV/POLL XML File '
                       'to evaluate')
    group.add_argument('--xml_dir', type=str, help='Directory containing '
                       'POV/POLL XML Files to evaluate')

    args = parser.parse_args()

    xml_files = args.xml
    if args.xml is None:
        assert os.path.isdir(args.xml_dir), "not a directory: " + args.xml_dir
        xml_files = glob.glob("%s/*.xml" % args.xml_dir)
        xml_files += glob.glob("%s/*.pov" % args.xml_dir)
        xml_files.sort()

    assert len(xml_files), "No Polls or POVs"

    if args.cb_seed_skip is not None:
        assert args.cb_seed_skip > 0, "Invalid cb_seed_skip: %d" % args.cb_seed_skip

    for filename in xml_files:
        assert filename.endswith('.pov') or filename.endswith('.xml'), 'invalid POV/Poll'

    if args.cb_seed is not None:
        assert len(args.cb_seed) == 96, "--cb_seed must be 96 bytes"
    
    if args.pov_seed is not None:
        assert len(args.pov_seed) == 96, "--pov_seed must be 96 bytes"

    logger = logging.getLogger()
    logger.setLevel(logging.WARNING)
    if args.debug:
        logger.setLevel(logging.DEBUG)

    if args.log is not None:
        log_fh = open(args.log, 'w')
    else:
        log_fh = sys.stdout

    log_handler = logging.StreamHandler(log_fh)
    log_handler.setFormatter(CB_Formatter('# %(message)s'))

    logger.addHandler(log_handler)

    if log_fh != sys.stdout:
        error_handler = logging.StreamHandler(sys.stdout)
        error_handler.setLevel(logging.ERROR)
        error_handler.setFormatter(logging.Formatter('# %(message)s'))
        logger.addHandler(error_handler)

    ret = -1
    runner = Runner(args.port, args.cb, xml_files, args.pcap, args.wrapper,
                    args.directory, args.should_core, args.failure_ok,
                    args.debug, args.timeout, log_fh, args.cb_seed,
                    args.cb_seed_skip, args.max_send, args.concurrent,
                    args.negotiate_seed, args.pov_seed, args.cb_no_attach, 
                    args.cb_seed_munge, args.cb_env)

    if args.ids_rules:
        runner.set_ids_rules(args.ids_rules)

    if args.enable_remote:
        runner.enable_remote(args.remote_nodes[0], args.remote_nodes[1],
                             args.remote_nodes[2])

    try:
        ret = runner.run()
    except KeyboardInterrupt:
        logger.warning('interrupted')
    finally:
        runner.cleanup()
    return ret

if __name__ == "__main__":
    exit(main())