pythonforward.py

Go to the documentation of this file.

import select
import socket
import threading
from robot.utils import PY2, WINDOWS
from .logger import logger
if PY2 and WINDOWS:
    import win_inet_pton
try:
    import SocketServer
except ImportError:
    import socketserver as SocketServer
 
 
def check_if_ipv6(ip):
    try:
        socket.inet_pton(socket.AF_INET6, ip)
        return True
    except socket.error:
        return False
 
 
class LocalPortForwarding:
    def __init__(self, port, host, transport, bind_address):
        self.serverserver = None
        self.portport = port
        self.hosthost = host
        self.transporttransport = transport
        self.bind_addressbind_address = bind_address
 
    def forward(self, local_port):
        class SubHandler(LocalPortForwardingHandler):
            port = self.portport
            host = self.hosthost
            ssh_transport = self.transporttransport
 
        self.serverserver = ForwardServer((self.bind_addressbind_address or '', local_port), SubHandler, ipv6=check_if_ipv6(self.hosthost))
        t = threading.Thread(target=self.serverserver.serve_forever)
        t.setDaemon(True)
        t.start()
        logger.info("Now forwarding port %d to %s:%d ..." % (local_port, self.hosthost, self.portport))
 
    def close(self):
        if self.serverserver:
            self.serverserver.shutdown()
            try:
                logger.log_background_messages()
            except AttributeError:
                pass
 
 
class ForwardServer(SocketServer.ThreadingTCPServer):
    daemon_threads = True
    allow_reuse_address = True
 
    def __init__(self, server_address, RequestHandlerClass, ipv6=False):
        if ipv6:
            ForwardServer.address_family = socket.AF_INET6
        SocketServer.ThreadingTCPServer.__init__(self, server_address, RequestHandlerClass, bind_and_activate=True)
 
 
class LocalPortForwardingHandler(SocketServer.BaseRequestHandler):
    host, port, ssh_transport = None, None, None
 
    def handle(self):
        try:
            chan = self.ssh_transportssh_transport.open_channel('direct-tcpip', (self.hosthost, self.portport),
                                                   self.request.getpeername())
        except Exception as e:
            logger.info("Incoming request to %s:%d failed: %s" % (self.hosthost, self.portport, repr(e)))
            return
        if chan is None:
            logger.info("Incoming request to %s:%d was rejected by the SSH server." % (self.hosthost, self.portport))
            return
        logger.info("Connected! Tunnel open %r -> %r -> %r" % (self.request.getpeername(),
                                                               chan.getpeername(),
                                                               (self.hosthost, self.portport)))
        while True:
            r, w, x = select.select([self.request, chan], [], [])
            if self.request in r:
                data = self.request.recv(1024)
                if len(data) == 0:
                    break
                chan.send(data)
            if chan in r:
                data = chan.recv(1024)
                if len(data) == 0:
                    break
                self.request.send(data)
        peername = self.request.getpeername()
        chan.close()
        self.request.close()
        logger.info("Tunnel closed from %r" % (peername,))