← Back to branch summary

~landscape/zope3/ztk-1.1.3

« back to all changes in this revision

Viewing changes to src/twisted/protocols/socks.py

  • Committer: Andreas Hasenack
  • Date: 2009-07-20 17:49:16 UTC
  • Revision ID: andreas@canonical.com-20090720174916-g2tn6qmietz2hn0u
Revert twisted removal, it breaks several dozen tests [trivial]

Show diffs side-by-side

added added

removed removed

Lines of Context:
src/twisted/protocols/socks.py
 
1
# Copyright (c) 2001-2004 Twisted Matrix Laboratories.
 
2
# See LICENSE for details.
 
3
 
 
4
 
 
5
"""
 
6
Implementation of the SOCKSv4 protocol.
 
7
"""
 
8
 
 
9
# twisted imports
 
10
from twisted.internet import reactor, protocol, defer
 
11
from twisted.python import log
 
12
 
 
13
# python imports
 
14
import struct
 
15
import string
 
16
import socket
 
17
import time
 
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
class SOCKSv4Incoming(protocol.Protocol):
 
42
    
 
43
    def __init__(self,socks):
 
44
        self.socks=socks
 
45
        self.socks.otherConn=self
 
46
    
 
47
    def connectionLost(self, reason):
 
48
        self.socks.transport.loseConnection()
 
49
    
 
50
    def dataReceived(self,data):
 
51
        self.socks.write(data)
 
52
    
 
53
    def write(self,data):
 
54
        self.socks.log(self,data)
 
55
        self.transport.write(data)
 
56
 
 
57
 
 
58
class SOCKSv4(protocol.Protocol):
 
59
    
 
60
    def __init__(self,logging=None):
 
61
        self.logging=logging
 
62
    
 
63
    def connectionMade(self):
 
64
        self.buf=""
 
65
        self.otherConn=None
 
66
 
 
67
    def dataReceived(self,data):
 
68
        if self.otherConn:
 
69
            self.otherConn.write(data)
 
70
            return
 
71
        self.buf=self.buf+data
 
72
        if '\000' in self.buf[8:]:
 
73
            head,self.buf=self.buf[:8],self.buf[8:]
 
74
            try:
 
75
                version,code,port=struct.unpack("!BBH",head[:4])
 
76
            except struct.error:
 
77
                raise RuntimeError, "struct error with head='%s' and buf='%s'"%(repr(head),repr(self.buf))
 
78
            user,self.buf=string.split(self.buf,"\000",1)
 
79
            if head[4:7]=="\000\000\000": # domain is after
 
80
                server,self.buf=string.split(self.buf,'\000',1)
 
81
                #server=gethostbyname(server)
 
82
            else:
 
83
                server=socket.inet_ntoa(head[4:8])
 
84
            assert version==4, "Bad version code: %s"%version
 
85
            if not self.authorize(code,server,port,user):
 
86
                self.makeReply(91)
 
87
                return
 
88
            if code==1: # CONNECT
 
89
                d = self.connectClass(server, port, SOCKSv4Outgoing, self)
 
90
                d.addErrback(lambda result, self=self: self.makeReply(91))
 
91
            elif code==2: # BIND
 
92
                ip = socket.gethostbyname(server)
 
93
                d = self.listenClass(0, SOCKSv4IncomingFactory, self, ip)
 
94
                d.addCallback(lambda (h, p), self=self: self.makeReply(90, 0, p, h))
 
95
            else:
 
96
                raise RuntimeError, "Bad Connect Code: %s" % code
 
97
            assert self.buf=="","hmm, still stuff in buffer... %s" % repr(self.buf)
 
98
 
 
99
    def connectionLost(self, reason):
 
100
        if self.otherConn:
 
101
            self.otherConn.transport.loseConnection()
 
102
 
 
103
    def authorize(self,code,server,port,user):
 
104
        log.msg("code %s connection to %s:%s (user %s) authorized" % (code,server,port,user))
 
105
        return 1
 
106
 
 
107
    def connectClass(self, host, port, klass, *args):
 
108
        return protocol.ClientCreator(reactor, klass, *args).connectTCP(host,port)
 
109
 
 
110
    def listenClass(self, port, klass, *args):
 
111
        serv = reactor.listenTCP(port, klass(*args))
 
112
        return defer.succeed(serv.getHost()[1:])
 
113
 
 
114
    def makeReply(self,reply,version=0,port=0,ip="0.0.0.0"):
 
115
        self.transport.write(struct.pack("!BBH",version,reply,port)+socket.inet_aton(ip))
 
116
        if reply!=90: self.transport.loseConnection()
 
117
 
 
118
    def write(self,data):
 
119
        self.log(self,data)
 
120
        self.transport.write(data)
 
121
 
 
122
    def log(self,proto,data):
 
123
        if not self.logging: return
 
124
        peer = self.transport.getPeer()
 
125
        their_peer = self.otherConn.transport.getPeer()
 
126
        f=open(self.logging,"a")
 
127
        f.write("%s\t%s:%d %s %s:%d\n"%(time.ctime(),
 
128
                                        peer.host,peer.port,
 
129
                                        ((proto==self and '<') or '>'),
 
130
                                        their_peer.host,their_peer.port))
 
131
        while data:
 
132
            p,data=data[:16],data[16:]
 
133
            f.write(string.join(map(lambda x:'%02X'%ord(x),p),' ')+' ')
 
134
            f.write((16-len(p))*3*' ')
 
135
            for c in p:
 
136
                if len(repr(c))>3: f.write('.')
 
137
                else: f.write(c)
 
138
            f.write('\n')
 
139
        f.write('\n')
 
140
        f.close()
 
141
 
 
142
 
 
143
class SOCKSv4Factory(protocol.Factory):
 
144
    """A factory for a SOCKSv4 proxy.
 
145
    
 
146
    Constructor accepts one argument, a log file name.
 
147
    """
 
148
    
 
149
    def __init__(self, log):
 
150
        self.logging = log
 
151
    
 
152
    def buildProtocol(self, addr):
 
153
        return SOCKSv4(self.logging)
 
154
 
 
155
 
 
156
class SOCKSv4IncomingFactory(protocol.Factory):
 
157
    """A utility class for building protocols for incoming connections."""
 
158
    
 
159
    def __init__(self, socks, ip):
 
160
        self.socks = socks
 
161
        self.ip = ip
 
162
    
 
163
    def buildProtocol(self, addr):
 
164
        if addr[0] == self.ip:
 
165
            self.ip = ""
 
166
            self.socks.makeReply(90, 0)
 
167
            return SOCKSv4Incoming(self.socks)
 
168
        elif self.ip == "":
 
169
            return None
 
170
        else:
 
171
            self.socks.makeReply(91, 0)
 
172
            self.ip = ""
 
173
            return None