Made TCP inputs handle restrictToHost.

Author: Frederick Ross

splunklib/client.py

@@ -797,6 +797,9 @@ def _proper_namespace(self, owner=None, app=None, sharing=None):
                 sharing = self._state.access.sharing
         return (owner, app, sharing)
 
+    def delete(self):
+        owner, app, sharing = self._proper_namespace()
+        return self.service.delete(self.path, owner=owner, app=app, sharing=sharing)
 
     def get(self, path_segment="", owner=None, app=None, sharing=None, **query):
         owner, app, sharing = self._proper_namespace(owner, app, sharing)
@@ -1679,6 +1682,21 @@ def __init__(self, service, path, kind=None, **kwargs):
         if self.kind == 'tcp/cooked':
             self.kind = 'splunktcp'
 
+    def update(self, **kwargs):
+        if 'restrictToHost' in kwargs and self.service.splunk_version < (5,):
+            raise IllegalOperationException("Updating restrictToHost has no effect before Splunk 5.0")
+        port = self.name.split(':', 1)[-1]
+        result = super(Input, self).update(**kwargs)
+        if 'restrictToHost' in kwargs:
+            if self.path.endswith('/'):
+                base_path, name = self.path.rsplit('/', 2)[:-1]
+            else:
+                base_path, name = self.path.rsplit('/', 1)
+            self.path = base_path + '/' + kwargs['restrictToHost'] + ':' + port
+        return result
+
+
+
 # Inputs is a "kinded" collection, which is a heterogenous collection where
 # each item is tagged with a kind, that provides a single merged view of all
 # input kinds.
