Rename tor-control.py to TorControl.py; begin making it into a useful library instead of a lame testing script.

svn:r3879

2 files changed, 309 insertions, 255 deletions

diff --git a/contrib/TorControl.py b/contrib/TorControl.py
new file mode 100755
index 0000000000..155603644d
--- /dev/null
+++ b/contrib/TorControl.py
@@ -0,0 +1,309 @@
+#!/usr/bin/python
+#$Id$
+
+import socket
+import struct
+import sys
+
+class _Enum:
+    def __init__(self, start, names):
+        self.nameOf = {}
+        idx = start
+        for name in names:
+            setattr(self,name,idx)
+            self.nameOf[idx] = name
+            idx += 1
+
+
+MSG_TYPE = _Enum(0x0000,
+                 ["ERROR",
+                  "DONE",
+                  "SETCONF",
+                  "GETCONF",
+                  "CONFVALUE",
+                  "SETEVENTS",
+                  "EVENT",
+                  "AUTH",
+                  "SAVECONF",
+                  "SIGNAL",
+                  "MAPADDRESS",
+                  "GETINFO",
+                  "INFOVALUE",
+                  "EXTENDCIRCUIT",
+                  "ATTACHSTREAM",
+                  "POSTDESCRIPTOR",
+                  "FRAGMENTHEADER",
+                  "FRAGMENT",
+                  "REDIRECTSTREAM",
+                  "CLOSESTREAM",
+                  "CLOSECIRCUIT",
+                  ])
+
+assert MSG_TYPE.SAVECONF = 0x0008
+assert MSG_TYPE.CLOSECIRCUIT = 0x0014
+
+EVENT_TYPE = _ENUM(0x0001,
+                   ["CIRCSTATUS",
+                    "STREAMSTATUS",
+                    "ORCONNSTATUS",
+                    "BANDWIDTH",
+                    "WARN",
+                    "NEWDESC"])
+
+ERR_CODES = {
+  0x0000 : "Unspecified error",
+  0x0001 : "Internal error",
+  0x0002 : "Unrecognized message type",
+  0x0003 : "Syntax error",
+  0x0004 : "Unrecognized configuration key",
+  0x0005 : "Invalid configuration value",
+  0x0006 : "Unrecognized byte code",
+  0x0007 : "Unauthorized",
+  0x0008 : "Failed authentication attempt",
+  0x0009 : "Resource exhausted",
+  0x000A : "No such stream",
+  0x000B : "No such circuit",
+  0x000C : "No such OR"
+}
+
+class TorCtlError(Exception):
+  pass
+
+class ProtocolError(TorCtlError):
+  pass
+
+class ErrorReply(TorCtlError):
+  pass
+
+def parseHostAndPort(h):
+    host, port = "localhost", 9051
+    if ":" in h:
+        i = h.index(":")
+        host = h[:i]
+        try:
+            port = int(h[i+1:])
+        except ValueError:
+            print "Bad hostname %r"%h
+            sys.exit(1)
+    elif h:
+        try:
+            port = int(h)
+        except ValueError:
+            host = h
+
+    return host, port
+
+
+def _unpack_msg(msg):
+    "return None, minLength, body or type,body,rest"
+    if len(msg) < 4:
+        return None, 4, msg
+    length,type = struct.unpack("!HH",msg)
+    if len(msg) >= 4+length:
+        return type,msg[4:4+length],msg[4+length:]
+    else:
+        return None,4+length,msg
+
+def unpack_msg(msg):
+    "returns as for _unpack_msg"
+    tp,body,rest = _unpack_msg(msg)
+    if tp != MSG_TYPE.FRAGMENTHEADER:
+        return tp, body, rest
+
+    if len(body) < 6:
+        raise ProtocolError("FRAGMENTHEADER message too short")
+
+    realType,realLength = struct.unpack("!HL", body[:6])
+
+    # Okay; could the message _possibly_ be here?
+    minPackets,minSlop = divmod(realLength+6,65535)
+    minLength = (minPackets*(65535+4))+4+minSlop
+    if len(msg) < minLength:
+        return None,  minLength, msg
+
+    # Okay; optimistically try to build up the msg.
+    soFar = [ body[6:] ]
+    lenSoFarLen = len(body)-6
+    while rest and lenSoFar < realLength:
+        ln, tp = struct.unpack("!HH" rest[:4])
+        if tp != MSG_TYPE.FRAGMENT:
+            raise ProtocolError("Missing FRAGMENT message")
+        soFar.append(rest[4:4+ln])
+        lenSoFar += ln
+        rest = rest[4+ln:]
+
+    if lenSoFar == realLength:
+        return realType, "".join(soFar), rest
+    elif lenSoFar > realLength:
+        raise ProtocolError("Bad fragmentation: message longer than declared")
+    else:
+        return None, len(msg)+(realLength-lenSoFar), msg
+
+def receive_message(s):
+    length, tp, body = _receive_msg(s)
+    if tp != MSG_TYPE.FRAGMENTHEADER:
+        return length, tp, body
+    if length < 6:
+        raise ProtocolError("FRAGMENTHEADER message too short")
+    realType,realLength = struct.unpack("!HL", body[:6])
+    data = [ body[6:] ]
+    soFar = len(data[0])
+    while 1:
+        length, tp, body = _receive_msg(s)
+        if tp != MSG_TYPE.FRAGMENT:
+            raise ProtocolError("Missing FRAGMENT message")
+        soFar += length
+        data.append(body)
+        if soFar == realLength:
+            return realLength, realType, "".join(data)
+        elif soFar > realLengtH:
+            raise ProtocolError("FRAGMENT message too long!")
+
+_event_handler = None
+def receive_reply(s, expected=None):
+    while 1:
+        _, tp, body = receive_message(s)
+        if tp == MSG_TYPE.EVENT:
+            if _event_handler is not None:
+                _event_handler(tp, body)
+        elif tp == MSG_TYPE.ERROR:
+            if len(body)<2:
+                raise ProtocolError("(Truncated error message)")
+            errCode, = struct.unpack("!H", body[:2])
+            raise ErrorReply((errCode,
+                              ERR_CODES.get(errCode,"[unrecognized]"),
+                              body[2:]))
+        elif (expected is not None) and (tp not in expected):
+            raise ProtocolError("Unexpected message type 0x%04x"%tp)
+        else:
+            return tp, body
+
+def pack_message(type, body=""):
+    length = len(body)
+    if length < 65536:
+        reqheader = struct.pack("!HH", length, type)
+        return "%s%s"%(reqheader,body)
+
+    fragheader = struct.pack("!HHHL",
+                             65535, MSG_TYPE.FRAGMENTHEADER, type, length)
+    msgs = [ fragheader, body[:65535-6] ]
+    body = body[65535-6:]
+    while body:
+        if len(body) > 65535:
+            fl = 65535
+        else:
+            fl = len(body)
+        fragheader = struct.pack("!HH", MSG_TYPE.FRAGMENT, fl)
+        msgs.append(fragheader)
+        msgs.append(body[:fl])
+        body = body[fl:]
+
+    return "".join(msgs)
+
+def send_message(s, type, body=""):
+    s.sendall(pack_message(type, body))
+
+def authenticate(s):
+    send_message(s,MSG_TYPE.AUTH)
+    type,body = receive_reply(s)
+    return
+
+def _parseKV(body,sep=" ",term="\n"):
+    res = []
+    for line in body.split(term):
+        if not line: continue
+        print repr(line)
+        k, v = line.split(sep,1)
+        res.append((k,v))
+    return res
+
+def get_option(s,name):
+    send_message(s,MSG_TYPE.GETCONF,name)
+    tp,body = receive_reply(s,[MSG_TYPE.CONFVALUE])
+    return _parseKV(body)
+
+def set_option(s,msg):
+    send_message(s,MSG_TYPE.SETCONF,msg)
+    tp,body = receive_reply(s,[MSG_TYPE.DONE])
+
+def get_info(s,name):
+    send_message(s,MSG_TYPE.GETINFO,name)
+    tp,body = receive_reply(s,[MSG_TYPE.INFOVALUE])
+    kvs = body.split("\0")
+    d = {}
+    for i in xrange(0,len(kvs)-1,2):
+        d[kvs[i]] = kvs[i+1]
+    return d
+
+def set_events(s,events):
+    send_message(s,MSG_TYPE.SETEVENTS,
+                 "".join([struct.pack("!H", event) for event in events]))
+    type,body = receive_reply(s,[MSG_TYPE.DONE])
+    return
+
+def save_conf(s):
+    send_message(s,MSG_TYPE.SAVECONF)
+    receive_reply(s,[MSG_TYPE.DONE])
+
+def send_signal(s, sig):
+    send_message(s,MSG_TYPE.SIGNAL,struct.pack("B",sig))
+    receive_reply(s,[MSG_TYPE.DONE])
+
+def map_address(s, kv):
+    msg = [ "%s %s\n"%(k,v) for k,v in kv ]
+    send_message(s,MSG_TYPE.MAPADDRESS,"".join(msg))
+    tp, body = receive_reply(s,[MSG_TYPE.DONE])
+    return _parseKV(body)
+
+def extend_circuit(s, circid, hops):
+    msg = struct.pack("!L",circid) + ",".join(hops) + "\0"
+    send_message(s,MSG_TYPE.EXTENDCIRCUIT,msg)
+    tp, body = receive_reply(s,[MSG_TYPE.DONE])
+    return body
+
+def listen_for_events(s):
+    while(1):
+        _,type,body = receive_message(s)
+        print "event",type
+    return
+
+def do_main_loop(host,port):
+    print "host is %s:%d"%(host,port)
+    s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
+    s.connect((host,port))
+    authenticate(s)
+    print "nick",`get_option(s,"nickname")`
+    print get_option(s,"DirFetchPeriod\n")
+    print `get_info(s,"version")`
+    #print `get_info(s,"desc/name/moria1")`
+    print `get_info(s,"network-status")`
+    print `get_info(s,"addr-mappings/all")`
+    print `get_info(s,"addr-mappings/config")`
+    print `get_info(s,"addr-mappings/cache")`
+    print `get_info(s,"addr-mappings/control")`
+    print `map_address(s, [("0.0.0.0", "Foobar.com"),
+                           ("1.2.3.4", "foobaz.com"),
+                           ("frebnitz.com", "5.6.7.8"),
+                           (".", "abacinator.onion")])`
+    print `extend_circuit(s,0,["moria1"])`
+    send_signal(s,1)
+    #save_conf(s)
+
+
+    #set_option(s,"1")
+    #set_option(s,"bandwidthburstbytes 100000")
+    #set_option(s,"runasdaemon 1")
+    #set_events(s,[EVENT_TYPE.WARN])
+    set_events(s,[EVENT_TYPE.WARN,EVENT_TYPE.STREAMSTATUS])
+
+    listen_for_events(s)
+
+    return
+
+if __name__ == '__main__':
+    if len(sys.argv) != 2:
+        print "Syntax: tor-control.py torhost:torport"
+        sys.exit(0)
+    sh,sp = parseHostAndPort(sys.argv[1])
+    do_main_loop(sh,sp)
+
diff --git a/contrib/tor-control.py b/contrib/tor-control.py
deleted file mode 100755
index 72333a4c3e..0000000000
--- a/contrib/tor-control.py
+++ /dev/null
@@ -1,255 +0,0 @@
-#!/usr/bin/python2
-#$Id$
-
-import socket
-import struct
-import sys
-
-MSG_TYPE_ERROR     = 0x0000
-MSG_TYPE_DONE      = 0x0001
-MSG_TYPE_SETCONF   = 0x0002
-MSG_TYPE_GETCONF   = 0x0003
-MSG_TYPE_CONFVALUE = 0x0004
-MSG_TYPE_SETEVENTS = 0x0005
-MSG_TYPE_EVENT     = 0x0006
-MSG_TYPE_AUTH      = 0x0007
-MSG_TYPE_SAVECONF  = 0x0008
-MSG_TYPE_SIGNAL    = 0x0009
-MSG_TYPE_MAPADDRESS     = 0x000A
-MSG_TYPE_GETINFO        = 0x000B
-MSG_TYPE_INFOVALUE      = 0x000C
-MSG_TYPE_EXTENDCIRCUIT  = 0x000D
-MSG_TYPE_ATTACHSTREAM   = 0x000E
-MSG_TYPE_POSTDESCRIPTOR = 0x000F
-MSG_TYPE_FRAGMENTHEADER = 0x0010
-MSG_TYPE_FRAGMENT       = 0x0011
-MSG_TYPE_REDIRECTSTREAM = 0x0012
-MSG_TYPE_CLOSESTREAM    = 0x0013
-MSG_TYPE_CLOSECIRCUIT   = 0x0014
-
-EVENT_TYPE_CIRCSTATUS   = 0x0001
-EVENT_TYPE_STREAMSTATUS = 0x0002
-EVENT_TYPE_ORCONNSTATUS = 0x0003
-EVENT_TYPE_BANDWIDTH    = 0x0004
-EVENT_TYPE_WARN         = 0x0005
-EVENT_TYPE_NEWDESC      = 0x0006
-
-ERR_CODES = {
-  0x0000 : "Unspecified error",
-  0x0001 : "Internal error",
-  0x0002 : "Unrecognized message type",
-  0x0003 : "Syntax error",
-  0x0004 : "Unrecognized configuration key",
-  0x0005 : "Invalid configuration value",
-  0x0006 : "Unrecognized byte code",
-  0x0007 : "Unauthorized",
-  0x0008 : "Failed authentication attempt",
-  0x0009 : "Resource exhausted",
-  0x000A : "No such stream",
-  0x000B : "No such circuit",
-  0x000C : "No such OR"
-}
-
-class TorCtlError(Exception):
-  pass
-
-class ProtocolError(TorCtlError):
-  pass
-
-class ErrorReply(TorCtlError):
-  pass
-
-def parseHostAndPort(h):
-    host, port = "localhost", 9051
-    if ":" in h:
-        i = h.index(":")
-        host = h[:i]
-        try:
-            port = int(h[i+1:])
-        except ValueError:
-            print "Bad hostname %r"%h
-            sys.exit(1)
-    elif h:
-        try:
-            port = int(h)
-        except ValueError:
-            host = h
-
-    return host, port
-
-def _receive_msg(s):
-  body = ""
-  header = s.recv(4)
-  length,type = struct.unpack("!HH",header)
-  if length:
-    body = s.recv(length)
-  return length,type,body
-
-def receive_message(s):
-  length, tp, body = _receive_msg(s)
-  if tp != MSG_TYPE_FRAGMENTHEADER:
-    return length, tp, body
-  if length < 6:
-    raise ProtocolError("FRAGMENTHEADER message too short")
-  realType,realLength = struct.unpack("!HL", body[:6])
-  data = [ body[6:] ]
-  soFar = len(data[0])
-  while 1:
-    length, tp, body = _receive_msg(s)
-    if tp != MSG_TYPE_FRAGMENT:
-      raise ProtocolError("Missing FRAGMENT message")
-    soFar += length
-    data.append(body)
-    if soFar == realLength:
-      return realLength, realType, "".join(data)
-    elif soFar > realLengtH:
-      raise ProtocolError("FRAGMENT message too long!")
-
-_event_handler = None
-def receive_reply(s, expected=None):
-  while 1:
-    _, tp, body = receive_message(s)
-    if tp == MSG_TYPE_EVENT:
-      if _event_handler is not None:
-        _event_handler(tp, body)
-    elif tp == MSG_TYPE_ERROR:
-      if len(body)<2:
-        raise ProtocolError("(Truncated error message)")
-      errCode, = struct.unpack("!H", body[:2])
-      raise ErrorReply((errCode,
-                        ERR_CODES.get(errCode,"[unrecognized]"),
-                        body[2:]))
-    elif (expected is not None) and (tp not in expected):
-      raise ProtocolError("Unexpected message type 0x%04x"%tp)
-    else:
-      return tp, body
-
-def pack_message(type, body=""):
-  length = len(body)
-  if length < 65536:
-    reqheader = struct.pack("!HH", length, type)
-    return "%s%s"%(reqheader,body)
-
-  fragheader = struct.pack("!HHHL",
-                           65535, MSG_TYPE_FRAGMENTHEADER, type, length)
-  msgs = [ fragheader, body[:65535-6] ]
-  body = body[65535-6:]
-  while body:
-    if len(body) > 65535:
-      fl = 65535
-    else:
-      fl = len(body)
-    fragheader = struct.pack("!HH", MSG_TYPE_FRAGMENT, fl)
-    msgs.append(fragheader)
-    msgs.append(body[:fl])
-    body = body[fl:]
-
-  return "".join(msgs)
-
-def send_message(s, type, body=""):
-  s.sendall(pack_message(type, body))
-
-def authenticate(s):
-  send_message(s,MSG_TYPE_AUTH)
-  type,body = receive_reply(s)
-  return
-
-def _parseKV(body,sep=" ",term="\n"):
-  res = []
-  for line in body.split(term):
-    if not line: continue
-    print repr(line)
-    k, v = line.split(sep,1)
-    res.append((k,v))
-  return res
-
-def get_option(s,name):
-  send_message(s,MSG_TYPE_GETCONF,name)
-  tp,body = receive_reply(s,[MSG_TYPE_CONFVALUE])
-  return _parseKV(body)
-
-def set_option(s,msg):
-  send_message(s,MSG_TYPE_SETCONF,msg)
-  tp,body = receive_reply(s,[MSG_TYPE_DONE])
-
-def get_info(s,name):
-  send_message(s,MSG_TYPE_GETINFO,name)
-  tp,body = receive_reply(s,[MSG_TYPE_INFOVALUE])
-  kvs = body.split("\0")
-  d = {}
-  for i in xrange(0,len(kvs)-1,2):
-    d[kvs[i]] = kvs[i+1]
-  return d
-
-def set_events(s,events):
-  send_message(s,MSG_TYPE_SETEVENTS,
-               "".join([struct.pack("!H", event) for event in events]))
-  type,body = receive_reply(s,[MSG_TYPE_DONE])
-  return
-
-def save_conf(s):
-  send_message(s,MSG_TYPE_SAVECONF)
-  receive_reply(s,[MSG_TYPE_DONE])
-
-def send_signal(s, sig):
-  send_message(s,MSG_TYPE_SIGNAL,struct.pack("B",sig))
-  receive_reply(s,[MSG_TYPE_DONE])
-
-def map_address(s, kv):
-  msg = [ "%s %s\n"%(k,v) for k,v in kv ]
-  send_message(s,MSG_TYPE_MAPADDRESS,"".join(msg))
-  tp, body = receive_reply(s,[MSG_TYPE_DONE])
-  return _parseKV(body)
-
-def extend_circuit(s, circid, hops):
-  msg = struct.pack("!L",circid) + ",".join(hops) + "\0"
-  send_message(s,MSG_TYPE_EXTENDCIRCUIT,msg)
-  tp, body = receive_reply(s,[MSG_TYPE_DONE])
-  return body
-
-def listen_for_events(s):
-  while(1):
-    _,type,body = receive_message(s)
-    print "event",type
-  return
-
-def do_main_loop(host,port):
-  print "host is %s:%d"%(host,port)
-  s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
-  s.connect((host,port))
-  authenticate(s)
-  print "nick",`get_option(s,"nickname")`
-  print get_option(s,"DirFetchPeriod\n")
-  print `get_info(s,"version")`
-  #print `get_info(s,"desc/name/moria1")`
-  print `get_info(s,"network-status")`
-  print `get_info(s,"addr-mappings/all")`
-  print `get_info(s,"addr-mappings/config")`
-  print `get_info(s,"addr-mappings/cache")`
-  print `get_info(s,"addr-mappings/control")`
-  print `map_address(s, [("0.0.0.0", "Foobar.com"),
-                         ("1.2.3.4", "foobaz.com"),
-                         ("frebnitz.com", "5.6.7.8"),
-                         (".", "abacinator.onion")])`
-  print `extend_circuit(s,0,["moria1"])`
-  send_signal(s,1)
-  #save_conf(s)
-
-
-  #set_option(s,"1")
-  #set_option(s,"bandwidthburstbytes 100000")
-  #set_option(s,"runasdaemon 1")
-  #set_events(s,[EVENT_TYPE_WARN])
-  set_events(s,[EVENT_TYPE_WARN,EVENT_TYPE_STREAMSTATUS])
-
-  listen_for_events(s)
-
-  return
-
-if __name__ == '__main__':
-  if len(sys.argv) != 2:
-    print "Syntax: tor-control.py torhost:torport"
-    sys.exit(0)
-  sh,sp = parseHostAndPort(sys.argv[1])
-  do_main_loop(sh,sp)
-