commit 7f89ff58d2ad6ecbde7f8225e7b2c02ec400be1b Author: Damian Johnson <atagar@torproject.org> Date: Mon Nov 21 08:36:20 2011 -0800 Function to query PROTOCOLINFO via a control port Implementation and integ testing to query a PROTOCOLINFO response via a control port. Next is to do the same for control sockets. --- run_tests.py | 2 +- stem/connection.py | 74 ++++++++++++++++++++++++++++++++- test/integ/connection/protocolinfo.py | 63 ++++++++++++++++++++-------- 3 files changed, 120 insertions(+), 19 deletions(-) diff --git a/run_tests.py b/run_tests.py index 2d12a57..76e1abb 100755 --- a/run_tests.py +++ b/run_tests.py @@ -41,7 +41,7 @@ UNIT_TESTS = (("stem.types.ControlMessage", test.unit.types.control_message.Test ) INTEG_TESTS = (("stem.types.ControlMessage", test.integ.types.control_message.TestControlMessage), - ("stem.connection.ProtocolInfoResponse", test.integ.connection.protocolinfo.TestProtocolInfoResponse), + ("stem.connection.ProtocolInfoResponse", test.integ.connection.protocolinfo.TestProtocolInfo), ("stem.util.conf", test.integ.util.conf.TestConf), ("stem.util.system", test.integ.util.system.TestSystem), ) diff --git a/stem/connection.py b/stem/connection.py index 3165fba..3e97e9e 100644 --- a/stem/connection.py +++ b/stem/connection.py @@ -28,6 +28,78 @@ LOGGER = logging.getLogger("stem") AuthMethod = stem.util.enum.Enum("NONE", "PASSWORD", "COOKIE", "UNKNOWN") +def get_protocolinfo_port(control_addr = "127.0.0.1", control_port = 9051, keep_alive = False): + """ + Issues a PROTOCOLINFO query to a control port, getting information about the + tor process running on it. + + Arguments: + control_addr (str) - ip address of the controller + control_port (int) - port number of the controller + keep_alive (bool) - keeps the socket used to issue the PROTOCOLINFO query + open if True, closes otherwise + + Returns: + ProtocolInfoResponse with the response given by the tor process + + Raises: + stem.types.ProtocolError if the PROTOCOLINFO response is malformed + stem.types.SocketError if problems arise in establishing or using the + socket + """ + + control_socket = socket.socket(socket.AF_INET, socket.SOCK_STREAM) + control_socket_file = control_socket.makefile() + protocolinfo_response, raised_exc = None, None + + try: + # initiates connection + control_socket.connect((control_addr, control_port)) + + # issues the PROTOCOLINFO query + control_socket_file.write("PROTOCOLINFO 1\r\n") + control_socket_file.flush() + + protocolinfo_response = stem.types.read_message(control_socket_file) + ProtocolInfoResponse.convert(protocolinfo_response) + except socket.error, exc: + raised_exc = stem.types.SocketError(exc) + except (stem.types.ProtocolError, stem.types.SocketError), exc: + raised_exc = exc + + control_socket_file.close() # done with the linked file + + if not keep_alive or raised_exc: + # shut down the socket we were using + try: control_socket.shutdown(socket.SHUT_RDWR) + except socket.error: pass + + control_socket.close() + else: + # if we're keeping the socket open then attach it to the response + protocolinfo_response.socket = control_socket + + if raised_exc: + raise raised_exc + else: + # if we have an unresolved cookie path then try to expand it again, using + # our port to infer a pid + + cookie_path = protocolinfo_response.cookie_file + if cookie_path and control_addr == "127.0.0.1" and stem.util.system.is_relative_path(cookie_path): + try: + tor_pid = stem.util.system.get_pid_by_port(control_port) + if not tor_pid: raise IOError("pid lookup failed") + + tor_cwd = stem.util.system.get_cwd(tor_pid) + if not tor_cwd: raise IOError("cwd lookup failed") + + protocolinfo_response.cookie_file = stem.util.system.expand_path(cookie_path, tor_cwd) + except IOError, exc: + LOGGER.debug("unable to expand relative tor cookie path by port: %s" % exc) + + return protocolinfo_response + class ProtocolInfoResponse(stem.types.ControlMessage): """ Version one PROTOCOLINFO query response. @@ -166,7 +238,7 @@ class ProtocolInfoResponse(stem.types.ControlMessage): self.cookie_file = stem.util.system.expand_path(self.cookie_file, tor_cwd) except IOError, exc: - LOGGER.debug("unable to expand relative tor cookie path: %s" % exc) + LOGGER.debug("unable to expand relative tor cookie path by name: %s" % exc) elif line_type == "VERSION": # Line format: # VersionLine = "250-VERSION" SP "Tor=" TorVersion OptArguments CRLF diff --git a/test/integ/connection/protocolinfo.py b/test/integ/connection/protocolinfo.py index 27b5f1f..73a7884 100644 --- a/test/integ/connection/protocolinfo.py +++ b/test/integ/connection/protocolinfo.py @@ -1,20 +1,21 @@ """ -Integration tests for the stem.connections.ProtocolInfoResponse class. +Integration tests for the stem.connections.ProtocolInfoResponse class and +related functions. """ -import socket import unittest import test.runner import stem.types import stem.connection -class TestProtocolInfoResponse(unittest.TestCase): +class TestProtocolInfo(unittest.TestCase): """ - Processes a ProtocolInfo query for a variety of setups. + Queries and parses PROTOCOLINFO. This should be run with the 'CONNECTION' + integ target to exercise the widest range of use cases. """ - def testProtocolInfoResponse(self): + def test_parsing(self): """ Makes a PROTOCOLINFO query and processes the response for our control connection. @@ -29,7 +30,7 @@ class TestProtocolInfoResponse(unittest.TestCase): control_socket = runner.get_tor_socket(False) control_socket_file = control_socket.makefile() - control_socket_file.write("PROTOCOLINFO\r\n") + control_socket_file.write("PROTOCOLINFO 1\r\n") control_socket_file.flush() protocolinfo_response = stem.types.read_message(control_socket_file) @@ -42,24 +43,52 @@ class TestProtocolInfoResponse(unittest.TestCase): self.assertNotEqual(None, protocolinfo_response.tor_version) self.assertNotEqual(None, protocolinfo_response.auth_methods) - self.assertEqual((), protocolinfo_response.unknown_auth_methods) self.assertEqual(None, protocolinfo_response.socket) + self.assert_protocolinfo_attr(protocolinfo_response, connection_type) + + def test_get_protocolinfo_port(self): + """ + Exercises the stem.connection.get_protocolinfo_port function. + """ + + connection_type = test.runner.get_runner().get_connection_type() + + if test.runner.OPT_PORT in test.runner.CONNECTION_OPTS[connection_type]: + protocolinfo_response = stem.connection.get_protocolinfo_port(control_port = test.runner.CONTROL_PORT) + self.assertEqual(None, protocolinfo_response.socket) + self.assert_protocolinfo_attr(protocolinfo_response, connection_type) + else: + # we don't have a control port + self.assertRaises(stem.types.SocketError, stem.connection.get_protocolinfo_port, "127.0.0.1", test.runner.CONTROL_PORT) + + def assert_protocolinfo_attr(self, protocolinfo_response, connection_type): + """ + Makes assertions that the protocolinfo response's attributes match those of + a given connection type. + """ + + # This should never have test.runner.TorConnection.NONE. If we somehow got + # a protocolinfo_response from that config then we have an issue. :) if connection_type == test.runner.TorConnection.NO_AUTH: - self.assertEqual((stem.connection.AuthMethod.NONE,), protocolinfo_response.auth_methods) - self.assertEqual(None, protocolinfo_response.cookie_file) + auth_methods = (stem.connection.AuthMethod.NONE,) elif connection_type == test.runner.TorConnection.PASSWORD: - self.assertEqual((stem.connection.AuthMethod.PASSWORD,), protocolinfo_response.auth_methods) - self.assertEqual(None, protocolinfo_response.cookie_file) + auth_methods = (stem.connection.AuthMethod.PASSWORD,) elif connection_type == test.runner.TorConnection.COOKIE: - self.assertEqual((stem.connection.AuthMethod.COOKIE,), protocolinfo_response.auth_methods) - self.assertEqual(runner.get_auth_cookie_path(), protocolinfo_response.cookie_file) + auth_methods = (stem.connection.AuthMethod.COOKIE,) elif connection_type == test.runner.TorConnection.MULTIPLE: - self.assertEqual((stem.connection.AuthMethod.COOKIE, stem.connection.AuthMethod.PASSWORD), protocolinfo_response.auth_methods) - self.assertEqual(runner.get_auth_cookie_path(), protocolinfo_response.cookie_file) + auth_methods = (stem.connection.AuthMethod.COOKIE, stem.connection.AuthMethod.PASSWORD) elif connection_type == test.runner.TorConnection.SOCKET: - self.assertEqual((stem.connection.AuthMethod.NONE,), protocolinfo_response.auth_methods) - self.assertEqual(None, protocolinfo_response.cookie_file) + auth_methods = (stem.connection.AuthMethod.NONE,) else: self.fail("Unrecognized connection type: %s" % connection_type) + + self.assertEqual((), protocolinfo_response.unknown_auth_methods) + self.assertEqual(auth_methods, protocolinfo_response.auth_methods) + + if test.runner.OPT_COOKIE in test.runner.CONNECTION_OPTS[connection_type]: + auth_cookie_path = test.runner.get_runner().get_auth_cookie_path() + self.assertEqual(auth_cookie_path, protocolinfo_response.cookie_file) + else: + self.assertEqual(None, protocolinfo_response.cookie_file)