#!/usr/bin/python
# -*- coding: utf-8; -*-
#
# (c) 2013 Mandriva, http://www.mandriva.com/
#
# This file is part of Pulse 2, http://pulse2.mandriva.org
#
# Pulse 2 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 2 of the License, or
# (at your option) any later version.
#
# Pulse 2 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 Pulse 2; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston,
# MA 02110-1301, USA.
"""
Script to establish two proxies to forward the traffic over the SSH tunnel
with checking of requestor IP. When the checking of requestor is enabled,
InputProxy accepts only the requests from passed requestor IP address.
Otherwise, both proxies are stopped.

+---------+ local_port  +----------+ proxy_port +----------+ to_port +------+
|Requestor|------------>|InputProxy|----------->|SSH tunnel|-------->|Client|
+---------+             +----------+            +----------+         +------+
Example:
192.168.20.14:8101 -> (8101:8102) -> 8102:127.0.0.1:5900 -> 192.168.20.199:5900
"""

import os
import sys
import logging
import logging.config

from twisted.internet import reactor, task
from twisted.protocols import portforward
from twisted.internet.endpoints import TCP4ServerEndpoint

from mmc.site import mmcconfdir
from pulse2.launcher.config import LauncherConfig
from pulse2.launcher.process_control import commandProtocol
from pulse2.launcher.txws import WebSocketFactoryMaker

LOCALHOST = "127.0.0.1"

config_file = os.path.join(mmcconfdir, "pulse2", "launchers", "launchers.ini")
config = LauncherConfig()
config.setup(config_file)
logging.config.fileConfig(config_file)

log = logging.getLogger()

class SSHForwardingProxy(object):
    """ Establishing of SSH tunnel by spawning of child process """

    def __init__(self, host, port, command, local_port, endback):
        """
        @param host: destinated machine
        @type host: str

        @param port: port number of destinated machine to open a VNC session
        @type port: int

        @param command: command to execute
        @type command: str

        @param local_port: port number of input of SSH tunnel
        @type local_port: int

        @param endback: called when our process with SSH tunnel is terminated
        @type endback: callable
        """
        self.host = host
        self.port = port
        self.command = command.replace(","," -o")

        self.local_port = local_port
        self.process = commandProtocol(self.command)
        self.process.endback = endback
        self.process.installEndBack()

    def local_to_remote(self):
        """ Establishing of local SSH forward"""
        return self._call_process("L")

    def remote_to_local(self):
        """ Establishing of remote SSH forward"""
        return self._call_process("R")

    def _call_process(self, switch):
        """
        Creating of a new child process.

        @param switch: switch to determine if local or remote forward {L|R}
        @type switch: str
        """

        args = "-%s %i:127.0.0.1:%i -N -n" % (switch,
                                              self.local_port,
                                              self.port)
        d = self._spawn(args)

        @d.addCallback
        def _cb(reason):
            log.debug("VNC Proxy: Output proxy - process successfully spawned")

        @d.addErrback
        def _eb(failure):
            log.warn("VNC Proxy: Output proxy - process spawn failed: %s" % str(failure))

        return d

    def _spawn(self, ssh_args):
        """
        Spawning the process

        @param ssh_args: passed arguments
        @type ssh_args: str
        """

        cmdline = "/usr/bin/ssh %s -o%s %s \n" % (ssh_args,
                                                  self.command,
                                                  self.host)
        cmds = cmdline.split()

        self.process.handler = reactor.spawnProcess(self.process, "/usr/bin/ssh", cmds)

        reactor.addSystemEventTrigger('before',
                                      'shutdown',
                                      self.process.sendSigKill)
        return self.process.deferred

