commit 8dc796d142bebb6d370188f7e5be6bf65e402743 Author: Damian Johnson <atagar@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()