#!/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
# Change path so we find rdpy
sys.path.insert(1, os.path.join(sys.path[0], '..'))

from rdpy.protocol.rdp import rdp
from twisted.internet import reactor

class ProxyServer(rdp.RDPServerObserver):
    """
    Server side of proxy
    """
    def __init__(self, controller):
        """
        @param controller: RDPServerController
        """
        rdp.RDPServerObserver.__init__(self, controller)
        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._controller.getColorDepth())
        
    def onReady(self):
        """
        Event use to inform state of server stack
        Use to connect client
        """
        width, height = self._controller.getScreen()
        reactor.connectTCP("si-hyperv-002", 3389, 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 connected
        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):
        rdp.ServerFactory.__init__(self, "/home/speyrefitte/dev/certificate/rdpy.key", "/home/speyrefitte/dev/certificate/rdpy.crt", 16)
        
    def buildObserver(self, controller):
        return ProxyServer(controller)
    
    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

if __name__ == '__main__':
    
    reactor.listenTCP(33389, ProxyServerFactory())
    reactor.run()