@@ -1776,8 +1794,15 @@ def create(self, kind, name, **kwargs):
         :return: The new input.
         """
         kindpath = self.kindpath(kind)
-        self.post(kindpath, name=name, **kwargs)
-        path = _path(self.path + kindpath, name)
+        if (kindpath == 'tcp/raw' or kindpath == 'tcp/cooked' or kindpath == 'udp') and \
+            'restrictToHost' in kwargs:
+            post_name = name
+            path_name = kwargs['restrictToHost'] + ':' + name
+        else:
+            post_name = name
+            path_name = name
+        self.post(kindpath, name=post_name, **kwargs)
+        path = _path(self.path + kindpath, path_name)
         return Input(self.service, path, kind)
 
     def delete(self, kind, name=None):
@@ -1864,7 +1889,7 @@ def list(self, *kinds, **kwargs):
             path = self.kindpath(kind)
             logging.debug("Path for inputs: %s", path)
             try:
-                response = self.service.get(path, **kwargs)
+                response = self.get(path, **kwargs)
             except HTTPError, he:
                 if he.status == 404: # No inputs of this kind
                     return []

tests/test_input.py

@@ -19,6 +19,85 @@
 
 import splunklib.client as client
 
+def highest_port(service, base_port, *kinds):
+    """Find the first port >= base_port not in use by any input in kinds."""
+    highest_port = base_port
+    for input in service.inputs.list(*kinds):
+        port = int(input.name.split(':')[-1])
+        highest_port = max(port, highest_port)
+    return highest_port
+
+class TestTcpInputNameHandling(testlib.SDKTestCase):
+    def setUp(self):
+        super(TestTcpInputNameHandling, self).setUp()
+        self.base_port = highest_port(self.service, 10000, 'tcp', 'splunktcp') + 1
+
+    def tearDown(self):
+        for input in self.service.inputs.list('tcp', 'splunktcp'):
+            port = int(input.name.split(':')[-1])
+            if port >= self.base_port:
+                input.delete()
+        super(TestTcpInputNameHandling, self).tearDown()
+
+    def test_create_tcp_port(self):
+        for kind in ['tcp', 'splunktcp']:
+            input = self.service.inputs.create(kind, str(self.base_port))
+            self.check_entity(input)
+            input.delete()
+
+    def test_cannot_create_with_restrictToHost_in_name(self):
+        self.assertRaises(
+            client.HTTPError,
+            lambda: self.service.inputs.create('tcp', 'boris:10000')
+        )
+
+    def test_create_tcp_ports_with_restrictToHost(self):
+        for kind in ['tcp', 'splunktcp']:
+            # Make sure we can create two restricted inputs on the same port
+            boris = self.service.inputs.create(kind, str(self.base_port), restrictToHost='boris')
+            natasha = self.service.inputs.create(kind, str(self.base_port), restrictToHost='natasha')
+            # And that they both function
+            boris.refresh()
+            natasha.refresh()
+            self.check_entity(boris)
+            self.check_entity(natasha)
+            boris.delete()
+            natasha.delete()
+
+    def test_restricted_to_unrestricted_collision(self):
+        for kind in ['tcp', 'splunktcp']:
+            restricted = self.service.inputs.create(kind, str(self.base_port), restrictToHost='boris')
+            self.assertRaises(
+                client.HTTPError,
+                lambda: self.service.inputs.create(kind, str(self.base_port))
+            )
+            restricted.delete()
+
+    def test_unrestricted_to_restricted_collision(self):
+        for kind in ['tcp', 'splunktcp']:
+            unrestricted = self.service.inputs.create(kind, str(self.base_port))
+            self.assertRaises(
+                client.HTTPError,
+                lambda: self.service.inputs.create(kind, str(self.base_port), restrictToHos='boris')
+            )
+            unrestricted.delete()
+
+    def test_update_restrictToHost(self):
+        for kind in ['tcp', 'splunktcp']:
+            boris = self.service.inputs.create(kind, str(self.base_port), restrictToHost='boris')
+            with self.fake_splunk_version((4,3)):
+                self.assertRaises(
+                    client.IllegalOperationException,
+                    lambda: boris.update(restrictToHost='hilda')
+                )
+            if self.service.splunk_version >= (5,):
+                boris.update(restrictToHost='hilda')
+                boris.refresh()
+                self.assertEqual('hilda:' + str(self.base_port), boris.name)
+                boris.refresh()
+                self.check_entity(boris)
+                boris.delete()
+
 class TestRead(testlib.SDKTestCase):
     def test_read(self):
         inputs = self.service.inputs
@@ -80,26 +159,20 @@ def test_oneshot_on_nonexistant_file(self):
         self.assertRaises(client.OperationFailedException,
             self.service.inputs.oneshot, name)
 
-
 class TestInput(testlib.SDKTestCase):
     def setUp(self):
         super(TestInput, self).setUp()
         inputs = self.service.inputs
-        test_inputs = [{'kind': 'tcp', 'name': '9999', 'host': 'sdk-test'},
-                       {'kind': 'udp', 'name': '9999', 'host': 'sdk-test'}]
+        tcp_port = str(highest_port(self.service, 10000, 'tcp', 'splunktcp')+1)
+        udp_port = str(highest_port(self.service, 10000, 'udp')+1)
+        test_inputs = [{'kind': 'tcp', 'name': tcp_port, 'host': 'sdk-test'},
+                       {'kind': 'udp', 'name': udp_port, 'host': 'sdk-test'}]
         self._test_entities = {}
 
-        base_port = 10000
-        while True:
-            if str(base_port) in inputs:
-                base_port += 1
-            else:
-                break
-
         self._test_entities['tcp'] = \
-            inputs.create('tcp', str(base_port), host='sdk-test')
+            inputs.create('tcp', str(tcp_port), host='sdk-test')
         self._test_entities['udp'] = \
-            inputs.create('udp', str(base_port), host='sdk-test')
+            inputs.create('udp', str(udp_port), host='sdk-test')
 
     def tearDown(self):
         super(TestInput, self).tearDown()
@@ -184,7 +257,9 @@ def test_delete(self):
             self.assertRaises(client.EntityDeletedException,
                               input_entity.refresh)
             remaining -= 1
-                              
+
+
+
 
 if __name__ == "__main__":
     import unittest