#!/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 spyer
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 twisted.internet import reactor
clientMacro = None
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
        """
        global clientMacro
        self._client = client
        clientMacro = client
        self._controller.setColorDepth(self._client._controller.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(ProxyClient(clientMacro._controller, self))
            return
        
        try:
            ip, port, domain, username, password = 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(ip, port, ProxyClientFactory(self, width, height))
        
    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._controller.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._controller.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._controller.sendPointerEvent(x, y, button, isPressed)
    
class ProxyClient(rdp.RDPClientObserver):
    """
    Client side of proxy
    """
    def __init__(self, controller, server):
        """
        @param controller: RDPClientObserver
        @param server: ProxyServer
        """
        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)

class ProxyServerFactory(rdp.ServerFactory):
    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):
    def __init__(self, server, width, height):
        self._server = server
        self._width = width
        self._height = height
        
    def buildObserver(self, controller):
        controller.setScreen(self._width, self._height)
        return ProxyClient(controller, self._server)
    
    def startedConnecting(self, connector):
        pass
    
    def clientConnectionLost(self, connector, reason):
        pass
        
    def clientConnectionFailed(self, connector, reason):
        pass

def help():
    """
    Print help in consoe
    """
    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

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 user['credentials']['ip'], user['credentials']['port'], user['credentials']['domain'], user['credentials']['username'], 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

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()
    
    reactor.listenTCP(int(args[0]), ProxyServerFactory(CredentialProvider(config)))
    reactor.run()