cb-replay-pov

#!/usr/bin/env python

"""
CB POV / Poll communication verification tool

Copyright (C) 2014 - 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 for deterministic communication to a CGC Challenge Binary
using a Challenge Binary as input.

1 - http://testanything.org/
"""

import subprocess
import multiprocessing
import random
import sys
import argparse
import os
import signal
import re
import socket
import struct
import time

class TimeoutException(Exception):
    """ Exception to be used by Timeout(), to allow catching of timeout
    exceptions """
    pass


class TestFailure(Exception):
    """ Exception to be used by Throw(), to allow catching of test failures """
    pass


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

    def __init__(self, seconds):
        self.seconds = seconds

    @staticmethod
    def cb_handle_timeout(signum, frame):
        """ SIGALRM signal handler callback """
        raise TimeoutException("timed out")

    def __enter__(self):
        if self.seconds > 0:
            signal.signal(signal.SIGALRM, self.cb_handle_timeout)
            signal.alarm(self.seconds)

    def __exit__(self, exit_type, exit_value, traceback):
        if self.seconds:
            signal.alarm(0)


def ptrace_traceme():
    from ctypes import cdll
    from ctypes.util import find_library
    from ctypes import c_long, c_ulong

    LIBC_FILENAME = find_library('c')
    libc = cdll.LoadLibrary(LIBC_FILENAME)

    _ptrace = libc.ptrace
    _ptrace.argtypes = (c_ulong, c_ulong, c_ulong, c_ulong)
    _ptrace.restype = c_ulong

    PTRACE_TRACEME = 0

    result = _ptrace(PTRACE_TRACEME, 0, 0, 0)
    result_signed = c_long(result).value
    return result_signed

def launch_gdb_proxy(pid, attach_port):
    gdb_pid = os.fork()
    if gdb_pid == 0:
        subprocess.call(['/usr/bin/gdbserver', ':%d' % attach_port, '--attach', '%d' % pid])
        exit(0)
    else:
        # ugh.  
        time.sleep(2)
        result = os.waitpid(gdb_pid, os.WNOHANG)
        if result != (0, 0):
            print "Unable to attach to the process"
            return 0

    return gdb_pid

