[tor-commits] [stem/master] Adding controller methods to check NEWNYM availability

[atagar at torproject.org]

commit c510c1b01599781b249f12c06062082309ea5f0e
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Jan 26 19:15:51 2014 -0800

    Adding controller methods to check NEWNYM availability
    
    Adding a couple methods to check if tor will respect a NEWNYM signal or not.
    This only works if signals are only sent through stem.
---
 docs/change_log.rst              |    1 +
 stem/control.py                  |   38 +++++++++++++++++++++++++++++++++++++-
 test/integ/control/controller.py |   17 +++++++++++++++++
 3 files changed, 55 insertions(+), 1 deletion(-)

diff --git a/docs/change_log.rst b/docs/change_log.rst
index fa8b638..ab52e5f 100644
--- a/docs/change_log.rst
+++ b/docs/change_log.rst
@@ -41,6 +41,7 @@ The following are only available within stem's `git repository
 
  * **Controller**
 
+  * Added :func:`~stem.control.Controller.is_newnym_available` and :func:`~stem.control.Controller.get_newnym_wait` methods to the :class:`~stem.control.Controller`
   * Added the id attribute to the :class:`~stem.response.events.ORConnEvent` (:spec:`6f2919a`)
   * Added `support for CONN_BW events <api/response.html#stem.response.events.ConnectionBandwidthEvent>`_ (:spec:`6f2919a`)
   * Added `support for CIRC_BW events <api/response.html#stem.response.events.CircuitBandwidthEvent>`_ (:spec:`6f2919a`)
diff --git a/stem/control.py b/stem/control.py
index adfc56f..22a1962 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -64,6 +64,8 @@ providing its own for interacting at a higher level.
     |- close_stream - close a stream
     |
     |- signal - sends a signal to the tor client
+    |- is_newnym_available - true if tor would presently accept a NEWNYM signal
+    |- get_newnym_wait - seconds until tor would accept a NEWNYM signal
     |- is_geoip_unavailable - true if we've discovered our geoip db to be unavailable
     +- map_address - maps one address to another such that connections to the original are replaced with the other
 
@@ -679,6 +681,7 @@ class Controller(BaseController):
 
     self._is_caching_enabled = True
     self._request_cache = {}
+    self._last_newnym = 0.0
 
     self._cache_lock = threading.RLock()
 
@@ -688,6 +691,7 @@ class Controller(BaseController):
     self._event_listeners_lock = threading.RLock()
 
     # number of sequential 'GETINFO ip-to-country/*' lookups that have failed
+
     self._geoip_failure_count = 0
     self._enabled_features = []
 
@@ -2258,12 +2262,44 @@ class Controller(BaseController):
     response = self.msg("SIGNAL %s" % signal)
     stem.response.convert("SINGLELINE", response)
 
-    if not response.is_ok():
+    if response.is_ok():
+      if signal == stem.Signal.NEWNYM:
+        self._last_newnym = time.time()
+    else:
       if response.code == "552":
         raise stem.InvalidArguments(response.code, response.message, [signal])
 
       raise stem.ProtocolError("SIGNAL response contained unrecognized status code: %s" % response.code)
 
+  def is_newnym_available(self):
+    """
+    Indicates if tor would presently accept a NEWNYM signal. This can only
+    account for signals sent via this controller.
+
+    .. versionadded:: 1.2.0
+
+    :returns: **True** if tor would presently accept a NEWNYM signal, **False**
+      otherwise
+    """
+
+    if self.is_alive():
+      return self.get_newnym_wait() == 0.0
+    else:
+      return False
+
+  def get_newnym_wait(self):
+    """
+    Provides the number of seconds until a NEWNYM signal would be respected.
+    This can only account for signals sent via this controller.
+
+    .. versionadded:: 1.2.0
+
+    :returns: **float** for the number of seconds until tor would respect
+      another NEWNYM signal
+    """
+
+    return max(0.0, self._last_newnym + 10 - time.time())
+
   def is_geoip_unavailable(self):
     """
     Provides **True** if we've concluded hat our geoip database is unavailable,
diff --git a/test/integ/control/controller.py b/test/integ/control/controller.py
index 1c0e468..d034853 100644
--- a/test/integ/control/controller.py
+++ b/test/integ/control/controller.py
@@ -654,6 +654,23 @@ class TestController(unittest.TestCase):
       # invalid signals
       self.assertRaises(stem.InvalidArguments, controller.signal, "FOOBAR")
 
+  def test_newnym_availability(self):
+    """
+    Test the is_newnym_available and get_newnym_wait methods.
+    """
+
+    if test.runner.require_control(self):
+      return
+
+    with test.runner.get_runner().get_tor_controller() as controller:
+      self.assertEqual(True, controller.is_newnym_available())
+      self.assertEqual(0.0, controller.get_newnym_wait())
+
+      controller.signal(stem.Signal.NEWNYM)
+
+      self.assertEqual(False, controller.is_newnym_available())
+      self.assertTrue(controller.get_newnym_wait() > 9.0)
+
   def test_extendcircuit(self):
     if test.runner.require_control(self):
       return