scrap-is-not-scrap/run.py

#! /usr/bin/env python3
# encoding: utf-8
import logging
import os
import shutil
import subprocess
import binascii
import random


class Sins():
    logger = logging.getLogger()

    def __init__(self, parent, seed=None, run_dir=None):
        self.parent = os.path.abspath(parent)
        if not os.path.isfile(parent):
            raise ValueError('Invalid executable image path.')

        self.seed = os.path.abspath(seed)

        self.run_dir = os.path.abspath(run_dir)

        if not os.path.isdir(self.run_dir):
            try:
                os.mkdir(self.run_dir)
            except:
                self.run_dir = os.path.dirname(self.seed)

        shutil.copy2(self.seed, self.run_dir)

        parsed = (self.parent, self.seed, self.run_dir)
        self.logger.info('init path: {}'.format(parsed))

        paths = (self.seed,)
        while True:
            for path in paths:
                scrap_path = os.path.join(self.run_dir, path)
                if os.path.isfile(scrap_path):
                    child = self.generation(
                        self.parent,
                        scrap_path,
                        self.run_dir
                    )
                    if child:
                        print(child)
            # paths = sorted(os.listdir(self.run_dir))

    def scrap_recent(self, run_dir):
        scraps = sorted(os.listdir(run_dir))

        if scraps:
            return os.path.join(run_dir, scraps[-1])

        return None

    def generation(self, parent, scrap, cwd):
        with open(scrap, 'rb') as scrap_file:
            scrap_bin = scrap_file.read()

        scrap_hex = binascii.b2a_hex(scrap_bin).upper()
        offset = random.randint(0, len(scrap_hex))
        flip = random.randint(0, 255)

        self.logger.debug(
            'generation: {}'.format((parent, scrap_hex, offset, flip, cwd))
        )

        proc = subprocess.run(
            [parent, scrap_hex, str(offset), str(flip)],
            cwd=cwd,
            stdout=subprocess.PIPE
        )

        if proc.stdout:
            child_hex = binascii.a2b_hex(proc.stdout)
            return child_hex


def hex_dumps(scrap_dir):
    scrap_dir = os.path.abspath(scrap_dir)
    dump_dir = os.path.join(scrap_dir, 'hex')

    if not os.path.isdir(dump_dir):
        os.mkdir(dump_dir)

    scraps = os.listdir(scrap_dir)

    for scrap in scraps:
        scrap_path = os.path.join(scrap_dir, scrap)
        hex_name = '{}.hex'.format(scrap)
        hex_path = os.path.join(dump_dir, hex_name)

        if os.path.isfile(scrap_path):
            with open(scrap_path, 'rb') as file_in:
                with open(hex_path, 'w') as file_out:
                    binary = bytes(file_in.read())
                    for each in binary:
                        file_out.writelines('\'\\x{:02X}\',\n'.format(each))


if __name__ == '__main__':
    import argparse
    import sys

    parser = argparse.ArgumentParser(
        description='position independent code (PIC) mutation experiment.'
    )
    parser.add_argument('--verbose', '-v', action='count')
    parser.add_argument('-build', action='store_true',
                        help='build parent and seed PIC image, exit.')
    parser.add_argument('-provision', action='store_true',
                        help='provision ubuntu for run, exit.')
    parser.add_argument('-logfile', help='log to file.')
    parser.add_argument('-seed', default='build/seed.asm.2.o',
                        help='path to PIC image.')
    parser.add_argument('-parent', default='build/generation',
                        help='path to parent process.')
    parser.add_argument('-dir', default='sandbox',
                        help='path to execution directory.')
    parser.add_argument('-dumps', action='store_true',
                        help='dump hex values of scraps in directory, exit.')
    args = parser.parse_args()

    logger = logging.getLogger()
    formatter = logging.Formatter(
        '# %(asctime)s %(levelname)s\n%(message)s\n'
    )

    if args.verbose:
        logger.setLevel(logging.DEBUG)
    else:
        print('Verbose not set, running in silence.')

    if args.logfile:
        file_handler = logging.FileHandler(args.logfile, 'a')
        file_handler.setFormatter(formatter)
        logger.setLevel(logging.DEBUG)
        logger.addHandler(file_handler)
    else:
        stream_handler = logging.StreamHandler()
        stream_handler.setFormatter(formatter)
        logger.setLevel(logging.INFO)
        logger.addHandler(stream_handler)

    logger.info('run: {}'.format(args))

    if args.provision:
        provision = ['sudo', 'sh', 'provision-ubuntu.sh']
        prov_proc = subprocess.run(provision)

        os.environ['CC'] = 'clang'
        waf = [sys.executable, 'waf.py', 'configure']
        waf_proc = subprocess.run(waf)

    elif args.build:
        waf = [sys.executable, 'waf.py', 'build']
        waf_proc = subprocess.run(waf)
        logger.info(waf_proc.stdout)

    elif args.dumps:
        hex_dumps(args.dir)
    else:
        sins = Sins(parent=args.parent, seed=args.seed, run_dir=args.dir)