class InputProxyServer(portforward.ProxyServer):
    """Establish a input proxy to check requestor IP"""

    def connectionMade(self):
        log.info("VNC Proxy: New connection from %s on port %d" % (self.from_ip, self.from_port))
        portforward.ProxyServer.connectionMade(self)

    def dataReceived(self, data):
        """Overriden method from ProxyServer"""
        # requestor IP
        client_ip = self.transport.client[0]
        if self.from_ip != client_ip and config.tcp_sproxy_check :
            log.warn("Client address <%s> is not matching with requestor ip <%s>" % (client_ip, self.from_ip))
            log.warn("Request from client address <%s> rejected" % client_ip)
            return

        portforward.ProxyServer.dataReceived(self, data)
        self.timeout_refresh()

    def kill_me(self):
        """ Terminate this proxy """
        self.transport.loseConnection()

    @classmethod
    def set_incomming_machine_info(cls, from_machine):
        """
        Client IP to check.

        @param from_machine: requestor IP address, requestor port
        @type from_machine: tuple
        """
        cls.from_ip, cls.from_port = from_machine

    @classmethod
    def set_establish_delay(cls, establish_delay):
        """
        Setting the timeout of proxy

        @param establish_delay: timeout to expire
        @type establish_delay: int
        """
        cls.establish_delay = establish_delay
        cls.timeout_refresh()

    @classmethod
    def register_stopper(cls, stopper):
        """
        Registering of global method to killing all the runtimes

        @param stopper: method to call if the traffic is inactive
        @type stopper: callable
        """
        cls.stopper = stopper

    @classmethod
    def timeout_refresh(cls):
        """
        Refreshing of the session.

        This metho is periodicaly called when data transport is processed
        """
        if hasattr(cls, "timeout_defer"):
            cls.timeout_defer.cancel()
        if hasattr(cls, "stopper"):
            cls.timeout_defer = task.deferLater(reactor,
                                             cls.establish_delay,
                                             cls.stopper)


class InputProxyFactory(portforward.ProxyFactory):
    """Factory proxy to provide port forwarding"""

    protocol = InputProxyServer

    def __init__(self, host, port, from_ip, from_port, establish_delay, stopper):
        """
        @param host: IP address of portforwarding proxy (usually localhost)
        @type host: str

        @param port: port to listen of incoming requests
        @type port: int

        @param from_ip: requestor IP address to check
        @type from_ip: str

        @param from_port: requestor port
        @type from_port: int

        @param establish_delay: input proxy timeout
        @type establish_delay: int

        @param stopper: handler to call when proxy is timeouted or stopped
        @type stopper: callable
        """
        self.host = host
        self.port = port

        self.protocol.set_incomming_machine_info((from_ip, from_port))
        self.protocol.set_establish_delay(establish_delay)
        self.protocol.register_stopper(stopper)


    def kill_me(self):
        """ calling the terminating method """
        self.protocol.transport.kill_me()

class timeout :
    """Decorator class to declare a timeout for a function"""

    def __init__(self, timeout, action):
        """
        @param timeout: time of expiration of decorated object
        @type timeout: int

        @param action: a callable object which be called after the timeout
        @type action: callable

        """
        self.timeout = timeout
        self.action = action

    def __call__(self, fnc, *args):

        def wrapped(*args):
            self.time_defer = task.deferLater(reactor,
                                         self.timeout,
                                         self.action)
            fnc_defer = fnc(*args)
            fnc_defer.addCallback(self.endback_fnc)

            return fnc_defer

        return wrapped

    def endback_fnc (self, result):
        self.time_defer.cancel()

