Merge pull request #54 from splunk/fross/restrictToHost

Fred Ross · Fred Ross · commit 30d20886999e · 2012-10-26T15:49:25.000-07:00
Add support for restrictToHost for TCP inputs
diff --git a/splunklib/client.py b/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)
@@ -1667,17 +1670,62 @@ def __init__(self, service, path, kind=None, **kwargs):
         Entity.__init__(self, service, path, **kwargs)
         if kind is None:
             path_segments = path.split('/')
-            i = path_segments.index('inputs')
-            self.kind = '/'.join(path_segments[i+1:-1])
-            assert self.kind is not None
+            i = path_segments.index('inputs') + 1
+            if path_segments[i] == 'tcp':
+                self.kind = path_segments[i] + '/' + path_segments[i+1]
+            else:
+                self.kind = path_segments[i]
         else:
             self.kind = kind
 
-        # Maintain compatibility with older kind labels
-        if self.kind == 'tcp/raw':
-            self.kind = 'tcp'
-        if self.kind == 'tcp/cooked':
-            self.kind = 'splunktcp'
+        # Handle old input kind names.
+        if self.kind == 'tcp':
+            self.kind = 'tcp/raw'
+        if self.kind == 'splunktcp':
+            self.kind = 'tcp/cooked'
+
+    @property
+    def restrictToHost(self):
+        if 'restrictToHost' in self._state.content:
+            return self._state.content['restrictToHost']
+        else:
+            return ''
+
+    def update(self, **kwargs):
+        if self.kind not in ['tcp', 'splunktcp', 'tcp/raw', 'tcp/cooked']:
+            result = super(Input, self).update(**kwargs)
+            return result
+        else:
+            # TCP inputs have a property 'restrictToHost' which requires special care.
+
+            # There is a bug in Splunk < 5.0. Don't bother trying to update restrictToHost.
+            if 'restrictToHost' in kwargs and self.service.splunk_version < (5,):
+                raise IllegalOperationException("Updating restrictToHost has no effect before Splunk 5.0")
+
+            # In Splunk 4.x, if you update without restrictToHost as one of the fields,
+            # restrictToHost keeps its previous state. In 5.0, it is set to empty string.
+            # Thus, we must pass it every time. This doesn't actually introduce a race
+            # condition because if someone else has set restrictToHost to a new value on this
+            # TCP input, our update request will fail, since our reference to it still uses
+            # the old path.
+            to_update = kwargs.copy()
+            to_update['restrictToHost'] = kwargs.get('restrictToHost', self['restrictToHost'])
+
+            # Do the actual update operation.
+            result = super(Input, self).update(**to_update)
+
+            # Now we must update the path in case it changed.
+            # The pieces we break it into are:
+            #  https://localhost:8089/services/data/inputs/tcp/raw/   boris:  10000
+            # |------------------ base_path -----------------------| | host || port|
+            assert self.path.endswith('/')
+            base_path = self.path.rsplit('/', 2)[0]
+            host = to_update['restrictToHost'] + ':' if to_update['restrictToHost'] != '' else ''
+            port = self.name.split(':', 1)[-1]
+            self.path = base_path + '/' + host + 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
@@ -1777,7 +1825,15 @@ def create(self, kind, name, **kwargs):
         """
         kindpath = self.kindpath(kind)
         self.post(kindpath, name=name, **kwargs)
-        path = _path(self.path + kindpath, name)
+
+        # If we created an input with restrictToHost set, then
+        # its path will be <restrictToHost>:<name>, not just <name>,
+        # and we have to adjust accordingly.
+        path = _path(
+            self.path + kindpath,
+            '%s:%s' % (kwargs['restrictToHost'], name) \
+                if kwargs.has_key('restrictToHost') else name
+        )
         return Input(self.service, path, kind)
 
     def delete(self, kind, name=None):
@@ -1807,7 +1863,8 @@ def _get_kind_list(self, subpath=[]):
             if entry.title == 'all' or this_subpath == ['tcp','ssl']:
                 continue
             elif 'create' in [x.rel for x in entry.link]:
-                kinds.append('/'.join(subpath + [entry.title]))
+                path = '/'.join(subpath + [entry.title])
+                kinds.append(path)
             else:
                 subkinds = self._get_kind_list(subpath + [entry.title])
                 kinds.extend(subkinds)
@@ -1864,7 +1921,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 []
diff --git a/tests/test_app.py b/tests/test_app.py
@@ -50,7 +50,7 @@ def tearDown(self):
             if app_name.startswith('delete-me'):
                 self.service.apps.delete(app_name)
                 self.assertEventuallyTrue(lambda: app_name not in self.service.apps)
-        self.clearRestartMessage()
+        self.clear_restart_message()
 
     def test_app_integrity(self):
         self.check_entity(self.app)
@@ -86,7 +86,7 @@ def test_delete(self):
         self.assertTrue(name in self.service.apps)
         self.service.apps.delete(name)
         self.assertFalse(name in self.service.apps)
-        self.clearRestartMessage() # We don't actually have to restart here.
+        self.clear_restart_message() # We don't actually have to restart here.
 
     def test_package(self):
         p = self.app.package()
diff --git a/tests/test_index.py b/tests/test_index.py
@@ -108,7 +108,7 @@ def test_submit_via_attached_socket(self):
         self.assertEventuallyTrue(lambda: self.totalEventCount() == eventCount+1, timeout=60)
 
     def test_upload(self):
-        self.installAppFromCollection("file_to_upload")
+        self.install_app_from_collection("file_to_upload")
 
         eventCount = int(self.index['totalEventCount'])
 
diff --git a/tests/test_input.py b/tests/test_input.py
@@ -19,6 +19,128 @@
 
 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_remove_host_restriction(self):
+        if self.service.splunk_version < (5,):
+            # We can't set restrictToHost before 5.0 due to a bug in splunkd.
+            return
+        input = self.service.inputs.create('tcp', str(self.base_port), restrictToHost='boris')
+        input.update(restrictToHost='')
+        input.refresh()
+        self.check_entity(input)
+        input.delete()
+
+    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']:
+            port = self.base_port
+            while True: # Find the next unbound port
+                try:
+                    boris = self.service.inputs.create(kind, str(port), restrictToHost='boris')
+                except client.HTTPError as he:
+                    if he.status == 400:
+                        port += 1
+                else:
+                    break
+
+            # No matter what version we're actually running against,
+            # we can check that on Splunk < 5.0, we get an exception
+            # from trying to update restrictToHost.
+            with self.fake_splunk_version((4,3)):
+                self.assertRaises(
+                    client.IllegalOperationException,
+                    lambda: boris.update(restrictToHost='hilda')
+                )
+
+            # And now back to real tests...
+            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()
+
+    def test_update_nonrestrictToHost(self):
+        for kind in ['tcp', 'splunktcp']:
+            port = self.base_port
+            while True: # Find the next unbound port
+                try:
+                    input = self.service.inputs.create(kind, str(port), restrictToHost='boris')
+                except client.HTTPError as he:
+                    if he.status == 400:
+                        port += 1
+                else:
+                    break
+            try:
+                input.update(host='meep')
+                input.refresh()
+                self.assertTrue(input.name.startswith('boris'))
+            except:
+                input.delete()
+                raise
+
 class TestRead(testlib.SDKTestCase):
     def test_read(self):
         inputs = self.service.inputs
@@ -59,7 +181,7 @@ def test_inputs_list_on_one_kind_with_search(self):
         self.assertEqual(expected, found)
 
     def test_oneshot(self):
-        self.installAppFromCollection('file_to_upload')
+        self.install_app_from_collection('file_to_upload')
 
         index_name = testlib.tmpname()
         index = self.service.indexes.create(index_name)
@@ -80,26 +202,23 @@ 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'}]
+        unrestricted_port = str(highest_port(self.service, 10000, 'tcp', 'splunktcp', 'udp')+1)
+        restricted_port = str(highest_port(self.service, int(unrestricted_port)+1, 'tcp', 'splunktcp')+1)
+        test_inputs = [{'kind': 'tcp', 'name': unrestricted_port, 'host': 'sdk-test'},
+                       {'kind': 'udp', 'name': unrestricted_port, 'host': 'sdk-test'},
+                       {'kind': 'tcp', 'name': 'boris:' + restricted_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', unrestricted_port, host='sdk-test')
         self._test_entities['udp'] = \
-            inputs.create('udp', str(base_port), host='sdk-test')
+            inputs.create('udp', unrestricted_port, host='sdk-test')
+        self._test_entities['restrictedTcp'] = \
+            inputs.create('tcp', restricted_port, restrictToHost='boris')
 
     def tearDown(self):
         super(TestInput, self).tearDown()
@@ -124,7 +243,7 @@ def test_lists_modular_inputs(self):
         else:
             # Install modular inputs to list, and restart
             # so they'll show up.
-            self.installAppFromCollection("modular-inputs")
+            self.install_app_from_collection("modular-inputs")
             self.uncheckedRestartSplunk()
 
             inputs = self.service.inputs
@@ -159,11 +278,10 @@ def test_update(self):
         inputs = self.service.inputs
         for entity in self._test_entities.itervalues():
             kind, name = entity.kind, entity.name
-            kwargs = {'host': 'foo', 'sourcetype': 'bar'}
+            kwargs = {'host': 'foo'}
             entity.update(**kwargs)
             entity.refresh()
             self.assertEqual(entity.host, kwargs['host'])
-            self.assertEqual(entity.sourcetype, kwargs['sourcetype'])
 
     def test_delete(self):
         inputs = self.service.inputs
@@ -177,14 +295,17 @@ def test_delete(self):
                 inputs.delete(name)
                 self.assertFalse(name in inputs)
             else:
-                self.assertRaises(client.AmbiguousReferenceException,
-                                  inputs.delete, name)
+                if not name.startswith('boris'):
+                    self.assertRaises(client.AmbiguousReferenceException,
+                        inputs.delete, name)
                 self.service.inputs.delete(kind, name)
                 self.assertFalse((kind,name) in inputs)
             self.assertRaises(client.EntityDeletedException,
                               input_entity.refresh)
             remaining -= 1
-                              
+
+
+
 
 if __name__ == "__main__":
     import unittest
diff --git a/tests/test_job.py b/tests/test_job.py
@@ -114,7 +114,7 @@ def test_read_jobs(self):
 class TestJobWithDelayedDone(testlib.SDKTestCase):
     def setUp(self):
         super(TestJobWithDelayedDone, self).setUp()
-        self.installAppFromCollection("sleep_command")
+        self.install_app_from_collection("sleep_command")
         self.query = "search index=_internal | sleep done=100"
         self.job = self.service.jobs.create(
             query=self.query,
diff --git a/tests/test_modular_input_kinds.py b/tests/test_modular_input_kinds.py
@@ -20,7 +20,7 @@
 class ModularInputKindTestCase(testlib.SDKTestCase):
     def setUp(self):
         super(ModularInputKindTestCase, self).setUp()
-        self.installAppFromCollection("modular-inputs")
+        self.install_app_from_collection("modular-inputs")
         self.uncheckedRestartSplunk()
 
     def test_list_arguments(self):
diff --git a/tests/test_service.py b/tests/test_service.py
@@ -101,6 +101,10 @@ def test_splunk_version(self):
         for p in v:
             self.assertTrue(isinstance(p, int) and p >= 0)
 
+        for version in [(4,3,3), (5,), (5,0,1)]:
+            with self.fake_splunk_version(version):
+                self.assertEqual(version, self.service.splunk_version)
+
 class TestSettings(testlib.SDKTestCase):
     def test_read_settings(self):
         settings = self.service.settings
diff --git a/tests/testlib.py b/tests/testlib.py
