contrib/tor-resolve.py

   1 #!/usr/bin/python
   2 #$Id$
   3
   4 import socket
   5 import struct
   6 import sys
   7
   8 def socks4AResolveRequest(hostname):
   9     version = 4
  10     command = 0xF0
  11     port = 0
  12     addr = 0x0000001
  13     username = ""
  14     reqheader = struct.pack("!BBHL", version, command, port, addr)
  15     return "%s%s\x00%s\x00"%(reqheader,username,hostname)
  16
  17 def socks4AParseResponse(response):
  18     RESPONSE_LEN = 8
  19     if len(response) < RESPONSE_LEN:
  20         print "return none", len(response)
  21         return None
  22     assert len(response) >= RESPONSE_LEN
  23     version,status,port = struct.unpack("!BBH",response[:4])
  24     assert version == 0
  25     assert port == 0
  26     if status == 90:
  27         return "%d.%d.%d.%d"%tuple(map(ord, response[4:]))
  28     else:
  29         return "ERROR (status %d)"%status
  30
  31 def socks5Hello():
  32     return "\x05\x01\x00"
  33 def socks5ParseHello(response):
  34     if response != "\x05\x00":
  35         raise ValueError("Bizarre socks5 response")
  36 def socks5ResolveRequest(hostname):
  37     version = 5
  38     command = 0xF0
  39     rsv = 0
  40     port = 0
  41     atype = 0x03
  42     reqheader = struct.pack("!BBBB",version, command, rsv, atype)
  43     portstr = struct.pack("!H",port)
  44     return "%s%s\0%s"%(reqheader,hostname,port)
  45 def socks5ParseResponse(r):
  46     if len(r)<8: return None
  47     version, reply, rsv, atype = struct.unpack("!BBBB",r[:4])
  48     assert version==5
  49     assert rsv==0
  50     if reply != 0x00:
  51         return "ERROR"
  52     assert atype in (0x01,0x04)
  53     expected_len = 4 + ({1:4,4:16}[atype]) + 2
  54     if len(r) < expected_len:
  55         return None
  56     elif len(r) > expected_len:
  57         raise ValueError("Overlong socks5 reply!")
  58     addr = r[4:-2]
  59     if atype == 0x01:
  60         return "%d.%d.%d.%d"%tuple(map(ord,addr))
  61     else:
  62         # not really the right way to format IPv6
  63         return "IPv6: %s"%(":".join([hex(ord(c)) for c in addr]))
  64
  65 def parseHostAndPort(h):
  66     host, port = "localhost", 9050
  67     if ":" in h:
  68         i = h.index(":")
  69         host = h[:i]
  70         try:
  71             port = int(h[i+1:])
  72         except ValueError:
  73             print "Bad hostname %r"%h
  74             sys.exit(1)
  75     elif h:
  76         try:
  77             port = int(h)
  78         except ValueError:
  79             host = h
  80
  81     return host, port
  82
  83 def resolve(hostname, sockshost, socksport, socksver=4):
  84     assert socksver in (4,5)
  85     if socksver == 4:
  86         fmt = socks4AResolveRequest
  87         parse = socks4AParseResponse
  88     else:
  89         fmt = socks5ResolveRequest
  90         parse = socks5ParseResponse
  91     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
  92     s.connect((sockshost,socksport))
  93     if socksver == 5:
  94         s.send(socks5Hello())
  95         socksParseHello(s.recv(2))
  96     s.send(fmt(hostname))
  97     answer = s.recv(8)
  98     result = parse(answer)
  99     while result is None:
 100         more = s.recv(1)
 101         if not more:
 102             print "Connection closed; dying."
 103             return None
 104         answer += more
 105         result = parse(answer)
 106     print "Got answer",result
 107     m = s.recv(1)
 108     if m:
 109         print "Got extra data too! Ick."
 110     return result
 111
 112 if __name__ == '__main__':
 113     if len(sys.argv) not in (2,3,4):
 114         print "Syntax: resolve.py [-4|-5] hostname [sockshost:socksport]"
 115         sys.exit(0)
 116     if sys.argv[1] in ("-4", "-5"):
 117         socksver = int(sys.argv[1][1])
 118         del sys.argv[1]
 119     if len(sys.argv) == 4:
 120         print "Syntax: resolve.py [-4|-5] hostname [sockshost:socksport]"
 121         sys.exit(0)
 122     if len(sys.argv) == 3:
 123         sh,sp = parseHostAndPort(sys.argv[2])
 124     else:
 125         sh,sp = parseHostAndPort("")
 126     resolve(sys.argv[1], sh, sp)