class DispatchProxies :
    """ Main dispatcher to control our proxies. """
    from_ip = ""
    to_ip = ""
    to_port = 0
    command = ""
    proxy_port = 0
    local_port = 0
    connect_delay = 0
    establish_delay = 0
    session_length = 0
    stoppers = list()

    def __init__(self):
        self.parse_args()
        log.info("VNC Proxy: Request from %s for computer %s (%s)" % (self.from_ip,
                                                                      self.to_ip,
                                                                      self.machine_name,))

        self.start_output_proxy()
        self.start_input_proxy()

        # Session expiration
        d = task.deferLater(reactor,
                        self.session_length,
                        self.stop_proxies)

        d.addCallback(self._session_expired)

    def _session_expired(result):
        log.warn("VNC Proxy: Your session has expired.")


    def parse_args(self):
        """ Input arguments parsing """
        try :
            # requestor ip - an address to verify the source of request
            # in other words a machine requesting a VNC session through 
            # our proxies
            self.from_ip = sys.argv[1]
            # destinated machine to create a VNC session
            self.to_ip = sys.argv[2]
            # assigned port to create a VNC session on dest. machine
            self.to_port = int(sys.argv[3])
            # passed command with all options
            self.command = sys.argv[4]
            # port to communicate between input and output proxy 
            self.proxy_port = int(sys.argv[5])
            # port to listen from reqestor ip
            self.local_port = int(sys.argv[6])
            # timeout of input proxy
            self.establish_delay = int(sys.argv[7])
            # timeout of output proxy
            self.connect_delay = int(sys.argv[8])
            # valid time of session - will be expired after session_length  
            self.session_length = int(sys.argv[9])
            # name of destinated computer in Pulse 
            self.machine_name = sys.argv[10]
            # authentication key
            self.auth_key = sys.argv[11]

        except ValueError, val_exc:
            log.warn("VNC Proxy: Arguments parsing failure: %s" % str(val_exc))
            self.stop_proxies()
        except IndexError, idx_exc:
            log.warn("VNC Proxy: Arguments parsing failure: %s" % str(idx_exc))
            self.stop_proxies()



    def register_stopper(self, stopper):
        """
        Registering of stopping methods of proxies

        @param stopper: proxy method to terminate of processing
        @type stopper: callable
        """
        self.stoppers.append(stopper)

    def stop_proxies(self):
        """ Stopping of all the proxies and exit from script """
        for stopper in self.stoppers :
            try :
                if callable(stopper):
                    stopper()
            except AttributeError :
                log.debug("VNC Proxy: service probably stopped, going to next")
            except Exception, exc:
                log.warn("VNC Proxy: Stopping failed: %s" % str(exc))

        log.info("VNC Proxy: Proxy for %s (%s) closed" % (self.to_ip, self.machine_name))

        reactor.stop()

    # stop_proxies called when connect_delay passed
    @timeout(connect_delay, stop_proxies)
    def start_output_proxy(self):
        """ Start a process spawn which create a SSH tunnel"""

        self.out_proxy = SSHForwardingProxy(self.to_ip,
                                            self.to_port,
                                            self.command,
                                            self.local_port,
                                            self._endback_output_proxy)

        d = task.deferLater(reactor, 0, self.out_proxy.local_to_remote)

        d.addCallback(self._cb_output_proxy)
        d.addErrback(self._eb_output_proxy)

        return d

    def _cb_output_proxy(self, result):
        log.debug("VNC Proxy: Output proxy successfully started")
        try:
            if hasattr(self.out_proxy, "process"):
                self.register_stopper(self.out_proxy.process.sendSigKill)
        except Exception, e:
            log.warn("VNC Proxy: %s" %str(e))

    def _eb_output_proxy(self, failure):
        log.warn("VNC Proxy: Start of output proxy failed: %s" % str(failure))
        self.stop_proxies()

    def _endback_output_proxy(self, result):
        log.debug("VNC Proxy: Output proxy - SSH tunnel building finished")


    def _start_input_proxy(self):
        """
        Creating of input proxy that is checking requestor IP
        and forwarding TCP trafic to binded SSH tunnel.
        """
        server = TCP4ServerEndpoint(reactor, self.proxy_port)
        self.in_proxy = InputProxyFactory(LOCALHOST,
                                          self.local_port,
                                          self.from_ip,
                                          self.proxy_port,
                                          self.establish_delay,
                                          self.stop_proxies)
        if config.create_web_proxy:
            WebSocketFactory = WebSocketFactoryMaker(self.auth_key)
            server.listen(WebSocketFactory(self.in_proxy))
        else:
            server.listen(self.in_proxy)

    @timeout(establish_delay, stop_proxies)
    def start_input_proxy(self):
        """
        Creating of input proxy that is checking requestor IP
        and forwarding TCP trafic to binded SSH tunnel.
        """
        d = task.deferLater(reactor, 1, self._start_input_proxy)
        d.addCallback(self._cb_input_proxy)
        d.addErrback(self._eb_input_proxy)

        return d

    def _cb_input_proxy(self, result):
        log.debug("VNC Proxy: Input proxy successfully started")
        message = "VNC Proxy: Proxy established for %s (%s):%d on %s:%d => Connection allowed for %s on port %d"
        log.info(message % (self.to_ip,
                            self.machine_name,
                            self.to_port,
                            LOCALHOST,
                            self.local_port,
                            self.from_ip,
                            self.proxy_port))

        self.register_stopper(self.in_proxy.kill_me)

    def _eb_input_proxy(self, failure):
        log.warn("VNC Proxy: Start of input proxy failed: %s" % str(failure))
        self.stop_proxies()



if __name__ == "__main__" :

    proxy_type = config.create_web_proxy and 'WebSocket' or 'TCP'
    log.info('VNC Proxy: Starting new ' + proxy_type + ' proxy: '+' '.join(sys.argv))
    reactor.callWhenRunning(DispatchProxies)
    reactor.run()
    sys.exit(0)
