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

[atagar at torproject.org]

commit d682c8cd13a0f6bffcd056a2a938db2dfa65f500
Author: Sean Robinson <seankrobinson at gmail.com>
Date:   Thu Nov 22 08:09:42 2012 -0700

    Add a close_circuit method to Controller
    
    Signed-off-by: Sean Robinson <seankrobinson at gmail.com>
---
 stem/control.py                  |   21 +++++++++++++++++++++
 test/integ/control/controller.py |   26 ++++++++++++++++++++++++++
 2 files changed, 47 insertions(+), 0 deletions(-)

diff --git a/stem/control.py b/stem/control.py
index a2d408b..0b7baee 100644
--- a/stem/control.py
+++ b/stem/control.py
@@ -29,6 +29,7 @@ providing its own for interacting at a higher level.
     |- signal - sends a signal to the tor client
     |- new_circuit - create new circuits
     |- extend_circuit - create new circuits and extend existing ones
+    |- close_circuit - close a circuit
     |- repurpose_circuit - change a circuit's purpose
     |- map_address - maps one address to another such that connections to the original are replaced with the other
     |- get_version - convenience method to get tor version
@@ -1268,6 +1269,26 @@ class Controller(BaseController):
     
     return int(new_circuit)
   
+  def close_circuit(self, circuit_id, flag = ''):
+    """
+    Closes the specified circuit.
+
+    :param int circuit_id: id of the circuit to be closed
+    :param str flag: optional value to modify closing, the only flag available
+      is "IfUnused" which will not close the circuit unless it is unused
+
+    :raises: :class:`stem.InvalidArguments` if the circuit doesn't exist or not enough information is provided
+    """
+
+    response = self.msg("CLOSECIRCUIT %s %s"% (str(circuit_id), flag))
+    stem.response.convert("SINGLELINE", response)
+
+    if not response.is_ok():
+      if response.code in ('512', '552'):
+        raise stem.InvalidRequest(response.code, response.message)
+      else:
+        raise stem.ProtocolError("CLOSECIRCUIT returned unexpected response code: %s" % response.code)
+
   def map_address(self, mapping):
     """
     Map addresses to replacement addresses. Tor replaces subseqent connections
diff --git a/test/integ/control/controller.py b/test/integ/control/controller.py
index 571d3ab..58208b5 100644
--- a/test/integ/control/controller.py
+++ b/test/integ/control/controller.py
@@ -406,6 +406,32 @@ class TestController(unittest.TestCase):
       self.assertRaises(stem.InvalidRequest, controller.repurpose_circuit, 'f934h9f3h4', "fooo")
       self.assertRaises(stem.InvalidRequest, controller.repurpose_circuit, '4', "fooo")
   
+  def test_close_circuit(self):
+    """
+    Tests Controller.close_circuit with valid and invalid input.
+    """
+
+    if test.runner.require_control(self): return
+    elif test.runner.require_online(self): return
+
+    runner = test.runner.get_runner()
+
+    with runner.get_tor_controller() as controller:
+      circ_id = controller.new_circuit()
+      controller.close_circuit(circ_id)
+      circuit_output = controller.get_info("circuit-status")
+      circ = filter(re.compile("^%i " % circ_id).match, circuit_output.splitlines())[0]
+      self.assertFalse(circ_id in circ)
+
+      circ_id = controller.new_circuit()
+      controller.close_circuit(circ_id, "IfUnused")
+      circuit_output = controller.get_info("circuit-status")
+      circ = filter(re.compile("^%i " % circ_id).match, circuit_output.splitlines())[0]
+      self.assertFalse(circ_id in circ)
+
+      circ_id = controller.new_circuit()
+      self.assertRaises(stem.InvalidRequest, controller.close_circuit, circ_id + 1024)
+
   def test_mapaddress(self):
     if test.runner.require_control(self): return
     elif test.runner.require_online(self): return