#! /usr/bin/env python3
from argparse import ArgumentParser
from datetime import datetime
from pathlib import Path
from random import randint
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 shell_func(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

    return function


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.')
    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()

    seed_data = bytearray(seed_data)

    offset = randint(0, len(seed_data))
    flip = randint(0, 255)
    seed_data[offset] ^= flip

    seed_len = ctypes.c_uint(len(seed_data))
    logger.info(f'seed_data\n{binascii.b2a_hex(seed_data)}')
    logger.info(f'seed_len: {seed_len}')

    seed = shell_func(seed_shell)
    ret_val = seed(seed_len)

    logger.info(f'ret_val:  {ret_val}')