twisted/protocols/socks.py

   1 # -*- test-case-name: twisted.test.test_socks -*-
   2 # Copyright (c) Twisted Matrix Laboratories.
   3 # See LICENSE for details.
   4
   5 """
   6 Implementation of the SOCKSv4 protocol.
   7 """
   8
   9 # python imports
  10 import struct
  11 import string
  12 import socket
  13 import time
  14
  15 # twisted imports
  16 from twisted.internet import reactor, protocol, defer
  17 from twisted.python import log
  18
  19
  20 class SOCKSv4Outgoing(protocol.Protocol):
  21
  22     def __init__(self,socks):
  23         self.socks=socks
  24
  25     def connectionMade(self):
  26         peer = self.transport.getPeer()
  27         self.socks.makeReply(90, 0, port=peer.port, ip=peer.host)
  28         self.socks.otherConn=self
  29
  30     def connectionLost(self, reason):
  31         self.socks.transport.loseConnection()
  32
  33     def dataReceived(self,data):
  34         self.socks.write(data)
  35
  36     def write(self,data):
  37         self.socks.log(self,data)
  38         self.transport.write(data)
  39
  40
  41
  42 class SOCKSv4Incoming(protocol.Protocol):
  43
  44     def __init__(self,socks):
  45         self.socks=socks
  46         self.socks.otherConn=self
  47
  48     def connectionLost(self, reason):
  49         self.socks.transport.loseConnection()
  50
  51     def dataReceived(self,data):
  52         self.socks.write(data)
  53
  54     def write(self,data):
  55         self.socks.log(self,data)
  56         self.transport.write(data)
  57
  58
  59 class SOCKSv4(protocol.Protocol):
  60     """
  61     An implementation of the SOCKSv4 protocol.
  62
  63     @type logging: C{str} or C{None}
  64     @ivar logging: If not C{None}, the name of the logfile to which connection
  65         information will be written.
  66
  67     @type reactor: object providing L{twisted.internet.interfaces.IReactorTCP}
  68     @ivar reactor: The reactor used to create connections.
  69
  70     @type buf: C{str}
  71     @ivar buf: Part of a SOCKSv4 connection request.
  72
  73     @type otherConn: C{SOCKSv4Incoming}, C{SOCKSv4Outgoing} or C{None}
  74     @ivar otherConn: Until the connection has been established, C{otherConn} is
  75         C{None}. After that, it is the proxy-to-destination protocol instance
  76         along which the client's connection is being forwarded.
  77     """
  78     def __init__(self, logging=None, reactor=reactor):
  79         self.logging = logging
  80         self.reactor = reactor
  81
  82     def connectionMade(self):
  83         self.buf = ""
  84         self.otherConn = None
  85
  86     def dataReceived(self, data):
  87         """
  88         Called whenever data is received.
  89
  90         @type data: C{str}
  91         @param data: Part or all of a SOCKSv4 packet.
  92         """
  93         if self.otherConn:
  94             self.otherConn.write(data)
  95             return
  96         self.buf = self.buf + data
  97         completeBuffer = self.buf
  98         if "\000" in self.buf[8:]:
  99             head, self.buf = self.buf[:8], self.buf[8:]
 100             version, code, port = struct.unpack("!BBH", head[:4])
 101             user, self.buf = self.buf.split("\000", 1)
 102             if head[4:7] == "\000\000\000" and head[7] != "\000":
 103                 # An IP address of the form 0.0.0.X, where X is non-zero,
 104                 # signifies that this is a SOCKSv4a packet.
 105                 # If the complete packet hasn't been received, restore the
 106                 # buffer and wait for it.
 107                 if "\000" not in self.buf:
 108                     self.buf = completeBuffer
 109                     return
 110                 server, self.buf = self.buf.split("\000", 1)
 111                 d = self.reactor.resolve(server)
 112                 d.addCallback(self._dataReceived2, user,
 113                               version, code, port)
 114                 d.addErrback(lambda result, self = self: self.makeReply(91))
 115                 return
 116             else:
 117                 server = socket.inet_ntoa(head[4:8])
 118
 119             self._dataReceived2(server, user, version, code, port)
 120
 121     def _dataReceived2(self, server, user, version, code, port):
 122         """
 123         The second half of the SOCKS connection setup. For a SOCKSv4 packet this
 124         is after the server address has been extracted from the header. For a
 125         SOCKSv4a packet this is after the host name has been resolved.
 126
 127         @type server: C{str}
 128         @param server: The IP address of the destination, represented as a
 129             dotted quad.
 130
 131         @type user: C{str}
 132         @param user: The username associated with the connection.
 133
 134         @type version: C{int}
 135         @param version: The SOCKS protocol version number.
 136
 137         @type code: C{int}
 138         @param code: The comand code. 1 means establish a TCP/IP stream
 139             connection, and 2 means establish a TCP/IP port binding.
 140
 141         @type port: C{int}
 142         @param port: The port number associated with the connection.
 143         """
 144         assert version == 4, "Bad version code: %s" % version
 145         if not self.authorize(code, server, port, user):
 146             self.makeReply(91)
 147             return
 148         if code == 1: # CONNECT
 149             d = self.connectClass(server, port, SOCKSv4Outgoing, self)
 150             d.addErrback(lambda result, self = self: self.makeReply(91))
 151         elif code == 2: # BIND
 152             d = self.listenClass(0, SOCKSv4IncomingFactory, self, server)
 153             d.addCallback(lambda (h, p),
 154                           self = self: self.makeReply(90, 0, p, h))
 155         else:
 156             raise RuntimeError, "Bad Connect Code: %s" % code
 157         assert self.buf == "", "hmm, still stuff in buffer... %s" % repr(
 158             self.buf)
 159
 160     def connectionLost(self, reason):
 161         if self.otherConn:
 162             self.otherConn.transport.loseConnection()
 163
 164     def authorize(self,code,server,port,user):
 165         log.msg("code %s connection to %s:%s (user %s) authorized" % (code,server,port,user))
 166         return 1
 167
 168     def connectClass(self, host, port, klass, *args):
 169         return protocol.ClientCreator(reactor, klass, *args).connectTCP(host,port)
 170
 171     def listenClass(self, port, klass, *args):
 172         serv = reactor.listenTCP(port, klass(*args))
 173         return defer.succeed(serv.getHost()[1:])
 174
 175     def makeReply(self,reply,version=0,port=0,ip="0.0.0.0"):
 176         self.transport.write(struct.pack("!BBH",version,reply,port)+socket.inet_aton(ip))
 177         if reply!=90: self.transport.loseConnection()
 178
 179     def write(self,data):
 180         self.log(self,data)
 181         self.transport.write(data)
 182
 183     def log(self,proto,data):
 184         if not self.logging: return
 185         peer = self.transport.getPeer()
 186         their_peer = self.otherConn.transport.getPeer()
 187         f=open(self.logging,"a")
 188         f.write("%s\t%s:%d %s %s:%d\n"%(time.ctime(),
 189                                         peer.host,peer.port,
 190                                         ((proto==self and '<') or '>'),
 191                                         their_peer.host,their_peer.port))
 192         while data:
 193             p,data=data[:16],data[16:]
 194             f.write(string.join(map(lambda x:'%02X'%ord(x),p),' ')+' ')
 195             f.write((16-len(p))*3*' ')
 196             for c in p:
 197                 if len(repr(c))>3: f.write('.')
 198                 else: f.write(c)
 199             f.write('\n')
 200         f.write('\n')
 201         f.close()
 202
 203
 204
 205 class SOCKSv4Factory(protocol.Factory):
 206     """
 207     A factory for a SOCKSv4 proxy.
 208
 209     Constructor accepts one argument, a log file name.
 210     """
 211
 212     def __init__(self, log):
 213         self.logging = log
 214
 215     def buildProtocol(self, addr):
 216         return SOCKSv4(self.logging, reactor)
 217
 218
 219
 220 class SOCKSv4IncomingFactory(protocol.Factory):
 221     """
 222     A utility class for building protocols for incoming connections.
 223     """
 224
 225     def __init__(self, socks, ip):
 226         self.socks = socks
 227         self.ip = ip
 228
 229
 230     def buildProtocol(self, addr):
 231         if addr[0] == self.ip:
 232             self.ip = ""
 233             self.socks.makeReply(90, 0)
 234             return SOCKSv4Incoming(self.socks)
 235         elif self.ip == "":
 236             return None
 237         else:
 238             self.socks.makeReply(91, 0)
 239             self.ip = ""
 240             return None