scrap-is-not-scrap/run.py

#! /usr/bin/env python3

from datetime import datetime
from random import randint
from time import sleep
import binascii
import logging
import os
import socket
import subprocess


class Server:
    _children = list()

    host = ('127.0.0.1', 5000)
    sock = socket.socket()
    sock.bind(host)
    sock.listen()
    client = None
    logging.info('binding [{}]'.format(host))

    def __enter__(self):
        return self

    def __exit__(self, *args):
        return

    def __del__(self):
        self.sock.shutdown(socket.SHUT_RD)
        self.sock.close()

    def accept(self):
        self.client, address = self.sock.accept()
        logging.info('address accpeted [{}]'.format(address))

    def recv(self, size: int=int()) -> list:
        if not self.client:
            raise BrokenPipeError('client connection missing')

        recvd = self.client.recv(size)

        if recvd:
            logging.info('recvd [{}]'.format(len(recvd)))

        return recvd

    def send(self, data: bytes):
        logging.info('send [{}]'.format(binascii.b2a_hex(data)))

        if not self.client:
            raise BrokenPipeError('client connection missing')

        self.client.sendall(data)

    @property
    def children(self):
        children = self._children
        self._children = list()
        return children


class Genorator:
    seed = bytes()

    def __init__(self, *, server: Server, parent: str, seed: str=str(), run_dir: str=str()):
        self.server = server
        self.parent = os.path.abspath(parent)

        if seed:
            with open(seed, 'rb') as seed_file:
                self.seed = seed_file.read()

        try:
            os.mkdir(run_dir)
        except FileNotFoundError or PermissionError:
            run_dir = os.path.dirname(self.seed)
        except FileExistsError:
            pass

        self.run_dir = run_dir

        parent_args = {
            'parent': self.parent,
            'scrap': self.seed,
            'cwd': self.run_dir}

        self.parenting(parent=self.parent, cwd=self.run_dir)
        self.seeding(scrap=self.seed)

    def __del__(self):
        if self.proc:
            self.proc.terminate()

    def scrap_recent(self, *, run_dir: str) -> str:
        scraps = sorted(os.listdir(run_dir))

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

        return None

    def parenting(self, *, parent: str, cwd: str):
        cmd = [parent]
        logging.info('parent {}'.format(cmd))

        self.proc = subprocess.Popen(
            [parent],
            cwd=cwd)

    def seeding(self, *, scrap: bytes):
        offset = randint(0, len(scrap))
        flip = randint(0, 255)

        logging.info('scrap\n{}'.format(binascii.b2a_hex(scrap)))

        self.server.accept()

        scrap_len = bytes([len(scrap)])
        logging.info('send child len [{}]'.format(scrap_len))

        self.server.send(scrap_len)
        # self.server.send(scrap)

        logging.info('child send')

        child = self.server.recv()
        if child:
            raise Exception(child)


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))


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

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


def build():
    waf = [sys.executable, 'waf.py', 'build']
    waf_proc = subprocess.run(waf)

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('-log', default='log_sins', 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.elf',
                        help='path to generation lib.')
    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()
    logger.setLevel(
        logging.DEBUG) if args.verbose else logger.setLevel(logging.INFO)
    formatter = logging.Formatter('# %(filename)s:%(lineno)s\n%(message)s')

    stream_handler = logging.StreamHandler()
    stream_handler.setFormatter(formatter)
    logger.addHandler(stream_handler)

    file_handler = logging.FileHandler(args.log, 'a')
    file_handler.setFormatter(formatter)
    logger.addHandler(file_handler)

    logging.info(args)

    if args.provision:
        provision()

    if args.build:
        build()

    if args.dumps:
        hex_dumps(args.dir)

    with Server() as server:
        gen = Genorator(server=server, parent=args.parent,
                        seed=args.seed, run_dir=args.dir)