[tor-commits] [stem/master] Function for writting to control sockets
atagar at torproject.org
atagar at torproject.org
Wed Nov 23 18:06:36 UTC 2011
commit 8dc796d142bebb6d370188f7e5be6bf65e402743
Author: Damian Johnson <atagar at torproject.org>
Date: Wed Nov 23 09:55:16 2011 -0800
Function for writting to control sockets
Writing directly to the socket file isn't hard (it's just a write and flush).
However, this is nicer since it wrap the control formatting, logging, and
exception quirks. Functions still need unit tests and I might just wrap the
socket object completely...
---
stem/connection.py | 3 +-
stem/types.py | 80 +++++++++++++++++++++++++++++++-
test/integ/connection/protocolinfo.py | 4 +-
test/integ/types/control_message.py | 48 ++++++-------------
test/runner.py | 7 +--
5 files changed, 97 insertions(+), 45 deletions(-)
diff --git a/stem/connection.py b/stem/connection.py
index 476c9c3..5fc06c1 100644
--- a/stem/connection.py
+++ b/stem/connection.py
@@ -107,8 +107,7 @@ def _get_protocolinfo_impl(control_socket, connection_args, keep_alive):
control_socket.connect(connection_args)
# issues the PROTOCOLINFO query
- control_socket_file.write("PROTOCOLINFO 1\r\n")
- control_socket_file.flush()
+ stem.types.write_message(control_socket_file, "PROTOCOLINFO 1")
protocolinfo_response = stem.types.read_message(control_socket_file)
ProtocolInfoResponse.convert(protocolinfo_response)
diff --git a/stem/types.py b/stem/types.py
index 882ca99..66aa205 100644
--- a/stem/types.py
+++ b/stem/types.py
@@ -7,6 +7,8 @@ ControllerError - Base exception raised when using the controller.
|- SocketError - Socket used for controller communication errored.
+- SocketClosed - Socket terminated.
+write_message - Writes a message to a control socket.
+format_write_message - Performs the formatting expected from sent messages.
read_message - Reads a ControlMessage from a control socket.
ControlMessage - Message from the control socket.
|- content - provides the parsed message content
@@ -58,18 +60,90 @@ class SocketError(ControllerError):
"Error arose while communicating with the control socket."
pass
-class SocketClosed(ControllerError):
+class SocketClosed(SocketError):
"Control socket was closed before completing the message."
pass
+def write_message(control_file, message, raw = False):
+ """
+ Sends a message to the control socket, adding the expected formatting for
+ single verses multiline messages. Neither message type should contain an
+ ending newline (if so it'll be treated as a multi-line message with a blank
+ line at the end). If the message doesn't contain a newline then it's sent
+ as...
+
+ <message>\r\n
+
+ and if it does contain newlines then it's split on \n and sent as...
+
+ +<line 1>\r\n
+ <line 2>\r\n
+ <line 3>\r\n
+ .\r\n
+
+ Arguments:
+ control_file (file) - file derived from the control socket (see the
+ socket's makefile() method for more information)
+ message (str) - message to be sent on the control socket
+ raw (bool) - leaves the message formatting untouched, passing it
+ to the socket as-is
+
+ Raises:
+ SocketError if a problem arises in using the socket
+ """
+
+ if not raw: message = format_write_message(message)
+
+ try:
+ LOGGER.debug("Sending message:\n" + message.replace("\r\n", "\n").rstrip())
+ control_file.write(message)
+ control_file.flush()
+ except socket.error, exc:
+ LOGGER.info("Failed to send message: %s" % exc)
+ raise SocketError(exc)
+ except AttributeError:
+ # This happens after the file's close() method has been called, the flush
+ # causing...
+ # AttributeError: 'NoneType' object has no attribute 'sendall'
+
+ LOGGER.info("Failed to send message: file has been closed")
+ raise SocketError("file has been closed")
+
+def format_write_message(message):
+ """
+ Performs the formatting expected of control messages (for more information
+ see the write_message function).
+
+ Arguments:
+ message (str) - message to be formatted
+
+ Returns:
+ str of the message wrapped by the formatting expected from controllers
+ """
+
+ # From 'Commands from controller to Tor' (section 2.2) of the control spec...
+ #
+ # Command = Keyword OptArguments CRLF / "+" Keyword OptArguments CRLF CmdData
+ # Keyword = 1*ALPHA
+ # OptArguments = [ SP *(SP / VCHAR) ]
+ #
+ # A command is either a single line containing a Keyword and arguments, or a
+ # multiline command whose initial keyword begins with +, and whose data
+ # section ends with a single "." on a line of its own.
+
+ if "\n" in message:
+ return "+%s\r\n.\r\n" % message.replace("\n", "\r\n")
+ else:
+ return message + "\r\n"
+
def read_message(control_file):
"""
Pulls from a control socket until we either have a complete message or
encounter a problem.
Arguments:
- control_file - file derived from the control socket (see the socket's
- makefile() method for more information)
+ control_file (file) - file derived from the control socket (see the
+ socket's makefile() method for more information)
Returns:
stem.types.ControlMessage read from the socket
diff --git a/test/integ/connection/protocolinfo.py b/test/integ/connection/protocolinfo.py
index 1035fe6..e2e9a7b 100644
--- a/test/integ/connection/protocolinfo.py
+++ b/test/integ/connection/protocolinfo.py
@@ -31,9 +31,7 @@ class TestProtocolInfo(unittest.TestCase):
control_socket = runner.get_tor_socket(False)
control_socket_file = control_socket.makefile()
- control_socket_file.write("PROTOCOLINFO 1\r\n")
- control_socket_file.flush()
-
+ stem.types.write_message(control_socket_file, "PROTOCOLINFO 1")
protocolinfo_response = stem.types.read_message(control_socket_file)
stem.connection.ProtocolInfoResponse.convert(protocolinfo_response)
diff --git a/test/integ/types/control_message.py b/test/integ/types/control_message.py
index 4aa878f..392cd25 100644
--- a/test/integ/types/control_message.py
+++ b/test/integ/types/control_message.py
@@ -27,8 +27,7 @@ class TestControlMessage(unittest.TestCase):
# PROTOCOLINFO then tor will give an 'Authentication required.' message and
# hang up.
- control_socket_file.write("GETINFO version\r\n")
- control_socket_file.flush()
+ stem.types.write_message(control_socket_file, "GETINFO version")
auth_required_response = stem.types.read_message(control_socket_file)
self.assertEquals("Authentication required.", str(auth_required_response))
@@ -39,11 +38,10 @@ class TestControlMessage(unittest.TestCase):
# The socket's broken but doesn't realize it yet. Send another message and
# it should fail with a closed exception. With a control port we won't get
# an error until we read from the socket. However, with a control socket
- # the flush will raise a socket.error.
+ # the write will cause a SocketError.
try:
- control_socket_file.write("GETINFO version\r\n")
- control_socket_file.flush()
+ stem.types.write_message(control_socket_file, "GETINFO version")
except: pass
self.assertRaises(stem.types.SocketClosed, stem.types.read_message, control_socket_file)
@@ -51,8 +49,7 @@ class TestControlMessage(unittest.TestCase):
# Additional socket usage should fail, and pulling more responses will fail
# with more closed exceptions.
- control_socket_file.write("GETINFO version\r\n")
- self.assertRaises(socket.error, control_socket_file.flush)
+ self.assertRaises(stem.types.SocketError, stem.types.write_message, control_socket_file, "GETINFO version")
self.assertRaises(stem.types.SocketClosed, stem.types.read_message, control_socket_file)
self.assertRaises(stem.types.SocketClosed, stem.types.read_message, control_socket_file)
self.assertRaises(stem.types.SocketClosed, stem.types.read_message, control_socket_file)
@@ -61,22 +58,17 @@ class TestControlMessage(unittest.TestCase):
# an impact.
control_socket.close()
- control_socket_file.write("GETINFO version\r\n")
- self.assertRaises(socket.error, control_socket_file.flush)
+ self.assertRaises(stem.types.SocketError, stem.types.write_message, control_socket_file, "GETINFO version")
self.assertRaises(stem.types.SocketClosed, stem.types.read_message, control_socket_file)
- # Closing the file handler, however, will cause a different type of error.
- # This seems to depend on the python version, in 2.6 we get an
- # AttributeError and in 2.7 the close() call raises...
- # error: [Errno 32] Broken pipe
+ # Tries again with the file explicitely closed. In python 2.7 the close
+ # call will raise...
+ # error: [Errno 32] Broken pipe
- try:
- control_socket_file.close()
- control_socket_file.write("GETINFO version\r\n")
+ try: control_socket_file.close()
except: pass
- # receives: AttributeError: 'NoneType' object has no attribute 'sendall'
- self.assertRaises(AttributeError, control_socket_file.flush)
+ self.assertRaises(stem.types.SocketError, stem.types.write_message, control_socket_file, "GETINFO version")
# receives: stem.types.SocketClosed: socket file has been closed
self.assertRaises(stem.types.SocketClosed, stem.types.read_message, control_socket_file)
@@ -90,9 +82,7 @@ class TestControlMessage(unittest.TestCase):
if not control_socket: self.skipTest("(no control socket)")
control_socket_file = control_socket.makefile()
- control_socket_file.write("blarg\r\n")
- control_socket_file.flush()
-
+ stem.types.write_message(control_socket_file, "blarg")
unrecognized_command_response = stem.types.read_message(control_socket_file)
self.assertEquals('Unrecognized command "blarg"', str(unrecognized_command_response))
self.assertEquals(['Unrecognized command "blarg"'], list(unrecognized_command_response))
@@ -111,9 +101,7 @@ class TestControlMessage(unittest.TestCase):
if not control_socket: self.skipTest("(no control socket)")
control_socket_file = control_socket.makefile()
- control_socket_file.write("GETINFO blarg\r\n")
- control_socket_file.flush()
-
+ stem.types.write_message(control_socket_file, "GETINFO blarg")
unrecognized_key_response = stem.types.read_message(control_socket_file)
self.assertEquals('Unrecognized key "blarg"', str(unrecognized_key_response))
self.assertEquals(['Unrecognized key "blarg"'], list(unrecognized_key_response))
@@ -135,9 +123,7 @@ class TestControlMessage(unittest.TestCase):
if not control_socket: self.skipTest("(no control socket)")
control_socket_file = control_socket.makefile()
- control_socket_file.write("GETINFO config-file\r\n")
- control_socket_file.flush()
-
+ stem.types.write_message(control_socket_file, "GETINFO config-file")
config_file_response = stem.types.read_message(control_socket_file)
self.assertEquals("config-file=%s\nOK" % torrc_dst, str(config_file_response))
self.assertEquals(["config-file=%s" % torrc_dst, "OK"], list(config_file_response))
@@ -174,9 +160,7 @@ class TestControlMessage(unittest.TestCase):
if not control_socket: self.skipTest("(no control socket)")
control_socket_file = control_socket.makefile()
- control_socket_file.write("GETINFO config-text\r\n")
- control_socket_file.flush()
-
+ stem.types.write_message(control_socket_file, "GETINFO config-text")
config_text_response = stem.types.read_message(control_socket_file)
# the response should contain two entries, the first being a data response
@@ -206,9 +190,7 @@ class TestControlMessage(unittest.TestCase):
if not control_socket: self.skipTest("(no control socket)")
control_socket_file = control_socket.makefile()
- control_socket_file.write("SETEVENTS BW\r\n")
- control_socket_file.flush()
-
+ stem.types.write_message(control_socket_file, "SETEVENTS BW")
setevents_response = stem.types.read_message(control_socket_file)
self.assertEquals("OK", str(setevents_response))
self.assertEquals(["OK"], list(setevents_response))
diff --git a/test/runner.py b/test/runner.py
index e1c9f0b..bc8e8ec 100644
--- a/test/runner.py
+++ b/test/runner.py
@@ -341,13 +341,12 @@ class Runner:
auth_cookie_contents = auth_cookie.read()
auth_cookie.close()
- control_socket_file.write("AUTHENTICATE %s\r\n" % binascii.b2a_hex(auth_cookie_contents))
+ stem.types.write_message(control_socket_file, "AUTHENTICATE %s" % binascii.b2a_hex(auth_cookie_contents))
elif OPT_PASSWORD in conn_opts:
- control_socket_file.write("AUTHENTICATE \"%s\"\r\n" % CONTROL_PASSWORD)
+ stem.types.write_message(control_socket_file, "AUTHENTICATE \"%s\"" % CONTROL_PASSWORD)
else:
- control_socket_file.write("AUTHENTICATE\r\n")
+ stem.types.write_message(control_socket_file, "AUTHENTICATE")
- control_socket_file.flush()
authenticate_response = stem.types.read_message(control_socket_file)
control_socket_file.close()
More information about the tor-commits
mailing list