[tor-commits] [stem/master] Adding function for getting the system's resolvers

[atagar at torproject.org]

commit 87b0d512e2b83a2dcced709e33128eaed10cc439
Author: Damian Johnson <atagar at torproject.org>
Date:   Sun Sep 22 12:25:00 2013 -0700

    Adding function for getting the system's resolvers
    
    Adding a get_system_resolvers() function that provides the connection resolvers
    that are likely to work on a given platform.
---
 stem/util/connection.py      |   47 +++++++++++++++++++++++++++++++++++++++++-
 test/unit/util/connection.py |   27 ++++++++++++++++++++++++
 2 files changed, 73 insertions(+), 1 deletion(-)

diff --git a/stem/util/connection.py b/stem/util/connection.py
index 0e4c190..e0b08df 100644
--- a/stem/util/connection.py
+++ b/stem/util/connection.py
@@ -21,13 +21,58 @@ but for now just moving the parts we need.
 import hashlib
 import hmac
 import os
+import platform
 import re
 
-CRYPTOVARIABLE_EQUALITY_COMPARISON_NONCE = os.urandom(32)
+import stem.util.proc
+
+from stem.util import enum
+
+Resolver = enum.Enum(
+  ('PROC', 'proc'),
+  ('NETSTAT', 'netstat'),
+  ('SS', 'ss'),
+  ('LSOF', 'lsof'),
+  ('SOCKSTAT', 'sockstat'),
+  ('BSD_SOCKSTAT', 'sockstat (bsd)'),
+  ('BSD_PROCSTAT', 'procstat (bsd)')
+)
 
 FULL_IPv4_MASK = "255.255.255.255"
 FULL_IPv6_MASK = "FFFF:FFFF:FFFF:FFFF:FFFF:FFFF:FFFF:FFFF"
 
+CRYPTOVARIABLE_EQUALITY_COMPARISON_NONCE = os.urandom(32)
+
+
+def get_system_resolvers(system = None):
+  """
+  Provides the types of connection resolvers likely to be available on this platform.
+
+  :param str system: system to get resolvers for, this is determined by
+    platform.system() if not provided
+
+  :returns: **list** of Resolvers likely to be available on this platform
+  """
+
+  if system is None:
+    system = platform.system()
+
+  if system == 'Windows':
+    resolvers = []
+  elif system in ('Darwin', 'OpenBSD'):
+    resolvers = [Resolver.LSOF]
+  elif system == 'FreeBSD':
+    resolvers = [Resolver.BSD_SOCKSTAT, Resolver.BSD_PROCSTAT, Resolver.LSOF]
+  else:
+    resolvers = [Resolver.NETSTAT, Resolver.SOCKSTAT, Resolver.LSOF, Resolver.SS]
+
+  # proc resolution, by far, outperforms the others so defaults to this is able
+
+  if stem.util.proc.is_available():
+    resolvers = [Resolver.PROC] + resolvers
+
+  return resolvers
+
 
 def is_valid_ipv4_address(address):
   """
diff --git a/test/unit/util/connection.py b/test/unit/util/connection.py
index 5d48f28..40a2e02 100644
--- a/test/unit/util/connection.py
+++ b/test/unit/util/connection.py
@@ -2,12 +2,39 @@
 Unit tests for the stem.util.connection functions.
 """
 
+import platform
 import unittest
 
+from mock import patch
+
 import stem.util.connection
 
+from stem.util.connection import Resolver
 
 class TestConnection(unittest.TestCase):
+  @patch('stem.util.proc.is_available')
+  def test_get_system_resolvers(self, proc_mock):
+    """
+    Checks the get_system_resolvers function.
+    """
+
+    proc_mock.return_value = False
+
+    self.assertEqual([], stem.util.connection.get_system_resolvers('Windows'))
+    self.assertEqual([Resolver.LSOF], stem.util.connection.get_system_resolvers('Darwin'))
+    self.assertEqual([Resolver.LSOF], stem.util.connection.get_system_resolvers('OpenBSD'))
+    self.assertEqual([Resolver.BSD_SOCKSTAT, Resolver.BSD_PROCSTAT, Resolver.LSOF], stem.util.connection.get_system_resolvers('FreeBSD'))
+    self.assertEqual([Resolver.NETSTAT, Resolver.SOCKSTAT, Resolver.LSOF, Resolver.SS], stem.util.connection.get_system_resolvers('Linux'))
+
+    proc_mock.return_value = True
+
+    self.assertEqual([Resolver.PROC, Resolver.NETSTAT, Resolver.SOCKSTAT, Resolver.LSOF, Resolver.SS], stem.util.connection.get_system_resolvers('Linux'))
+
+    # check that calling without an argument is equivilant to calling for this
+    # platform
+
+    self.assertEqual(stem.util.connection.get_system_resolvers(platform.system()), stem.util.connection.get_system_resolvers())
+
   def test_is_valid_ipv4_address(self):
     """
     Checks the is_valid_ipv4_address function.