shellcode_gen.py

""" Built-in modules """
import os
import re
import shlex
import sys
from pathlib import Path
from subprocess import Popen, TimeoutExpired


def objdump_run(filename: str, tmp_file: Path):
    """
    Takes the passed in binary executable, runs objdump utility, and writes the output to file.

    :param filename:  The file to be run through the objdump utility.
    :param tmp_file:  The temp file where the objdump output will be saved to.
    :return:  Nothing
    """
    # Shell escape filename for execution #
    file = shlex.quote(filename)

    try:
        # Open the temp file in write mode #
        with tmp_file.open('w', encoding='utf-8') as out_file:
            # Create objdump command process, storing output in open file #
            with Popen(['objdump', '-M', 'intel', '-D', file], stdout=out_file,
                       stderr=out_file) as command:
                command.communicate()

    # If error occurs during file operation #
    except OSError as io_err:
        print_err(f'Error occurred writing {file} objdump to {tmp_file}: {io_err}')
        sys.exit(2)


def print_err(msg: str):
    """
    Prints a timed error message via stderr.

    :param msg:  The error message to be displayed.
    :return:  Nothing
    """
    print(f'\n* [ERROR]: {msg} *\n', file=sys.stderr)


def prompt_user(re_exe) -> str:
    """
    Prompt the user for input, handle errors accordingly.

    :param re_exe:  The regular expression match format.
    :return:  Validated user input.
    """
    cmd = shlex.quote('clear')

    while True:
        # Clear the display #
        os.system(cmd)
        # Prompt user for input #
        prompt = input('[+] Invalid or no args provided .. enter file to grab bytes: ')

        # If input validation regex fails #
        if not re.search(re_exe, prompt):
            print_err('Improper input .. try again')
            continue

        return prompt


def main():
    """
    Take input file, run in objdump, and parse output into shellcode.

    :return:  Nothing
    """
    # Compile path regex #
    re_exe = re.compile(r'[a-zA-Z\d_\"\' .,\-]{1,20}')

    # If an arg was passed in #
    if len(sys.argv) > 1:
        # Check if arg was passed in #
        arg_check = re.search(re_exe, sys.argv[1])

        # If regex did not match #
        if not arg_check:
            # Prompt the user for input #
            filename = prompt_user(re_exe)
        else:
            # Set the passed in arg as filename #
            filename = sys.argv[1]
    # If no args were passed in #
    elif len(sys.argv) == 1:
        # Prompt the user for binary tp execute on #
        filename = prompt_user(re_exe)
    # If unexpected error occurs parsing args #
    else:
        print_err('Unknown error occurred on program startup')
        # Exit program on error #
        sys.exit(1)

    out_path = Path('/tmp')
    # Set temp file name for pre-parsing #
    tmp_file = out_path / f'{filename}.txt'

    # Run objdump utility on binary file and save to tmp file #
    objdump_run(filename, tmp_file)

    # Compile regular expression for byte formatting #
    re_byte = re.compile(r'\s([a-f\d]{2}\s){1,7}')
    re_space_strip = re.compile(r'\s[^\S\r\n]')
    re_raw_byte = re.compile(r'\s|\t')
    shellcode = ''

    try:
        # Open the temporary file in read mode #
        with tmp_file.open('r', encoding='utf-8') as re_file:
            # Iterate through file line by line #
            for line in re_file:
                # Attempt to match bytes in line #
                byte_grab = re.search(re_byte, line)
                # If regex matches #
                if byte_grab:
                    # Strip whitespace from match #
                    strip = re.sub(re_space_strip, r'', byte_grab.group(0))
                    # Format raw byte like \xXX #
                    raw_byte = re.sub(re_raw_byte, r'\\x', strip)
                    # Append raw byte to shellcode string #
                    shellcode += raw_byte[:-2]

        print(f'Shellcode: {shellcode}')

    # If error occurs during file operation #
    except OSError as io_err:
        print_err(f'Error occurred writing shellcode to {tmp_file.name}: {io_err}')
        sys.exit(3)

    # Unlink the temporary parsing file #
    tmp_file.unlink()

    sys.exit(0)


if __name__ == '__main__':
    main()