[tor-commits] [stem/master] Add get_socks_ports method to Controller

[atagar at torproject.org]

commit 3a2875f22d1bf5dba6171e98807ef22f2a02c729
Author: Sean Robinson <seankrobinson at gmail.com>
Date:   Sat Dec 22 06:16:56 2012 -0700

    Add get_socks_ports method to Controller
    
    Adds a method to find the SOCKS listeners on a controlled tor.  This
    method is tested on tor versions 0.2.1.32 - 0.2.4.6-alpha and provides a
    similar response regardless of the underlying tor process.  Also includes
    unit and integration tests for the new method.
    
    Signed-off-by: Sean Robinson <seankrobinson at gmail.com>
---
 stem/control.py                  |   27 +++++++++++++++++++
 test/integ/control/controller.py |   11 ++++++++
 test/unit/control/controller.py  |   53 +++++++++++++++++++++++++++++++++++++-
 3 files changed, 90 insertions(+), 1 deletions(-)

diff --git a/stem/control.py b/stem/control.py
index 15b762e..8a36285 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -1320,6 +1320,33 @@ class Controller(BaseController):
     else:
       raise stem.ProtocolError("SAVECONF returned unexpected response code")
   
+  def get_socks_ports(self):
+    """
+    Returns a list of SOCKS proxy ports open on the controlled tor instance.
+    
+    :returns: list of **(host, port)** tuples or an empty list if there are no
+      SOCKS listeners
+    """
+    
+    try:
+      raw_addrs = self.get_info("net/listeners/socks").split()
+      # remove the surrounding quotes from each listener
+      raw_addrs = [x.replace("\"", "") for x in raw_addrs]
+    except stem.InvalidArguments:
+      # tor version is old (pre-tor-0.2.2.26-beta); use get_conf()
+      socks_port = self.get_conf('SocksPort')
+      raw_addrs = []
+      for listener in self.get_conf('SocksListenAddress', multiple = True):
+        if listener.count(':') == 0:
+          listener = listener + ":" + socks_port
+        raw_addrs.append(listener)
+    # both processes above give a list of strings of the form host:port
+    proxy_addrs = []
+    for proxy in raw_addrs:
+      proxy_pair = proxy.split(":")
+      proxy_addrs.append(tuple((proxy_pair[0], int(proxy_pair[1]))))
+    return proxy_addrs
+  
   def is_feature_enabled(self, feature):
     """
     Checks if a control connection feature is enabled. These features can be
diff --git a/test/integ/control/controller.py b/test/integ/control/controller.py
index 4388cdc..5c146e0 100644
--- a/test/integ/control/controller.py
+++ b/test/integ/control/controller.py
@@ -458,6 +458,17 @@ class TestController(unittest.TestCase):
         controller.load_conf(oldconf)
         controller.save_conf()
   
+  def test_get_socks_ports(self):
+    """
+    Test Controller.get_socks_ports against a running tor instance.
+    """
+    if test.runner.require_control(self): return
+    
+    runner = test.runner.get_runner()
+    
+    with runner.get_tor_controller() as controller:
+      self.assertEqual([('127.0.0.1', 1112)], controller.get_socks_ports())
+  
   def test_enable_feature(self):
     """
     Test Controller.enable_feature with valid and invalid inputs.
diff --git a/test/unit/control/controller.py b/test/unit/control/controller.py
index 400d795..f3849eb 100644
--- a/test/unit/control/controller.py
+++ b/test/unit/control/controller.py
@@ -8,7 +8,7 @@ import unittest
 import stem.socket
 import stem.version
 
-from stem import InvalidRequest, ProtocolError
+from stem import InvalidArguments, InvalidRequest, ProtocolError
 from stem.control import _parse_circ_path, Controller, EventType
 from test import mocking
 
@@ -76,4 +76,55 @@ class TestControl(unittest.TestCase):
     
     # EventType.SIGNAL was added in tor version 0.2.3.1-alpha
     self.assertRaises(InvalidRequest, self.controller.add_event_listener, mocking.no_op(), EventType.SIGNAL)
+  
+  def test_socks_port_old_tor(self):
+    """
+    Exercises the get_socks_ports method as if talking to an old tor process.
+    """
+    
+    # An old tor raises stem.InvalidArguments for get_info about socks, but
+    #  get_socks_ports returns the socks information, anyway.
+    mocking.mock_method(Controller, "get_info", mocking.raise_exception(InvalidArguments))
+    mocking.mock_method(Controller, "get_conf", mocking.return_for_args({
+      ("SocksPort",): "9050",
+      ("SocksListenAddress", "multiple=True"): ["127.0.0.1"]
+    }, method = True))
+    self.assertEqual([('127.0.0.1', 9050)], self.controller.get_socks_ports())
+    
+    # Again, an old tor, but SocksListenAddress overrides the port number.
+    mocking.mock_method(Controller, "get_conf", mocking.return_for_args({
+      ("SocksPort",): "9050",
+      ("SocksListenAddress", "multiple=True"): ["127.0.0.1:1112"]
+    }, method = True))
+    self.assertEqual([('127.0.0.1', 1112)], self.controller.get_socks_ports())
+    
+    # Again, an old tor, but multiple listeners
+    mocking.mock_method(Controller, "get_conf", mocking.return_for_args({
+      ("SocksPort",): "9050",
+      ("SocksListenAddress", "multiple=True"): ["127.0.0.1:1112", "127.0.0.1:1114"]
+    }, method = True))
+    self.assertEqual([('127.0.0.1', 1112), ('127.0.0.1', 1114)], self.controller.get_socks_ports())
+    
+    # Again, an old tor, but no SOCKS listeners
+    mocking.mock_method(Controller, "get_conf", mocking.return_for_args({
+      ("SocksPort",): "0",
+      ("SocksListenAddress", "multiple=True"): []
+    }, method = True))
+    self.assertEqual([], self.controller.get_socks_ports())
+  
+  def test_socks_port_new_tor(self):
+    """
+    Exercises the get_socks_ports method as if talking to a newer tor process.
+    """
+    
+    # multiple SOCKS listeners
+    mocking.mock_method(Controller, "get_info", mocking.return_value(
+      "\"127.0.0.1:1112\" \"127.0.0.1:1114\""
+    ))
+    self.assertEqual([('127.0.0.1', 1112), ('127.0.0.1', 1114)],
+        self.controller.get_socks_ports())
+    
+    # no SOCKS listeners
+    mocking.mock_method(Controller, "get_info", mocking.return_value(""))
+    self.assertEqual([], self.controller.get_socks_ports())