Repository: hkraal/ssht

Issue:
Use more fields for searching
Searching based on IP does not yield any result. Although it's not very useful `ssht <ip>` it should be possible since;
- a user might be tempted to use `ssht` instead of `ssh`
- a user might want to get feedback about the server it's connecting to

Comments:

Patch:
diff --git a/ssht/plugins.py b/ssht/plugins.py
index e321935..0edfb7a 100644
--- a/ssht/plugins.py
+++ b/ssht/plugins.py
@@ -34,6 +34,13 @@ class Host(object):
             return '{0}@{1}'.format(self.user, self.hostname)
         return self.hostname
 
+    def match(self, needle):
+        for search_field in ['hostname', 'ipv4', 'ipv6']:
+            value = getattr(self, search_field, None)
+            if value is not None and needle in value:
+                return True
+        return False
+
     def __repr__(self):
         if self.ipv4 is not None:
             return '<Host: hostname={0}, ipv4={1}>'.format(
@@ -54,8 +61,15 @@ class Parser(object):
     def get_files(self, ext='.json'):
         return [x for x in os.listdir(self._path) if x.endswith(ext)]
 
-    def search(self, name):
-        return [x for x in self._hosts if name in x.hostname]
+    def search(self, needle):
+        '''
+        Return Host objects which have properties matching the needle
+        '''
+        results = []
+        for host in self._hosts:
+            if host.match(needle):
+                results.append(host)
+        return results
 
 
 class JsonParser(Parser):   # pragma: nocover


Test patch:
diff --git a/tests/test_parser.py b/tests/test_parser.py
index 315a4ee..6409173 100644
--- a/tests/test_parser.py
+++ b/tests/test_parser.py
@@ -15,13 +15,27 @@ class TestParser:
         assert parser.get_files('.mysql') == ['test.mysql']
         assert parser.get_files('.fake') == []
 
-    def test_search(self):
+    def test_search_hostname(self):
         parser = Parser('/tmp')
         parser._hosts = [Host('host01.example.com'),
                          Host('host02.example.com')]
         assert parser.search('host01')[0].hostname == 'host01.example.com'
         assert len(parser.search('example.com')) == 2
 
+    def test_search_ipv4(self):
+        parser = Parser('/tmp')
+        parser._hosts = [Host('host01.example.com', ipv4='192.168.0.1'),
+                         Host('host02.example.com', ipv4='192.168.0.2')]
+        assert parser.search('192.168.0.1')[0].hostname == 'host01.example.com'
+        assert len(parser.search('192.168.0')) == 2
+
+    def test_search_ipv6(self):
+        parser = Parser('/tmp')
+        parser._hosts = [Host('host01.example.com', ipv6='fe80::41:dead:beef:cafe'),
+                         Host('host02.example.com', ipv6='fe80::41:dead:beef:daff')]
+        assert parser.search('dead:beef:cafe')[0].hostname == 'host01.example.com'
+        assert len(parser.search('dead:beef:')) == 2
+
 
 class TestJsonParser:
     pass
diff --git a/tests/test_ssht.py b/tests/test_ssht.py
index 3292f48..28f8e76 100644
--- a/tests/test_ssht.py
+++ b/tests/test_ssht.py
@@ -90,8 +90,8 @@ class TestConnect:
 class TestSsht():
 
     def test_select_host(self, mocker):
-        hosts = [Host(hostname='host01.example.com'),
-                 Host(hostname='host02.example.com', user='admin')]
+        hosts = [Host(hostname='host01.example.com', ipv4='127.0.0.1'),
+                 Host(hostname='host02.example.com', ipv4='127.0.0.2', user='admin')]
 
         mocker.patch('ssht.ssht.get_answer', return_value='1')
         assert select_host(hosts).hostname == 'host01.example.com'
@@ -100,7 +100,7 @@ class TestSsht():
         mocker.patch('ssht.ssht.get_answer', return_value='2')
         assert select_host(hosts).hostname == 'host02.example.com'
         assert select_host(hosts).user == 'admin'
-        
+
         mocker.patch('ssht.ssht.get_answer', return_value='')
         assert select_host(hosts) is None