#!/usr/bin/python
#
# Copyright (c) 2014 Sylvain Peyrefitte
#
# This file is part of rdpy.
#
# rdpy 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.
#
# This program 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 this program. If not, see <http://www.gnu.org/licenses/>.
#

"""
RDP proxy recorder and spy function
Proxy RDP protocol
"""

import sys, os, getopt, json
# Change path so we find rdpy
sys.path.insert(1, os.path.join(sys.path[0], '..'))

from rdpy.base import log, error
from rdpy.protocol.rdp import rdp
from rdpy.ui import widget
from twisted.internet import reactor
from PyQt4 import QtCore, QtGui

class IProxyClient(object):
    def getColorDepth(self):
        pass
    def sendKeyEventScancode(self, code, isPressed):
        pass
    def sendKeyEventUnicode(self, code, isPressed):
        pass
    def sendPointerEvent(self, x, y, button, isPressed):
        pass

class ProxyServer(rdp.RDPServerObserver):
    """
    Server side of proxy
    """
    def __init__(self, controller, credentialProvider):
        """
        @param controller: RDPServerController
        @param credentialProvider: CredentialProvider
        """
        rdp.RDPServerObserver.__init__(self, controller)
        self._credentialProvider = credentialProvider
        self._client = None
        
    def clientConnected(self, client):
        """
        Event throw by client when it's ready
        @param client: ProxyClient
        """
        self._client = client
        self._controller.setColorDepth(self._client.getColorDepth())
        
    def onReady(self):
        """
        Event use to inform state of server stack
        Use to connect client
        """
        domain, username, password = self._controller.getCredentials()
        
        if self._credentialProvider.isAdmin(domain, username, password):
            self.clientConnected(ProxyAdmin(self))
            return
        
        try:
            dstIp, dstPort, dstDomain, dstUsername, dstPassword = self._credentialProvider.getCredentials(domain, username, password)
        except error.InvalidExpectedDataException as e:
            log.info(e.message)
            #self._controller.close()
            return
        
        width, height = self._controller.getScreen()
        reactor.connectTCP(dstIp, dstPort, ProxyClientFactory(self, width, height, dstDomain, dstUsername, dstPassword, "%s/%s -> %s %s/%s"%(domain, username, dstIp, dstDomain, dstUsername)))
        
    def onKeyEventScancode(self, code, isPressed):
        """
        Event call when a keyboard event is catch in scan code format
        @param code: scan code of key
        @param isPressed: True if key is down
        """
        #no client connected
        if self._client is None:
            return
        
        self._client.sendKeyEventScancode(code, isPressed)
    
    def onKeyEventUnicode(self, code, isPressed):
        """
        Event call when a keyboard event is catch in unicode format
        @param code: unicode of key
        @param isPressed: True if key is down
        """
        #no client connectedomaind
        if self._client is None:
            return
        
        self._client.sendKeyEventUnicode(code, isPressed)
        
    def onPointerEvent(self, x, y, button, isPressed):
        """
        Event call on mouse event
        @param x: x position
        @param y: y position
        @param button: 1, 2 or 3 button
        @param isPressed: True if mouse button is pressed
        """
        #no client connected
        if self._client is None:
            return
        
        self._client.sendPointerEvent(x, y, button, isPressed)
    
class ProxyClient(rdp.RDPClientObserver, IProxyClient):
    """
    Client side of proxy
    """
    def __init__(self, controller, server):
        """
        @param controller: RDPClientObserver
        @param server: ProxyServer
        @param name: name of session
        """
        rdp.RDPClientObserver.__init__(self, controller)
        self._server = server
        
    def onReady(self):
        """
        Event use to signal that RDP stack is ready
        Inform proxy server that i'm connected
        implement RDPClientObserver
        """
        self._server.clientConnected(self)
        
    def onUpdate(self, destLeft, destTop, destRight, destBottom, width, height, bitsPerPixel, isCompress, data):
        """
        Event use to inform bitmap update
        implement RDPClientObserver
        @param destLeft: xmin position
        @param destTop: ymin position
        @param destRight: xmax position because RDP can send bitmap with padding
        @param destBottom: ymax position because RDP can send bitmap with padding
        @param width: width of bitmap
        @param height: height of bitmap
        @param bitsPerPixel: number of bit per pixel
        @param isCompress: use RLE compression
        @param data: bitmap data
        """
        self._server._controller.sendUpdate(destLeft, destTop, destRight, destBottom, width, height, bitsPerPixel, isCompress, data)
    def getColorDepth(self):
        return self._controller.getColorDepth() 
    def sendKeyEventScancode(self, code, isPressed):
        self._controller.sendKeyEventScancode(code, isPressed)
    def sendKeyEventUnicode(self, code, isPressed):
        self._controller.sendKeyEventUnicode(code, isPressed)
    def sendPointerEvent(self, x, y, button, isPressed):
        self._controller.sendPointerEvent(x, y, button, isPressed)

class ProxyServerFactory(rdp.ServerFactory):
    """
    Factory on listening events
    """
    def __init__(self, credentialProvider):
        """
        @param config: rdp-proxy configuration
        @param credentialProvider: CredentialProvider
        """
        rdp.ServerFactory.__init__(self, "/home/sylvain/dev/certificate/rdpy.key", "/home/sylvain/dev/certificate/rdpy.crt", 16)
        self._credentialProvider = credentialProvider
        
    def buildObserver(self, controller):
        """
        Implement rdp.ServerFactory
        @param controller: rdp.RDPServerController
        """
        return ProxyServer(controller, self._credentialProvider)
    
    def startedConnecting(self, connector):
        pass
    
    def clientConnectionLost(self, connector, reason):
        pass
        
    def clientConnectionFailed(self, connector, reason):
        pass

