commit b3aa2fb6fbf2159f56b1bab64fa2d9c1ba238947 Author: icodemachine gauthamnekk@gmail.com Date: Mon Feb 16 22:17:40 2015 +0530
Added Windows Resolver support for get_connections as well as pid_by_name --- stem/util/connection.py | 11 ++++++++--- stem/util/system.py | 13 ++++++++++--- test/integ/util/system.py | 23 +++++++++++++++++++---- test/unit/util/connection.py | 36 +++++++++++++++++++++++++++++++----- test/unit/util/system.py | 26 +++++++++++++++++++++++++- 5 files changed, 93 insertions(+), 16 deletions(-)
diff --git a/stem/util/connection.py b/stem/util/connection.py index 70ba609..93e137a 100644 --- a/stem/util/connection.py +++ b/stem/util/connection.py @@ -64,7 +64,7 @@ LOG_CONNECTION_RESOLUTION = False Resolver = enum.Enum( ('PROC', 'proc'), ('NETSTAT', 'netstat'), - ('NETSTAT_WINDOWS', 'netstat'), + ('NETSTAT_WINDOWS', 'netstat (windows)'), ('SS', 'ss'), ('LSOF', 'lsof'), ('SOCKSTAT', 'sockstat'), @@ -187,12 +187,17 @@ def get_connections(resolver, process_pid = None, process_name = None): return [Connection(*conn) for conn in stem.util.proc.connections(process_pid)]
resolver_command = RESOLVER_COMMAND[resolver].format(pid = process_pid) - + + #In case, process_name is only specified + if resolver == Resolver.NETSTAT_WINDOWS: + if not process_pid and process_name: + process_pid = stem.util.system.pid_by_name(process_name)[0] + try: results = stem.util.system.call(resolver_command) except OSError as exc: raise IOError("Unable to query '%s': %s" % (resolver_command, exc)) - + resolver_regex_str = RESOLVER_FILTER[resolver].format( protocol = '(?P<protocol>\S+)', local_address = '(?P<local_address>[0-9.]+)', diff --git a/stem/util/system.py b/stem/util/system.py index ea8d8a1..98b64f7 100644 --- a/stem/util/system.py +++ b/stem/util/system.py @@ -454,6 +454,9 @@ def pid_by_name(process_name, multiple = False): tasklist_regex_str = '^\s*' + process_name + '\s+(?P<pid>[0-9]*)' tasklist_regex = re.compile(tasklist_regex_str) + if not results: + raise IOError("No results found for tasklist") + for line in results: match = tasklist_regex.search(line) if match: @@ -461,13 +464,17 @@ def pid_by_name(process_name, multiple = False): id = int(attr['pid']) process_ids.append(id) - return process_ids + if process_ids == []: + raise IOError("Process Name not Found : %s" % process_name) + + if multiple: + return process_ids + elif len(process_ids) > 0: + return process_ids[0] except OSError as exc: log.debug("failed to query '%s': %s" % (command, exc)) raise IOError("Unable to query '%s': %s" % (command, exc)) - -
log.debug("failed to resolve a pid for '%s'" % process_name) diff --git a/test/integ/util/system.py b/test/integ/util/system.py index 1ffe7cd..3971e61 100644 --- a/test/integ/util/system.py +++ b/test/integ/util/system.py @@ -79,10 +79,7 @@ class TestSystem(unittest.TestCase): will fail if there's other tor instances running. """
- if stem.util.system.is_windows(): - test.runner.skip(self, '(unavailable on windows)') - return - elif self._is_extra_tor_running(): + if self._is_extra_tor_running(): test.runner.skip(self, '(multiple tor instances)') return
@@ -218,6 +215,24 @@ class TestSystem(unittest.TestCase):
if len(all_tor_pids) == 1: self.assertEqual(our_tor_pid, all_tor_pids[0]) + + def test_pid_by_name_tasklist(self): + """ + Tests the pid_by_name function with a tasklist response. + """ + + runner = test.runner.get_runner() + if self._is_extra_tor_running(): + test.runner.skip(self, '(multiple tor instances)') + return + elif not stem.util.system.is_available('tasklist'): + test.runner.skip(self, '(tasklist unavailable)') + return + + tor_pid = test.runner.get_runner().get_pid() + tor_cmd = test.runner.get_runner().get_tor_command(True) + self.assertEqual(tor_pid, stem.util.system.pid_by_name(tor_cmd)) +
def test_pid_by_port(self): """ diff --git a/test/unit/util/connection.py b/test/unit/util/connection.py index 731fa5b..233cb29 100644 --- a/test/unit/util/connection.py +++ b/test/unit/util/connection.py @@ -32,6 +32,18 @@ unix 3 [ ] STREAM CONNECTED 34164276 15843/tor unix 3 [ ] STREAM CONNECTED 7951 - """
+NETSTAT_WINDOWS_OUTPUT = """\ +Active Connections + + Proto Local Address Foreign Address State PID + TCP 0.0.0.0:135 0.0.0.0:0 LISTENING 852 + TCP 0.0.0.0:445 0.0.0.0:0 LISTENING 4 + TCP 0.0.0.0:902 0.0.0.0:0 LISTENING 2076 + TCP 0.0.0.0:912 0.0.0.0:0 LISTENING 2076 + TCP 192.168.0.1:44284 38.229.79.2:443 ESTABLISHED 15843 + TCP 0.0.0.0:37782 0.0.0.0:0 LISTENING 4128 +""" + SS_OUTPUT = """\ Netid State Recv-Q Send-Q Local Address:Port Peer Address:Port tcp CLOSE-WAIT 1 0 192.168.0.1:43780 53.203.145.45:443 users:(("firefox",20586,118)) @@ -105,19 +117,17 @@ BSD_PROCSTAT_OUTPUT = """\
class TestConnection(unittest.TestCase): - @patch('os.access') @patch('stem.util.system.is_available') @patch('stem.util.proc.is_available') - def test_system_resolvers(self, proc_mock, is_available_mock, os_mock): + def test_system_resolvers(self, proc_mock, is_available_mock): """ Checks the system_resolvers function. """
is_available_mock.return_value = True proc_mock.return_value = False - os_mock.return_value = True
- self.assertEqual([], stem.util.connection.system_resolvers('Windows')) + self.assertEqual([Resolver.NETSTAT_WINDOWS], stem.util.connection.system_resolvers('Windows')) self.assertEqual([Resolver.LSOF], stem.util.connection.system_resolvers('Darwin')) self.assertEqual([Resolver.LSOF], stem.util.connection.system_resolvers('OpenBSD')) self.assertEqual([Resolver.BSD_SOCKSTAT, Resolver.BSD_PROCSTAT, Resolver.LSOF], stem.util.connection.system_resolvers('FreeBSD')) @@ -126,7 +136,7 @@ class TestConnection(unittest.TestCase): proc_mock.return_value = True self.assertEqual([Resolver.PROC, Resolver.NETSTAT, Resolver.SOCKSTAT, Resolver.LSOF, Resolver.SS], stem.util.connection.system_resolvers('Linux'))
- # check that calling without an argument is equivalent to calling for this + # check that calling without an argument is equivilant to calling for this # platform
self.assertEqual(stem.util.connection.system_resolvers(platform.system()), stem.util.connection.system_resolvers()) @@ -185,6 +195,22 @@ class TestConnection(unittest.TestCase): self.assertRaises(IOError, stem.util.connection.get_connections, Resolver.NETSTAT, process_pid = 1111)
@patch('stem.util.system.call') + def test_get_connections_by_windows_netstat(self, call_mock): + """ + Checks the get_connections function with the Windows netstat resolver. + """ + + call_mock.return_value = NETSTAT_WINDOWS_OUTPUT.split('\n') + expected = [Connection('192.168.0.1', 44284, '38.229.79.2', 443, 'tcp')] + self.assertEqual(expected, stem.util.connection.get_connections(Resolver.NETSTAT_WINDOWS, process_pid = 15843, process_name = 'tor')) + + #self.assertRaises(IOError, stem.util.connection.get_connections, Resolver.NETSTAT, process_pid = 15843, process_name = 'stuff') + self.assertRaises(IOError, stem.util.connection.get_connections, Resolver.NETSTAT_WINDOWS, process_pid = 1111, process_name = 'tor') + + call_mock.side_effect = OSError('Unable to call netstat') + self.assertRaises(IOError, stem.util.connection.get_connections, Resolver.NETSTAT_WINDOWS, process_pid = 1111) + + @patch('stem.util.system.call') def test_get_connections_by_ss(self, call_mock): """ Checks the get_connections function with the ss resolver. diff --git a/test/unit/util/system.py b/test/unit/util/system.py index 3b4f035..94bb7fa 100644 --- a/test/unit/util/system.py +++ b/test/unit/util/system.py @@ -45,6 +45,16 @@ GET_PID_BY_NAME_PS_BSD_MULTIPLE = [ ' 10 ?? Ss 0:09.97 kextd', ' 41 ?? Ss 9:00.22 launchd']
+ +GET_PID_BY_NAME_TASKLIST_RESULTS = [ + 'Image Name PID Session Name Session# Mem Usage', + 'System Idle Process 0 Services 0 20 K', + 'svchost.exe 872 Services 0 8,744 K', + 'hpservice.exe 1112 Services 0 3,828 K', + 'tor.exe 3712 Console 1 29,976 K', + 'tor.exe 3713 Console 1 21,976 K', + 'conhost.exe 3012 Console 1 4,652 K'] + GET_PID_BY_PORT_NETSTAT_RESULTS = [ 'Active Internet connections (only servers)', 'Proto Recv-Q Send-Q Local Address Foreign Address State PID/Program name', @@ -53,7 +63,7 @@ GET_PID_BY_PORT_NETSTAT_RESULTS = [ 'tcp6 0 0 ::1:631 :::* LISTEN - ', 'udp 0 0 0.0.0.0:5353 0.0.0.0:* - ', 'udp6 0 0 fe80::7ae4:ff:fe2f::123 :::* - '] - + GET_PID_BY_PORT_SOCKSTAT_RESULTS = [ '_tor tor 4397 7 tcp4 51.64.7.84:9051 *:*', '_tor tor 4397 12 tcp4 51.64.7.84:54011 80.3.121.7:9051', @@ -252,6 +262,20 @@ class TestSystem(unittest.TestCase):
@patch('stem.util.system.call') @patch('stem.util.system.is_available', Mock(return_value = True)) + def test_pid_by_name_tasklist(self, call_mock): + """ + Tests the pid_by_name function with tasklist responses. + """ + + call_mock.return_value = GET_PID_BY_NAME_TASKLIST_RESULTS + self.assertEqual(3712, system.pid_by_name('tor')) + self.assertEqual(None, system.pid_by_name('DirectoryService')) + self.assertEqual(None, system.pid_by_name('blarg')) + + self.assertEqual([3712, 3713], system.pid_by_name('tor', multiple = True)) + + @patch('stem.util.system.call') + @patch('stem.util.system.is_available', Mock(return_value = True)) def test_pid_by_port_netstat(self, call_mock): """ Tests the pid_by_port function with a netstat response.