#!/usr/bin/env python2
import argparse, sys, os, types
import pwnlib
from pwnlib import asm, log, util
from pwnlib.context import __possible__ as contexts
import pwnlib.term.text as text

r = text.red
g = text.green
b = text.blue
banner = '\n'.join(['  ' + r('____') + '  ' + g('_') + '          ' + r('_') + ' ' + g('_') + '                 ' + b('__') + ' ' + r('_'),
                    ' ' + r('/ ___|') + g('| |__') + '   ' + b('___') + r('| |') + ' ' + g('|') + ' ' + b('___') + ' ' + r('_ __') + ' ' + g('__ _') + ' ' + b('/ _|') + ' ' + r('|_'),
                    ' ' + r('\___ \\') + g('| \'_ \\') + ' ' + b('/ _ \\') + ' ' + r('|') + ' ' + g('|') + b('/ __|') + ' ' + r('\'__/') + ' ' + g('_` |') + ' ' + b('|_') + r('| __|'),
                    '  ' + r('___) |') + ' ' + g('| | |') + '  ' + b('__/') + ' ' + r('|') + ' ' + g('|') + ' ' + b('(__') + r('| |') + ' ' + g('| (_| |') + '  ' + b('_|') + ' ' + r('|_'),
                    ' ' + r('|____/') + g('|_| |_|') + b('\\___|') + r('_|') + g('_|') + b('\\___|') + r('_|') + '  ' + g('\\__,_|') + b('_|') + '  ' + r('\\__|'),
                    ''
                    ])
#  ____  _          _ _                 __ _
# / ___|| |__   ___| | | ___ _ __ __ _ / _| |_
# \___ \| '_ \ / _ \ | |/ __| '__/ _` | |_| __|
#  ___) | | | |  __/ | | (__| | | (_| |  _| |_
# |____/|_| |_|\___|_|_|\___|_|  \__,_|_|  \__|

def _string(s):
    out = []
    for c in s:
        co = ord(c)
        if co >= 0x20 and co <= 0x7e and c not in '/$\'"`':
            out.append(c)
        else:
            out.append('\\x%02x' % co)
    return '"' + ''.join(out) + '"\n'

def _carray(s):
    out = []
    for c in s:
        out.append('0x' + util.fiddling.enhex(c))
    return '{' + ', '.join(out) + '};\n'

def _hex(s):
    return pwnlib.util.fiddling.enhex(s) + '\n'

p = argparse.ArgumentParser(
    description = banner + '\nMicrowave shellcode -- Easy, fast and delicious',
    formatter_class = argparse.RawDescriptionHelpFormatter,
)

p.add_argument(
    'shellcode',
    nargs = '?',
    default = '',
    metavar = '<shellcode>',
    help = 'The shellcode you want',
)

p.add_argument(
    'args',
    nargs = '*',
    metavar = '<arg>',
    default = (),
    help = 'Argument to the chosen shellcode',
)

p.add_argument(
    '-?', '--show',
    action = 'store_true',
    help = 'Show shellcode documentation',
)

p.add_argument(
    '-o', '--out',
    metavar = '<file>',
    type = argparse.FileType('w'),
    default = sys.stdout,
    help = 'Output file (default: stdout)',
)

p.add_argument(
    '-f', '--format',
    metavar = '<format>',
    choices = ['r', 'raw',
               's', 'str', 'string',
               'c',
               'h', 'hex',
               'a', 'asm', 'assembly',
               'p',
               'i', 'hexii'],
    default = 'hex',
    help = 'Output format (default: hex), choose from {r}aw, {s}tring, {c}-style array, {h}ex string, hex{i}i, {a}ssembly code, {p}reprocssed code',
)

args = p.parse_args()

def get_tree(path, val, result):
    if path:
        path += '.'

    if path.startswith('common.'):
        return

    mods = []

    for k in sorted(dir(val)):
        if k and k[0] != '_':
            cur = getattr(val, k)
            if isinstance(cur, types.ModuleType):
                mods.append((path + k, cur))
            else:
                result.append((path + k, cur))

    for path, val in mods:
        get_tree(path, val, result)
    return result

vals = get_tree('', pwnlib.shellcraft, [])
if args.shellcode:
    vals = [(k, val) for k, val in vals if k.startswith(args.shellcode + '.') or k == args.shellcode]

if len(vals) == 0:
    log.fatal("Cannot find subtree by the name of %r" % args.shellcode)
elif len(vals) > 1:
    for k, _ in vals:
        print k
    exit()
else:
    func = vals[0][1]

if args.show:
    print func.__doc__
    exit()

defargs = len(func.func_defaults or ())
reqargs = func.func_code.co_argcount - defargs
if len(args.args) < reqargs:
    if defargs > 0:
        log.fatal('%s takes at least %d arguments' % (args.shellcode, reqargs))
    else:
        log.fatal('%s takes exactly %d arguments' % (args.shellcode, reqargs))

# Captain uglyness saves the day!
for i in range(len(args.args)):
    if args.args[i] in ['False', 'false']:
        args.args[i] = False

# And he strikes again!
os = arch = None
for k in args.shellcode.split('.')[:-1]:
    if k in contexts['arch']:
        arch = k
    elif k in contexts['os']:
        os = k

code = func(*args.args)

if args.format in ['a', 'asm', 'assembly']:
    print code
    exit()
if args.format == 'p':
    print pwnlib.asm.cpp(code, arch = arch, os = os)
    exit()

code = pwnlib.asm.asm(code, arch = arch, os = os)

if args.format in ['s', 'str', 'string']:
    code = _string(code)
elif args.format == 'c':
    code = _carray(code)
elif args.format in ['h', 'hex']:
    code = _hex(code)
elif args.format in ['i', 'hexii']:
    code = pwnlib.util.fiddling.hexii(code) + '\n'

sys.stdout.write(code)