class ProxyClientFactory(rdp.ClientFactory):
    """
    Factory for proxy client
    """
    _CLIENT_PROXY_ = {}
    def __init__(self, server, width, height, domain, username, password, name):
        """
        @param server: ProxyServer
        @param width: screen width
        @param height: screen height
        @param domain: domain session
        @param username: username session
        @param password: password session
        """
        self._server = server
        self._width = width
        self._height = height
        self._domain = domain
        self._username = username
        self._password = password
        self._name = name
        
    def buildObserver(self, controller):
        """
        Implement rdp.ClientFactory
        Build observer (ProxyClient)
        @param controller: rdp.RDPClientController
        """
        #set screen resolution
        controller.setScreen(self._width, self._height)
        #set credential
        controller.setDomain(self._domain)
        controller.setUsername(self._username)
        controller.setPassword(self._password)
        proxy = ProxyClient(controller, self._server)
        ProxyClientFactory._CLIENT_PROXY_[self._name] = proxy
        return proxy
    
    def startedConnecting(self, connector):
        pass
    
    def clientConnectionLost(self, connector, reason):
        if ProxyClientFactory._CLIENT_PROXY_.has_key(self._name):
            del ProxyClientFactory._CLIENT_PROXY_[self._name]
        
    def clientConnectionFailed(self, connector, reason):
        if ProxyClientFactory._CLIENT_PROXY_.has_key(self._name):
            del ProxyClientFactory._CLIENT_PROXY_[self._name]
    
class CredentialProvider(object):
    """
    Credential provider for proxy
    """
    def __init__(self, config):
        """
        @param config: rdp proxy config
        """
        self._config = config
        
    def getCredentials(self, domain, username, password):
        """
        @param condig: rdp config
        @param domain: domain to check
        @param username: username in domain
        @param password: password in domain
        @return: (ip, port, domain, username, password) or None if error
        """
        if not self._config['domain'].has_key(domain):
            raise error.InvalidExpectedDataException("Unknown domain %s"%(domain))
        for user in self._config['domain'][domain]:
            if user['username'] == username and user['password'] == password:
                return str(user['credentials']['ip']), user['credentials']['port'], str(user['credentials']['domain']), str(user['credentials']['username']), str(user['credentials']['password'])
        raise error.InvalidExpectedDataException("Unknown credential %s\%s"%(domain, username))

    def isAdmin(self, domain, username, password):
        """
        @return: True if credential match admin credential
        """
        return self._config['admin']['domain'] == domain and self._config['admin']['username'] == username and self._config['admin']['password'] == password

class ProxyAdmin(IProxyClient):
    def __init__(self, server):
        self._server = server
        #self._list = widget.List(ProxyClientFactory._CLIENT_PROXY_.keys(), 100, 100, self.onSelect, widget.Anchor(0, 0, widget.RDPWidgetListener(self._server._controller)))
        self._list = widget.List(["salut les copains"], 300, 300, self.onSelect, widget.Anchor(0, 0, widget.RDPWidgetListener(self._server._controller)))
    def getColorDepth(self):
        return 16
    def sendKeyEventScancode(self, code, isPressed):
        if isPressed:
            self._list.keyEvent(code)
    def sendKeyEventUnicode(self, code, isPressed):
        pass
    def sendPointerEvent(self, x, y, button, isPressed):
        pass
    def onSelect(self, name):
        if ProxyClientFactory._CLIENT_PROXY_.has_key(name):
            self._server.clientConnected(ProxyClient(ProxyClientFactory._CLIENT_PROXY_[name]._controller, self._server))
            
def help():
    """
    Print help in console
    """
    print "Usage: rdpy-rdpproxy -f config_file_path listen_port"
    
def loadConfig(configFilePath):
    """
    Load and check config file
    @param configFilePath: config file path
    """
    if not os.path.isfile(configFilePath):
        log.error("File %s doesn't exist"%configFilePath)
        return None
    
    f = open(configFilePath, 'r')
    config = json.load(f)
    f.close()
    
    if not config.has_key('domain'):
        log.error("Need domain definition in config file")
        return None
    
    #check admin account
    if not config.has_key('admin') or not config['admin'].has_key('domain') or not config['admin'].has_key('username') or not config['admin'].has_key('password'):
        log.error("Bad admin account definition in config file")
        return None
    
    return config

if __name__ == '__main__':
    configFilePath = None
    try:
        opts, args = getopt.getopt(sys.argv[1:], "hf:")
    except getopt.GetoptError:
        help()
    for opt, arg in opts:
        if opt == "-h":
            help()
            sys.exit()
        elif opt == "-f":
            configFilePath = arg
            
    if configFilePath is None:
        print "Config file is mandatory"
        help()
        sys.exit()
    
    #load config file
    config = loadConfig(configFilePath)
    if config is None:
        log.error("Bad configuration file")
        sys.exit()
    
    #use to init font
    app = QtGui.QApplication(sys.argv)
    
    reactor.listenTCP(int(args[0]), ProxyServerFactory(CredentialProvider(config)))
    reactor.run()