#! /usr/bin/env python3
from argparse import ArgumentParser
from datetime import datetime
from pathlib import Path
from random import randint
from multiprocessing import Process, Queue
from queue import Empty
import binascii
import ctypes
import logging
import mmap

template_shell = b''.join([
    b'\x55',              # push   rbp
    b'\x48\x89\xe5',      # mov    rbp,rsp
    b'\x48\x89\x7d\xf8',  # mov    QWORD [rbp-0x8],rdi
    b'\x48\x8b\x45\xf8',  # mov    rax,QWORD [rbp-0x8]
    b'\x5d',              # pop    rbp
    b'\xc3'])             # ret

seed_shell = b''.join([
    b'\x55',
    b'\x48\x89\xe5',
    b'\x90' * randint(8, 64),
    b'\x48\x89\x7d\xf8',
    b'\x90' * randint(8, 64),
    b'\x48\x8b\x45\xf8',
    b'\x5d',
    b'\xc3'])


def flip(shellcode: bytes):
    shellcode = bytearray(shellcode)
    offset = randint(0, len(shellcode) -1)
    flip = randint(0, 255)
    shellcode[offset] ^= flip
    return bytes(shellcode)


def generation(queue: Queue, shellcode: bytes):
    prot = mmap.PROT_READ | mmap.PROT_WRITE | mmap.PROT_EXEC
    flags = mmap.MAP_ANONYMOUS | mmap.MAP_PRIVATE
    exec_mem = mmap.mmap(-1, len(shellcode), prot=prot, flags=flags)

    exec_mem.write(shellcode)
    ctypes_buffer = ctypes.c_int.from_buffer(exec_mem)
    addr = ctypes.addressof(ctypes_buffer)

    function = ctypes.CFUNCTYPE(ctypes.c_uint, ctypes.c_uint)(addr)
    function._avoid_gc_for_mmap = exec_mem

    shellcode_len = ctypes.c_uint(len(shellcode))
    result = function(shellcode_len)

    queue.put(result)


def sins():
    now = '{0:%Y%m%dT%H%M%S}'.format(datetime.utcnow())
    parser = ArgumentParser(
        description='position independent code (PIC) mutation experiment.')
    parser.add_argument('-v', '--verbose', action='count')
    parser.add_argument('-s', '--seed', help='path to PIC image.')
    parser.add_argument('-o', '--output', help='path to results directory.')
    parser.add_argument('-l', '--lineage', default=10,
        help='max count of unsuccessful generation.')
    args = parser.parse_args()

    log_level = logging.INFO
    log_format = logging.Formatter('%(message)s')

    if args.verbose:
        log_level = logging.DEBUG
        log_format = logging.Formatter(
            '%(levelname)s %(filename)s:%(lineno)d\n%(message)s\n')

    logger = logging.getLogger('sins')
    logger.setLevel(log_level)

    stream_handler = logging.StreamHandler()
    stream_handler.setLevel(log_level)
    stream_handler.setFormatter(log_format)
    logger.addHandler(stream_handler)

    if args.output:
        log_path = f'{args.output}/sins-{now}.log'
        file_handler = logging.FileHandler(log_path)
        file_handler.setLevel(log_level)
        file_handler.setFormatter(log_format)
        logger.addHandler(file_handler)

    logger.info(now)

    seed_data = seed_shell

    if args.seed:
        seed = Path(args.seed)

        with seed.open('rb') as seed_file:
            seed_data = seed_file.read()

    logger.info(f'seed:\n{seed_data}')

    queue = Queue()

    while True:
        lineage = 0
        seed_flipped = flip(seed_data)

        while lineage < args.lineage:
            logger.info(f'lineage: {lineage}')
            result = None

            proc = Process(target=generation, args=(queue, seed_flipped))
            proc.start()
            try:
                result = queue.get(timeout=1)
            except Empty:
                lineage += 1
                continue

            if not result:
                lineage += 1
                continue

            logger.info(f'scrap:\n{seed_flipped}')
            logger.info(f'result: {result}')
            lineage = 0
            seed_flipped = flip(seed_flipped)