twisted/protocols/portforward.py

   1
   2 # Copyright (c) Twisted Matrix Laboratories.
   3 # See LICENSE for details.
   4
   5 """
   6 A simple port forwarder.
   7 """
   8
   9 # Twisted imports
  10 from twisted.internet import protocol
  11 from twisted.python import log
  12
  13 class Proxy(protocol.Protocol):
  14     noisy = True
  15
  16     peer = None
  17
  18     def setPeer(self, peer):
  19         self.peer = peer
  20
  21     def connectionLost(self, reason):
  22         if self.peer is not None:
  23             self.peer.transport.loseConnection()
  24             self.peer = None
  25         elif self.noisy:
  26             log.msg("Unable to connect to peer: %s" % (reason,))
  27
  28     def dataReceived(self, data):
  29         self.peer.transport.write(data)
  30
  31 class ProxyClient(Proxy):
  32     def connectionMade(self):
  33         self.peer.setPeer(self)
  34
  35         # Wire this and the peer transport together to enable
  36         # flow control (this stops connections from filling
  37         # this proxy memory when one side produces data at a
  38         # higher rate than the other can consume).
  39         self.transport.registerProducer(self.peer.transport, True)
  40         self.peer.transport.registerProducer(self.transport, True)
  41
  42         # We're connected, everybody can read to their hearts content.
  43         self.peer.transport.resumeProducing()
  44
  45 class ProxyClientFactory(protocol.ClientFactory):
  46
  47     protocol = ProxyClient
  48
  49     def setServer(self, server):
  50         self.server = server
  51
  52     def buildProtocol(self, *args, **kw):
  53         prot = protocol.ClientFactory.buildProtocol(self, *args, **kw)
  54         prot.setPeer(self.server)
  55         return prot
  56
  57     def clientConnectionFailed(self, connector, reason):
  58         self.server.transport.loseConnection()
  59
  60
  61 class ProxyServer(Proxy):
  62
  63     clientProtocolFactory = ProxyClientFactory
  64     reactor = None
  65
  66     def connectionMade(self):
  67         # Don't read anything from the connecting client until we have
  68         # somewhere to send it to.
  69         self.transport.pauseProducing()
  70
  71         client = self.clientProtocolFactory()
  72         client.setServer(self)
  73
  74         if self.reactor is None:
  75             from twisted.internet import reactor
  76             self.reactor = reactor
  77         self.reactor.connectTCP(self.factory.host, self.factory.port, client)
  78
  79
  80 class ProxyFactory(protocol.Factory):
  81     """Factory for port forwarder."""
  82
  83     protocol = ProxyServer
  84
  85     def __init__(self, host, port):
  86         self.host = host
  87         self.port = port