class Throw(object):
    """Throw - Perform the interactions with a CB

    This class implements the basic methods to interact with a CB, verifying
    the interaction works as expected.

    Usage:
        a = Throw((source_ip, source_port), (target_ip, target_port), POV,
                  timeout, should_debug, negotiate, cb_seed, attach_port)
        a.run()

    Attributes:
        source: touple of host and port for the outbound connection
        target: touple of host and port for the CB

        count: Number of actions performed

        debug: Is debugging enabled

        failed: Number of actions that did not work as expected

        passed: Number of actions that did worked as expected

        pov: POV, as defined by POV()

        sock: TCP Socket to the CB

        timeout: connection timeout

        values: Variable dictionary

        logs: all of the output from the interactions

        negotiate: Should the PRNG be negotiated with the CB

    """
    def __init__(self, source, target, pov, timeout, debug, negotiate, cb_seed, attach_port, max_send, pov_seed):
        self.times = 10
        self.source = source
        self.target = target
        self.count = 0
        self.failed = 0
        self.passed = 0
        self.pov = pov
        self.debug = debug
        self.sock = None
        self.timeout = timeout
        self.negotiate_fd_fd = negotiate
        self.negotiate = negotiate
        self.cb_seed = cb_seed
        self.logs = []
        self.attach_port = attach_port
        self.max_send = max_send
        self.pov_seed = pov_seed

        if self.cb_seed is None:
            self.cb_seed = os.urandom(48)
        else:
            self.cb_seed = self.cb_seed.decode('hex')

    def setup_negotiation(self):
        if not self.is_pov():
            return
        self.negotiate_fd = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.negotiate_fd.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.negotiate_fd.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack('ii', 1, 5))
        self.negotiate_fd.bind(('',0))
        self.negotiate_fd.listen(self.times)
        negotiate_address = self.negotiate_fd.getsockname()

        if self.debug:
            self.log("negotiation listen at: %s" % repr(negotiate_address))

        self.negotiate_fd_client = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.negotiate_fd_client.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.negotiate_fd_client.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER, struct.pack('ii', 1, 5))
        self.negotiate_fd_client.connect(negotiate_address)

    def log(self, data):
        print "# %s" % data

    def count_bits_set(self, value):
        count = 0
        while value:
            value &= value - 1
            count += 1
        return count

    def recv_all(self, sock, data_len):
        data = ''
        while len(data) < data_len:
            read = sock.recv(data_len - len(data))
            assert len(read)
            data += read
        return data

    def pov_negotiate(self):
        MIN_BITS = 20
        if not self.is_pov():
            return
        
        if self.debug:         
            self.log("negotiating")
        connection, address = self.negotiate_fd.accept()

        # as soon as we've accepted one, drop this server.  
        self.negotiate_fd.close()

        data = self.recv_all(connection, 4)
        pov_type = struct.unpack('<L', data)[0]
        assert pov_type in [1, 2], 'invalid pov type'

        self.log("negotiation type: %d" % pov_type)

        if pov_type == 1:
            data = self.recv_all(connection, 12)
            ipmask, regmask, regnum = struct.unpack('<LLL', data)
            accept = True

            if self.count_bits_set(ipmask) < MIN_BITS:
                self.log('not enough bits in IP mask')
                accept = False

            if self.count_bits_set(regmask) < MIN_BITS:
                self.log('not enough bits in REG mask')
                accept = False

            if regnum > 7:
                self.log('invalid GPR')
                accept = False

            if accept:
                ip = random.randint(0, 0xFFFFFFFF) & ipmask
                reg = random.randint(0, 0xFFFFFFFF) & regmask

                self.log('type 1 masks: %08x %08x' % (ipmask, regmask))
                self.log('type 1 pov: %08x %08x %d' % (ip, reg, regnum))

                pov_values = struct.pack('<LL', ip, reg)
                connection.sendall(pov_values)
 
        if pov_type == 2:
            PAGE_ADDR = 0x4347C000
            PAGE_LENGTH = 0x1000
            PAGE_BYTES = 4

            self.log("sending page location: %d, %d, %d" % (PAGE_ADDR, PAGE_LENGTH, PAGE_BYTES))
            pov_location = struct.pack('<LLL', PAGE_ADDR, PAGE_LENGTH, PAGE_BYTES)
            connection.sendall(pov_location)

            if self.debug:
                self.log("getting secret")
            data = self.recv_all(connection, 4)
            self.log("secret value: %s" % data.encode('hex'))

        if self.debug:
            self.log('done')
        connection.close()

    def is_pov(self):
        # XXX - if we expose CB based pollers, change this
        return True

    def _encode(self, records):
        """
            record is a list of records in the format (type, data)

            Current wire format:
            RECORD_COUNT (DWORD)
                record_0_type (DWORD)
                record_0_len (DWORD)
                record_0_data (record_0_len bytes)
                record_N_type (DWORD)
                record_N_len (DWORD)
                record_N_data (record_N_len bytes)
        """

        packed = []
        for record_type, data in records:
            packed.append(struct.pack('<LL', record_type, len(data)) + data)

        result = struct.pack('<L', len(packed)) + ''.join(packed)
        return result

    def _read_len(self, size):
        total_size = 0
        result = ''
        while total_size < size:
            data = self.sock.recv(size - total_size)
            assert len(data), "not enough data returned from cb-server"
            total_size += len(data)
            result += data

        return result

    def _send_all(self, data, max_send=None):
        total_sent = 0
        while total_sent < len(data):
            if max_send is not None:
                sent = self.sock.send(data[total_sent:total_sent+max_send])
                # allow the kernel a chance to forward the data
                time.sleep(0.00001)
            else:
                sent = self.sock.send(data[total_sent:])
            if sent == 0:
                return total_sent
            total_sent += sent

        return total_sent

    def cb_negotiate(self):
        """ Prior to starting the POV comms, setup the seeds with the CB server

        Args:
            None

        Returns:
            None

        Raises:
            None
        """

        if not self.negotiate:
            return 0

        request_seed = (1, self.cb_seed)
        self.log('using seed: %s' % self.cb_seed.encode('hex'))
        request = [request_seed]
        encoded = self._encode(request)
        sent = self._send_all(encoded)
        if sent != len(encoded):
            self.log_fail('negotiate failed.  expected to send %d, sent %d' % (len(encoded), sent))
            return -1

        response_packed = self._read_len(4)
        response = struct.unpack('<L', response_packed)[0]
        if response != 1:
            return -1

        return 0

    def run(self):
        """ Iteratively execute each of the actions within the POV

        Args:
            None

        Returns:
            None

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

        self.log('%s' % (self.pov))

        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
        self.sock.setsockopt(socket.SOL_SOCKET, socket.SO_LINGER,
                             struct.pack('ii', 1, 5))
        self.sock.bind(self.source)
        self.sock.connect(self.target)

        if self.debug:
            self.log('connected to %s' % repr(self.target))

        if self.is_pov():
            self.setup_negotiation()

        # handle PRNG negotiation with cb-server
        self.cb_negotiate()

        # get a socket pair

        queue = multiprocessing.Queue()
        gdb_pid = None

        pid = os.fork()
        if pid == 0:
            if self.timeout > 0 and not self.attach_port:
                signal.alarm(self.timeout)

            os.dup2(self.sock.fileno(), sys.stdin.fileno())
            os.dup2(self.sock.fileno(), sys.stdout.fileno())

            if not self.debug:
                null = os.open('/dev/null', 0)
                os.dup2(null, 2)
                os.close(null)

            if self.is_pov():
                os.dup2(self.negotiate_fd_client.fileno(), 3)

            if self.attach_port:
                ptrace_traceme()

            args = [self.pov]

            if self.max_send > 0:
                args.append('max_transmit=%d' % self.max_send)
                args.append('max_receive=%d' % self.max_send)

            if self.pov_seed:
                args.append('seed=%s' % self.pov_seed)

            queue.get(1)
            os.execv(self.pov, args)
            exit(-1)
        else:
            if self.attach_port:
                gdb_pid = launch_gdb_proxy(pid, self.attach_port)

            queue.put(1)

            if self.timeout > 0 and not self.attach_port:
                with Timeout(self.timeout + 5):
                    self.pov_negotiate()
            else:
                self.pov_negotiate()

        if self.debug:
            self.log('waiting')

        return os.waitpid(pid, 0)

def init_worker():
    signal.signal(signal.SIGINT, signal.SIG_IGN)

def run_pov(src, dst, pov, timeout, debug, negotiate, cb_seed, attach, max_send, pov_seed):
    """
    Parse and Throw a POV/Poll

    Arguments:
        src: IP/Port tuple for the source of the connection
        dst: IP/Port tuple for the destination of the connection
        pov: filename of the POV
        timeout: How long the POV communication is allowed to take
        debug: Flag to enable debug logs
        negotate: Should PRNG be negotiated with the CB
        cb_seed: seed to use in the CB
        attach: should the POV be run under gdbserver
        max_send: maximum amount of transmit/receive 
        pov_seed: the POV seed to use

    Returns:
        The number of passed tests
        The number of failed tests
        A list containing the logs

    Raises:
        Exception if parsing the POV times out
    """

    thrower = Throw(src, dst, pov, timeout, debug, negotiate, cb_seed, attach,
                    max_send, pov_seed)
    return thrower.run()

def main():
    """ Parse and Throw the POVs """
    parser = argparse.ArgumentParser(description='Send CB based CGC Polls and POVs')
    required = parser.add_argument_group(title='required arguments')
    required.add_argument('--host', required=True, type=str,
                          help='IP address of CB server')
    required.add_argument('--port', required=True, type=int,
                          help='PORT of the listening CB')
    required.add_argument('files', metavar='pov', type=str, nargs='+',
                          help='pov file')
    parser.add_argument('--source_host', required=False, type=str, default='',
                        help='Source IP address to use in connections')
    parser.add_argument('--source_port', required=False, type=int,
                        default=0, help='Source port to use in connections')
    parser.add_argument('--timeout', required=False, type=int, default=15,
                        help='Connect timeout')
    parser.add_argument('--max_send', required=False, type=int, default=0,
                        help='Maximum amount of data to send and receive at once')
    parser.add_argument('--debug', required=False, action='store_true',
                        default=False, help='Enable debugging output')
    parser.add_argument('--negotiate', required=False, action='store_true',
                        default=False, help='The CB seed should be negotiated')
    parser.add_argument('--cb_seed', required=False, type=str,
                        help='Specify the CB Seed')
    parser.add_argument('--pov_seed', required=False, type=str,
                        help='Specify the POV Seed')
    parser.add_argument('--attach_port', required=False, type=int, 
                        help='Attach with gdbserver prior to launching the '
                        'POV on the specified port')

    args = parser.parse_args()

    if args.cb_seed is not None and not args.negotiate:
        raise Exception('CB Seeds can only be set with seed negotiation')

    assert len(args.files)
    for filename in args.files:
        assert os.path.isfile(filename), "pov must be a file: %s" % repr(filename)
        assert filename.endswith('.pov'), "%s does not end in .pov" % repr(filename)

    pool_responses = []
    for pov in args.files:
        pid, status = run_pov((args.source_host, args.source_port),
                              (args.host, args.port), pov, args.timeout,
                              args.debug, args.negotiate, args.cb_seed,
                              args.attach_port, args.max_send, args.pov_seed)

    return status != 0


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