contrib/TorControl.py

   1 #!/usr/bin/python
   2 #$Id$
   3
   4 import socket
   5 import struct
   6 import sys
   7
   8 class _Enum:
   9     def __init__(self, start, names):
  10         self.nameOf = {}
  11         idx = start
  12         for name in names:
  13             setattr(self,name,idx)
  14             self.nameOf[idx] = name
  15             idx += 1
  16 class _Enum2:
  17     def __init__(self, **args):
  18         self.__dict__.update(args)
  19
  20 MSG_TYPE = _Enum(0x0000,
  21                  ["ERROR",
  22                   "DONE",
  23                   "SETCONF",
  24                   "GETCONF",
  25                   "CONFVALUE",
  26                   "SETEVENTS",
  27                   "EVENT",
  28                   "AUTH",
  29                   "SAVECONF",
  30                   "SIGNAL",
  31                   "MAPADDRESS",
  32                   "GETINFO",
  33                   "INFOVALUE",
  34                   "EXTENDCIRCUIT",
  35                   "ATTACHSTREAM",
  36                   "POSTDESCRIPTOR",
  37                   "FRAGMENTHEADER",
  38                   "FRAGMENT",
  39                   "REDIRECTSTREAM",
  40                   "CLOSESTREAM",
  41                   "CLOSECIRCUIT",
  42                   ])
  43
  44 assert MSG_TYPE.SAVECONF == 0x0008
  45 assert MSG_TYPE.CLOSECIRCUIT == 0x0014
  46
  47 EVENT_TYPE = _Enum(0x0001,
  48                    ["CIRCSTATUS",
  49                     "STREAMSTATUS",
  50                     "ORCONNSTATUS",
  51                     "BANDWIDTH",
  52                     "WARN",
  53                     "NEWDESC"])
  54
  55 CIRC_STATUS = _Enum(0x00,
  56                     ["LAUNCHED",
  57                      "BUILT",
  58                      "EXTENDED",
  59                      "FAILED",
  60                      "CLOSED"])
  61 STREAM_STATUS = _Enum(0x00,
  62                       ["SENT_CONNECT",
  63                        "SENT_RESOLVE",
  64                        "SUCCEEDED",
  65                        "FAILED",
  66                        "CLOSED",
  67                        "NEW_CONNECT",
  68                        "NEW_RESOLVE",
  69                        "DETACHED"])
  70 OR_CONN_STATUS = _Enum(0x00,
  71                        ["LAUNCHED","CONNECTED","FAILED","CLOSED"])
  72 SIGNAL = _Enum2(HUP=0x01,INT=0x02,USR1=0x0A,USR2=0x0C,TERM=0x0F)
  73
  74 ERR_CODES = {
  75   0x0000 : "Unspecified error",
  76   0x0001 : "Internal error",
  77   0x0002 : "Unrecognized message type",
  78   0x0003 : "Syntax error",
  79   0x0004 : "Unrecognized configuration key",
  80   0x0005 : "Invalid configuration value",
  81   0x0006 : "Unrecognized byte code",
  82   0x0007 : "Unauthorized",
  83   0x0008 : "Failed authentication attempt",
  84   0x0009 : "Resource exhausted",
  85   0x000A : "No such stream",
  86   0x000B : "No such circuit",
  87   0x000C : "No such OR"
  88 }
  89
  90 class TorCtlError(Exception):
  91   pass
  92
  93 class ProtocolError(TorCtlError):
  94   pass
  95
  96 class ErrorReply(TorCtlError):
  97   pass
  98
  99 def parseHostAndPort(h):
 100     host, port = "localhost", 9051
 101     if ":" in h:
 102         i = h.index(":")
 103         host = h[:i]
 104         try:
 105             port = int(h[i+1:])
 106         except ValueError:
 107             print "Bad hostname %r"%h
 108             sys.exit(1)
 109     elif h:
 110         try:
 111             port = int(h)
 112         except ValueError:
 113             host = h
 114
 115     return host, port
 116
 117 def _unpack_msg(msg):
 118     "return None, minLength, body or type,body,rest"
 119     if len(msg) < 4:
 120         return None, 4, msg
 121     length,type = struct.unpack("!HH",msg)
 122     if len(msg) >= 4+length:
 123         return type,msg[4:4+length],msg[4+length:]
 124     else:
 125         return None,4+length,msg
 126
 127 def _minLengthToPack(bytes):
 128     whole,left = divmod(bytes,65535)
 129     if left:
 130         return whole*(65535+4)+4+left
 131     else:
 132         return whole*(65535+4)
 133
 134 def unpack_msg(msg):
 135     "returns as for _unpack_msg"
 136     tp,body,rest = _unpack_msg(msg)
 137     if tp != MSG_TYPE.FRAGMENTHEADER:
 138         return tp, body, rest
 139
 140     if len(body) < 6:
 141         raise ProtocolError("FRAGMENTHEADER message too short")
 142
 143     realType,realLength = struct.unpack("!HL", body[:6])
 144
 145     # Okay; could the message _possibly_ be here?
 146     minLength = _minLengthToPack(realLength+6)
 147     if len(msg) < minLength:
 148         return None,  minLength, msg
 149
 150     # Okay; optimistically try to build up the msg.
 151     soFar = [ body[6:] ]
 152     lenSoFarLen = len(body)-6
 153     while len(rest)>=4 and lenSoFar < realLength:
 154         ln, tp = struct.unpack("!HH", rest[:4])
 155         if tp != MSG_TYPE.FRAGMENT:
 156             raise ProtocolError("Missing FRAGMENT message")
 157         soFar.append(rest[4:4+ln])
 158         lenSoFar += ln
 159         if 4+ln > len(rest):
 160             rest = ""
 161             leftInPacket = 4+ln-len(rest)
 162         else:
 163             rest = rest[4+ln:]
 164             leftInPacket=0
 165
 166     if lenSoFar == realLength:
 167         return realType, "".join(soFar), rest
 168     elif lenSoFar > realLength:
 169         raise ProtocolError("Bad fragmentation: message longer than declared")
 170     else:
 171         inOtherPackets = realLength-lenSoFar-leftInPacket
 172         minLength = _minLengthToPack(inOtherPackets)
 173         return None, len(msg)+leftInPacket+inOtherPackets, msg
 174
 175 def _receive_msg(s):
 176   body = ""
 177   header = s.recv(4)
 178   length,type = struct.unpack("!HH",header)
 179   if length:
 180     body = s.recv(length)
 181   return length,type,body
 182
 183 def receive_message(s):
 184     length, tp, body = _receive_msg(s)
 185     if tp != MSG_TYPE.FRAGMENTHEADER:
 186         return length, tp, body
 187     if length < 6:
 188         raise ProtocolError("FRAGMENTHEADER message too short")
 189     realType,realLength = struct.unpack("!HL", body[:6])
 190     data = [ body[6:] ]
 191     soFar = len(data[0])
 192     while 1:
 193         length, tp, body = _receive_msg(s)
 194         if tp != MSG_TYPE.FRAGMENT:
 195             raise ProtocolError("Missing FRAGMENT message")
 196         soFar += length
 197         data.append(body)
 198         if soFar == realLength:
 199             return realLength, realType, "".join(data)
 200         elif soFar > realLengtH:
 201             raise ProtocolError("FRAGMENT message too long!")
 202
 203 _event_handler = None
 204 def receive_reply(s, expected=None):
 205     while 1:
 206         _, tp, body = receive_message(s)
 207         if tp == MSG_TYPE.EVENT:
 208             if _event_handler is not None:
 209                 _event_handler(body)
 210         elif tp == MSG_TYPE.ERROR:
 211             if len(body)<2:
 212                 raise ProtocolError("(Truncated error message)")
 213             errCode, = struct.unpack("!H", body[:2])
 214             raise ErrorReply((errCode,
 215                               ERR_CODES.get(errCode,"[unrecognized]"),
 216                               body[2:]))
 217         elif (expected is not None) and (tp not in expected):
 218             raise ProtocolError("Unexpected message type 0x%04x"%tp)
 219         else:
 220             return tp, body
 221
 222 def pack_message(type, body=""):
 223     length = len(body)
 224     if length < 65536:
 225         reqheader = struct.pack("!HH", length, type)
 226         return "%s%s"%(reqheader,body)
 227
 228     fragheader = struct.pack("!HHHL",
 229                              65535, MSG_TYPE.FRAGMENTHEADER, type, length)
 230     msgs = [ fragheader, body[:65535-6] ]
 231     body = body[65535-6:]
 232     while body:
 233         if len(body) > 65535:
 234             fl = 65535
 235         else:
 236             fl = len(body)
 237         fragheader = struct.pack("!HH", MSG_TYPE.FRAGMENT, fl)
 238         msgs.append(fragheader)
 239         msgs.append(body[:fl])
 240         body = body[fl:]
 241
 242     return "".join(msgs)
 243
 244 def send_message(s, type, body=""):
 245     s.sendall(pack_message(type, body))
 246
 247 def authenticate(s):
 248     send_message(s,MSG_TYPE.AUTH)
 249     type,body = receive_reply(s)
 250     return
 251
 252 def _parseKV(body,sep=" ",term="\n"):
 253     res = []
 254     for line in body.split(term):
 255         if not line: continue
 256         print repr(line)
 257         k, v = line.split(sep,1)
 258         res.append((k,v))
 259     return res
 260
 261 def get_option(s,name):
 262     send_message(s,MSG_TYPE.GETCONF,name)
 263     tp,body = receive_reply(s,[MSG_TYPE.CONFVALUE])
 264     return _parseKV(body)
 265
 266 def set_option(s,msg):
 267     send_message(s,MSG_TYPE.SETCONF,msg)
 268     tp,body = receive_reply(s,[MSG_TYPE.DONE])
 269
 270 def get_info(s,name):
 271     send_message(s,MSG_TYPE.GETINFO,name)
 272     tp,body = receive_reply(s,[MSG_TYPE.INFOVALUE])
 273     kvs = body.split("\0")
 274     d = {}
 275     for i in xrange(0,len(kvs)-1,2):
 276         d[kvs[i]] = kvs[i+1]
 277     return d
 278
 279 def set_events(s,events):
 280     send_message(s,MSG_TYPE.SETEVENTS,
 281                  "".join([struct.pack("!H", event) for event in events]))
 282     type,body = receive_reply(s,[MSG_TYPE.DONE])
 283     return
 284
 285 def save_conf(s):
 286     send_message(s,MSG_TYPE.SAVECONF)
 287     receive_reply(s,[MSG_TYPE.DONE])
 288
 289 def send_signal(s, sig):
 290     send_message(s,MSG_TYPE.SIGNAL,struct.pack("B",sig))
 291     receive_reply(s,[MSG_TYPE.DONE])
 292
 293 def map_address(s, kv):
 294     msg = [ "%s %s\n"%(k,v) for k,v in kv ]
 295     send_message(s,MSG_TYPE.MAPADDRESS,"".join(msg))
 296     tp, body = receive_reply(s,[MSG_TYPE.DONE])
 297     return _parseKV(body)
 298
 299 def extend_circuit(s, circid, hops):
 300     msg = struct.pack("!L",circid) + ",".join(hops) + "\0"
 301     send_message(s,MSG_TYPE.EXTENDCIRCUIT,msg)
 302     tp, body = receive_reply(s,[MSG_TYPE.DONE])
 303     if len(body) != 4:
 304         raise ProtocolError("Extendcircuit reply too short or long")
 305     return struct.unpack("!L",body)[0]
 306
 307 def redirect_stream(s, streamid, newtarget):
 308     msg = struct.pack("!L",streamid) + newtarget + "\0"
 309     send_message(s,MSG_TYPE.REDIRECTSTREAM,msg)
 310     tp,body = receive_reply(s,[MSG_TYPE.DONE])
 311
 312 def attach_stream(s, streamid, circid):
 313     msg = struct.pack("!LL",streamid, circid)
 314     send_message(s,MSG_TYPE.ATTACHSTREAM,msg)
 315     tp,body = receive_reply(s,[MSG_TYPE.DONE])
 316
 317 def close_stream(s, streamid, reason=0, flags=0):
 318     msg = struct.pack("!LBB",streamid,reason,flags)
 319     send_message(s,MSG_TYPE.CLOSESTREAM,msg)
 320     tp,body = receive_reply(s,[MSG_TYPE.DONE])
 321
 322 def close_circuit(s, circid, flags=0):
 323     msg = struct.pack("!LB",circid,flags)
 324     send_message(s,MSG_TYPE.CLOSECIRCUIT,msg)
 325     tp,body = receive_reply(s,[MSG_TYPE.DONE])
 326
 327 def _unterminate(s):
 328     if s[-1] == '\0':
 329         return s[:-1]
 330     else:
 331         return s
 332
 333 def unpack_event(body):
 334     if len(body)<2:
 335         raise ProtocolError("EVENT body too short.")
 336     evtype, = struct.unpack("!H", body[:2])
 337     body = body[2:]
 338     if evtype == EVENT_TYPE.CIRCSTATUS:
 339         if len(body)<5:
 340             raise ProtocolError("CIRCUITSTATUS event too short.")
 341         status,ident = struct.unpack("!BL", body[:5])
 342         path = _unterminate(body[5:]).split(",")
 343         args = status, ident, path
 344     elif evtype == EVENT_TYPE.STREAMSTATUS:
 345         if len(body)<5:
 346             raise ProtocolError("CIRCUITSTATUS event too short.")
 347         status,ident = struct.unpack("!BL", body[:5])
 348         target = _unterminate(body[5:])
 349         args = status, ident, target
 350     elif evtype == EVENT_TYPE.ORCONNSTATUS:
 351         if len(body)<2:
 352             raise ProtocolError("CIRCUITSTATUS event too short.")
 353         status = ord(body[0])
 354         target = _unterminate(body[1:])
 355         args = status, target
 356     elif evtype == EVENT_TYPE.BANDWIDTH:
 357         if len(body)<8:
 358             raise ProtocolError("BANDWIDTH event too short.")
 359         read, written = struct.unpack("!LL",body[:8])
 360         args = read, written
 361     elif evtype == EVENT_TYPE.WARN:
 362         args = (_unterminate(body),)
 363     elif evtype == EVENT_TYPE.NEWDESC:
 364         args = (_unterminate(body).split(","),)
 365     else:
 366         args = (body,)
 367
 368     return evtype, args
 369
 370 def listen_for_events(s):
 371     while(1):
 372         _,type,body = receive_message(s)
 373         print "event",type
 374     return
 375
 376 def do_main_loop(host,port):
 377     print "host is %s:%d"%(host,port)
 378     s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
 379     s.connect((host,port))
 380     authenticate(s)
 381     print "nick",`get_option(s,"nickname")`
 382     print get_option(s,"DirFetchPeriod\n")
 383     print `get_info(s,"version")`
 384     #print `get_info(s,"desc/name/moria1")`
 385     print `get_info(s,"network-status")`
 386     print `get_info(s,"addr-mappings/all")`
 387     print `get_info(s,"addr-mappings/config")`
 388     print `get_info(s,"addr-mappings/cache")`
 389     print `get_info(s,"addr-mappings/control")`
 390     print `map_address(s, [("0.0.0.0", "Foobar.com"),
 391                            ("1.2.3.4", "foobaz.com"),
 392                            ("frebnitz.com", "5.6.7.8"),
 393                            (".", "abacinator.onion")])`
 394     print `extend_circuit(s,0,["moria1"])`
 395     send_signal(s,1)
 396     #save_conf(s)
 397
 398
 399     #set_option(s,"1")
 400     #set_option(s,"bandwidthburstbytes 100000")
 401     #set_option(s,"runasdaemon 1")
 402     #set_events(s,[EVENT_TYPE.WARN])
 403     set_events(s,[EVENT_TYPE.WARN,EVENT_TYPE.STREAMSTATUS])
 404
 405     listen_for_events(s)
 406
 407     return
 408
 409 if __name__ == '__main__':
 410     if len(sys.argv) != 2:
 411         print "Syntax: tor-control.py torhost:torport"
 412         sys.exit(0)
 413     sh,sp = parseHostAndPort(sys.argv[1])
 414     do_main_loop(sh,sp)
 415