#!/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/>.
#
from twisted.mail.pop3client import ERR

"""
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 view
from twisted.internet import reactor
from PyQt4 import QtCore, QtGui

class IProxyClient(object):
    """
    Interface use by Proxy server to interact with client
    """
    def close(self):
        """
        Close proxy client
        """
        raise error.CallPureVirtualFuntion("%s:%s defined by interface %s"%(self.__class__, "getColorDepth", "IProxyClient"))
    
    def getColorDepth(self):
        """
        Color depth client, Use to re-negociate color depth with server
        """
        raise error.CallPureVirtualFuntion("%s:%s defined by interface %s"%(self.__class__, "getColorDepth", "IProxyClient"))
    
    def sendKeyEventScancode(self, code, isPressed):
        """
        Key event as scan code
        @param code: scan code of key
        @param isPressed: True if key is down
        """
        raise error.CallPureVirtualFuntion("%s:%s defined by interface %s"%(self.__class__, "sendKeyEventScancode", "IProxyClient"))
    
    def sendKeyEventUnicode(self, code, isPressed):
        """
        Key event as unicode
        @param code: unicode of key
        @param isPressed: True if key is down
        """
        raise error.CallPureVirtualFuntion("%s:%s defined by interface %s"%(self.__class__, "sendKeyEventUnicode", "IProxyClient"))
    
    def sendPointerEvent(self, x, y, button, isPressed):
        """
        Mouse event
        @param x: x position
        @param y: y position
        @param isPressed: True if button is down
        """
        raise error.CallPureVirtualFuntion("%s:%s defined by interface %s"%(self.__class__, "sendPointerEvent", "IProxyClient"))
    
    def sendRefreshOrder(self, left, top, right, bottom):
        """
        Refresh zone
        @param left: left postion
        @param top: top position
        @param right: right position
        @param bottom: bottom position
        """
        raise error.CallPureVirtualFuntion("%s:%s defined by interface %s"%(self.__class__, "sendRefreshOrder", "IProxyClient"))
    

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
        """
        if self._client is None:
            #try a connection
            domain, username, password = self._controller.getCredentials()
            
            if self._credentialProvider.isAdmin(domain, username, password):
                self.clientConnected(ProxyAdmin(self))
                return
            
            try:
                ip, port = self._credentialProvider.getProxyPass(domain, username)
            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, domain, username, password, "%s\\%s on %s:%s"%(domain, username, ip, port)))
        else:
            #refresh client
            width, height = self._controller.getScreen()
            self._client.sendRefreshOrder(0, 0, width, height)
            
    def onClose(self):
        """
        Call when client close connection
        """
        if not self._client is None:
            self._client.close()
        
    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 connected domain
        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
    """
    _CONNECTED_ = {}
    def __init__(self, controller, server, name):
        """
        @param controller: RDPClientObserver
        @param server: ProxyServer
        @param name: name of session
        """
        rdp.RDPClientObserver.__init__(self, controller)
        self._server = server
        self._name = name
        
    def onReady(self):
        """
        Event use to signal that RDP stack is ready
        Inform proxy server that i'm connected
        implement RDPClientObserver
        """
        ProxyClient._CONNECTED_[self._name] = self
        self._server.clientConnected(self)
        
    def onClose(self):
        """
        Stack is closes
        """
        if ProxyClient._CONNECTED_.has_key(self._name):
            del ProxyClient._CONNECTED_[self._name]
            self._server._controller.close()
        
    def close(self):
        """
        Close proxy client
        """
        self._controller.close()
        
    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):
        """
        Color depth client, Use to re-negociate color depth with server
        """
        return self._controller.getColorDepth() 
    
    def sendKeyEventScancode(self, code, isPressed):
        """
        Key event as scan code
        @param code: scan code of key
        @param isPressed: True if key is down
        """
        self._controller.sendKeyEventScancode(code, isPressed)
        
    def sendKeyEventUnicode(self, code, isPressed):
        """
        Key event as uni code
        @param code: uni code of key
        @param isPressed: True if key is down
        """
        self._controller.sendKeyEventUnicode(code, isPressed)
        
    def sendPointerEvent(self, x, y, button, isPressed):
        """
        Mouse event
        @param x: x position
        @param y: y position
        @param isPressed: True if button is down
        """
        self._controller.sendPointerEvent(x, y, button, isPressed)
        
    def sendRefreshOrder(self, left, top, right, bottom):
        """
        Refresh zone
        @param left: left postion
        @param top: top position
        @param right: right position
        @param bottom: bottom position
        """
        self._controller.sendRefreshOrder(left, top, right, bottom)

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/speyrefitte/dev/certificate/rdpy.key", "/home/speyrefitte/dev/certificate/rdpy.crt", 16)
        self._credentialProvider = credentialProvider
        
    def buildObserver(self, controller):
        """
        Implement rdp.ServerFactory
        @param controller: rdp.RDPServerController
        """
        return ProxyServer(controller, self._credentialProvider)

class ProxyClientFactory(rdp.ClientFactory):
    """
    Factory for proxy client
    """
    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
        @param name: name of 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, self._name)
        return proxy
    
    def startedConnecting(self, connector):
        pass
    
    def clientConnectionLost(self, connector, reason):
        pass
        
    def clientConnectionFailed(self, connector, reason):
        pass
            
    
class ProxyAdmin(IProxyClient):
    """
    Use to manage client side of admin session
    Add GUI to select which session to see
    And manage see session
    Just escape key is authorized during see session
    """
    class State(object):
        """
        GUI state -> list of active session
        SPY state -> watch active session
        """
        GUI = 0
        SPY = 1
        
    def __init__(self, server):
        """
        @param server: rdp.RDPServerController
        """
        self._server = server
        self._spyProxy = None
        self.initView()
        self._state = ProxyAdmin.State.GUI
        
    def initView(self):
        """
        Init GUI view
        """
        width, height = self._server._controller.getScreen()
        self._window = view.WindowView(width, height, QtGui.QColor(24, 93, 123))
        self._window.addView(view.AnchorView(width / 2 - 250, height / 2 - 250, view.ListView(ProxyClient._CONNECTED_.keys(), 500, 500, self.onSelect, QtGui.QColor(24, 93, 123)))) 
        self._render = view.RDPRenderer(self._server._controller)
    
    def close(self):
        """
        Close proxy client
        """
        pass
    
    def getColorDepth(self):
        """
        Use same Color depth as server init
        @return color depth of GUI
        """
        if self._state == ProxyAdmin.State.GUI:
            return self._server._controller.getColorDepth()
        elif self._state == ProxyAdmin.State.SPY:
            return self._spyProxy.getColorDepth()
    
    def sendKeyEventScancode(self, code, isPressed):
        """
        IProxyClient implement is unauthorized during admin session
        Only for GUI
        @param code: scan code of key
        @param isPressed: True if key is down
        """
        if self._state == ProxyAdmin.State.GUI:
            if not isPressed:
                return
            self._window.keyEvent(code)
            self._window.update(self._render)
        elif code == 1:
            #escape button refresh GUI
            self._spyProxy._controller.removeClientObserver(self._spyProxy)
            self._state = ProxyAdmin.State.GUI
            self._server.clientConnected(self)
            
    def sendKeyEventUnicode(self, code, isPressed):
        """
        Key event as uni code is unauthorized during admin session
        @param code: uni code of key
        @param isPressed: True if key is down
        """
        pass
    
    def sendPointerEvent(self, x, y, button, isPressed):
        """
        Mouse event is unauthorized during admin session
        @param x: x position
        @param y: y position
        @param isPressed: True if button is down
        """
        pass
    
    def sendRefreshOrder(self, left, top, right, bottom):
        """
        Refresh zone
        @param left: left postion
        @param top: top position
        @param right: right position
        @param bottom: bottom position
        """
        if self._state == ProxyAdmin.State.GUI:
            self._window.update(self._render)
        elif self._state == ProxyAdmin.State.SPY:
            self._spyProxy.sendRefreshOrder(left, top, right, bottom)
        
    def onSelect(self, name):
        """
        Call back of list view
        @param name: name selected by user
        """
        if not ProxyClient._CONNECTED_.has_key(name):
            return
        self._spyProxy = ProxyClient(ProxyClient._CONNECTED_[name]._controller, self._server, "Admin")
        self._state = ProxyAdmin.State.SPY
        #reconnect me
        self._server.clientConnected(self)

class CredentialProvider(object):
    """
    Credential provider for proxy
    """
    def __init__(self, config):
        """
        @param config: rdp proxy config
        """
        self._config = config
        
    def getAccount(self, domain, username):
        if not self._config.has_key(domain) or not self._config[domain].has_key(username):
            raise error.InvalidExpectedDataException("Invalid credentials %s\\%s"%(domain, username))
        
        return self._config[domain][username]
        
    def getProxyPass(self, domain, username):
        """
        @param domain: domain to check
        @param username: username in domain
        @return: (ip, port) or None if error
        """
        account = self.getAccount(domain, username)
        if not account.has_key("ip") or not account.has_key("port"):
            raise error.InvalidExpectedDataException("Invalid credentials declaration %s\\%s"%(domain, username))
        return str(account['ip']), account['port']
    
    def isAdmin(self, domain, username, password):
        """
        @return: True if credential match admin credential
        """
        account = self.getAccount(domain, username)
        return account.has_key("admin") and account["admin"] and account.has_key("password") and str(account["password"]) == password
       
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()
    
    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()