DVPL-3313: splunklib | Improve the SDK to make it easier to talk back to Splunk's API from within a modular input or custom search command

Part 2: splunklib.modularinput

Added Script.service property that returns a service object or None. Verified with automated unit tests, ad-hoc integration tests, and code review (thanks Shakeel).

The service object is created from the Server URI and Session Key passed to the command invocation on the input stream. It is available
as soon as the :code:`Script.stream_events` method is  called.   The script.service property returns a value of None, if you call it before the Script.stream_events` method is called.

Signed-off-by: David Noble <[email protected]>

Author: David Noble

splunklib/modularinput/script.py

@@ -13,8 +13,10 @@
 # under the License.
 
 from abc import ABCMeta, abstractmethod
+from urlparse import urlsplit
 import sys
 
+from splunklib.client import Service
 from splunklib.modularinput.event_writer import EventWriter
 from splunklib.modularinput.input_definition import InputDefinition
 from splunklib.modularinput.validation_definition import ValidationDefinition
@@ -37,6 +39,10 @@ class Script(object):
     """
     __metaclass__ = ABCMeta
 
+    def __init__(self):
+        self._input_definition = None
+        self._service = None
+
     def run(self, args):
         """Runs this modular input
 
@@ -59,19 +65,22 @@ def run_script(self, args, event_writer, input_stream):
 
         try:
             if len(args) == 1:
-                # This script is running as an input. Input definitions will be passed on stdin
-                # as XML, and the script will write events on stdout and log entries on stderr.
-                input_definition = InputDefinition.parse(input_stream)
-                self.stream_events(input_definition, event_writer)
+                # This script is running as an input. Input definitions will be
+                # passed on stdin as XML, and the script will write events on
+                # stdout and log entries on stderr.
+                self._input_definition = InputDefinition.parse(input_stream)
+                self.stream_events(self._input_definition, event_writer)
                 event_writer.close()
                 return 0
 
             elif str(args[1]).lower() == "--scheme":
-                # Splunk has requested XML specifying the scheme for this modular input.
-                # Return it and exit.
+                # Splunk has requested XML specifying the scheme for this
+                # modular input Return it and exit.
                 scheme = self.get_scheme()
                 if scheme is None:
-                    event_writer.log(EventWriter.FATAL, "Modular input script returned a null scheme.")
+                    event_writer.log(
+                        EventWriter.FATAL,
+                        "Modular input script returned a null scheme.")
                     return 1
                 else:
                     event_writer.write_xml_document(scheme.to_xml())
@@ -97,6 +106,39 @@ def run_script(self, args, event_writer, input_stream):
             event_writer._err.write(err_string)
             return 1
 
+    @property
+    def service(self):
+        """ Returns a Splunk service object for this script invocation.
+
+        The service object is created from the Splunkd URI and session key
+        passed to the command invocation on the modular input stream. It is
+        available as soon as the :code:`Script.stream_events` method is
+        called.
+
+        :return: :class:splunklib.client.Service. A value of None is returned,
+        if you call this method before the :code:`Script.stream_events` method
+        is called.
+
+        """
+        if self._service is not None:
+            return self._service
+
+        if self._input_definition is None:
+            return None
+
+        splunkd_uri = self._input_definition.metadata["server_uri"]
+        session_key = self._input_definition.metadata["session_key"]
+
+        scheme, netloc, _, _, _ = urlsplit(splunkd_uri, allow_fragments=False)
+
+        splunkd_host, splunkd_port = netloc.split(':')
+
+        self._service = Service(
+            scheme=scheme, host=splunkd_host, port=splunkd_port,
+            token=session_key)
+
+        return self._service
+
     @abstractmethod
     def get_scheme(self):
         """The scheme defines the parameters understood by this modular input.
@@ -105,13 +147,16 @@ def get_scheme(self):
         """
 
     def validate_input(self, definition):
-        """Handles external validation for modular input kinds. When Splunk
-        calls a modular input script in validation mode, it will pass in an XML document
-        giving information about the Splunk instance (so you can call back into it if needed)
-        and the name and parameters of the proposed input.
+        """Handles external validation for modular input kinds.
+
+        When Splunk calls a modular input script in validation mode, it will
+        pass in an XML document giving information about the Splunk instance (so
+        you can call back into it if needed) and the name and parameters of the
+        proposed input.
 
-        If this function does not throw an exception, the validation is assumed to succeed.
-        Otherwise any errors thrown will be turned into a string and logged back to Splunk.
+        If this function does not throw an exception, the validation is assumed
+        to succeed. Otherwise any errors thrown will be turned into a string and
+        logged back to Splunk.
 
         The default implementation always passes.
 

tests/modularinput/test_script.py

@@ -13,6 +13,7 @@
 # under the License.
 
 from tests.modularinput.modularinput_testlib import unittest, xml_compare, data_open
+from splunklib.client import Service
 from splunklib.modularinput.argument import Argument
 from splunklib.modularinput.event import Event
 from splunklib.modularinput.event_writer import EventWriter
@@ -210,5 +211,42 @@ def stream_events(self, inputs, ew):
 
         self.assertTrue(xml_compare(expected, found))
 
+    def test_service_property(self):
+        """ Check that Script.service returns a valid Service instance as soon
+        as the stream_events method is called, but not before.
+
+        """
+
+        # Override abstract methods
+        class NewScript(Script):
+            def __init__(self, test):
+                super(NewScript, self).__init__()
+                self.test = test
+
+            def get_scheme(self):
+                return None
+
+            def stream_events(self, inputs, ew):
+                service = self.service
+                self.test.assertIsInstance(service, Service)
+                self.test.assertEqual(str(service.authority), inputs.metadata['server_uri'])
+
+        script = NewScript(self)
+        input_configuration = data_open("data/conf_with_2_inputs.xml")
+
+        out = StringIO()
+        err = StringIO()
+        ew = EventWriter(out, err)
+
+        self.assertEqual(script.service, None)
+
+        return_value = script.run_script(
+            [TEST_SCRIPT_PATH], ew, input_configuration)
+
+        self.assertEqual(0, return_value)
+        self.assertEqual("", err.getvalue())
+
+        return
+
 if __name__ == "__main__":
     unittest.main()