#! /usr/bin/env python3
from multiprocessing import Queue
from random import randint
import ctypes
import mmap

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

# use template to declutter output, todo
seed_shell = b''.join([
    b'\x55',              # push   rbp
    b'\x48\x89\xe5',      # mov    rbp,rsp
    b'\x90' * 8,          # nop
    b'\x48\x89\x7d\xf8',  # mov    QWORD [rbp-0x8],rdi
    b'\x90' * 8,          # nop
    b'\x48\x8b\x45\xf8',  # mov    rax,QWORD [rbp-0x8]
    b'\x5d',              # pop    rbp
    b'\xc3'])             # ret


def mutate(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 growth(*, shellcode: bytes, length: int) -> bytes:
    shellcode = bytearray(shellcode)

    # slow growth and stop shrinking
    if length > len(shellcode):
        growth = 1
    else:
        growth = 0

    shellcode = shellcode + (b'\x90' * growth)

    return bytes(shellcode)