Allows to specify index when using splunk HTTP Event collector

flow/record/adapter/splunk.py

@@ -21,18 +21,24 @@
 from flow.record.utils import boolean_argument, to_base64, to_bytes, to_str
 
 if TYPE_CHECKING:
+    from collections.abc import Callable
+
     from flow.record.base import Record
 
 __usage__ = """
 Splunk output adapter (writer only)
 ---
 Write usage: rdump -w splunk+[PROTOCOL]://[IP]:[PORT]?tag=[TAG]&token=[TOKEN]&sourcetype=[SOURCETYPE]
-[PROTOCOL]: Protocol to use for forwarding data. Can be tcp, http or https, defaults to tcp if omitted.
 [IP]:[PORT]: ip and port to a splunk instance
-[TAG]: optional value to add as "rdtag" output field when writing
 [TOKEN]: Authentication token for sending data over HTTP(S)
-[SOURCETYPE]: Set sourcetype of data. Defaults to records, but can also be set to JSON.
-[SSL_VERIFY]: Whether to verify the server certificate when sending data over HTTPS. Defaults to True.
+
+Optional arguments:
+    [PROTOCOL]: Protocol to use for forwarding data [tcp|http|https] (default: tcp).
+    [TAG]: optional value to add as "rdtag" output field when writing
+    [SOURCETYPE]: Set sourcetype of data [records|json] (default: records).
+    [SSL_VERIFY]: Whether to verify the server certificate when sending data over HTTPS (default: True).
+    [INDEX]: The name of the index by which the event data is to be indexed (default: None, use token default index).
+        Ignored when PROTOCOL is tcp.
 """
 
 log = logging.getLogger(__name__)
@@ -132,18 +138,20 @@ def record_to_splunk_json(packer: JsonRecordPacker, record: Record, tag: str | N
     return json_dict
 
 
-def record_to_splunk_http_api_json(packer: JsonRecordPacker, record: Record, tag: str | None = None) -> str:
+def record_to_splunk_http_api_json(
+    packer: JsonRecordPacker, record: Record, tag: str | None = None, index: str | None = None
+) -> str:
     ret = {}
 
-    indexer_fields = [
+    event_medatata = [
         ("host", "host"),
         ("host", "hostname"),
         ("time", "ts"),
     ]
 
     # When converting a record to json text for splunk, we distinguish between the 'event' (containing the data) and a
-    # few other fields that are splunk-specific for indexing. We add those 'indexer_fields' to the return object first.
-    for splunk_name, field_name in indexer_fields:
+    # few other fields that are splunk-specific for indexing. We add those 'event_medatata' to the return object first.
+    for splunk_name, field_name in event_medatata:
         if hasattr(record, field_name):
             val = getattr(record, field_name)
             if val:
@@ -152,12 +160,15 @@ def record_to_splunk_http_api_json(packer: JsonRecordPacker, record: Record, tag
                     ret[splunk_name] = val.timestamp()
                     continue
                 ret[splunk_name] = to_str(val)
-
+    if index:
+        ret["index"] = index
     ret["event"] = record_to_splunk_json(packer, record, tag)
     return json.dumps(ret, default=packer.pack_obj)
 
 
-def record_to_splunk_tcp_api_json(packer: JsonRecordPacker, record: Record, tag: str | None = None) -> str:
+def record_to_splunk_tcp_api_json(
+    packer: JsonRecordPacker, record: Record, tag: str | None = None, index: str | None = None
+) -> str:
     record_dict = record_to_splunk_json(packer, record, tag)
     return json.dumps(record_dict, default=packer.pack_obj)
 
@@ -173,6 +184,7 @@ def __init__(
         token: str | None = None,
         sourcetype: str | None = None,
         ssl_verify: bool = True,
+        index: str | None = None,
         **kwargs,
     ):
         # If the writer is initiated without a protocol, we assume we will be writing over tcp
@@ -192,10 +204,11 @@ def __init__(
         self.port = parsed_url.port
 
         self.tag = tag
+        self.index = index
         self.record_buffer = []
         self._warned = False
         self.packer = None
-        self.json_converter = None
+        self.json_converter: Callable[[JsonRecordPacker, Record, str | None, str | None], str] | None = None
 
         if self.protocol == Protocol.TCP:
             self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM, socket.SOL_TCP)
@@ -225,6 +238,8 @@ def __init__(
             endpoint = "event" if self.sourcetype != SourceType.RECORDS else "raw"
             port = f":{self.port}" if self.port else ""
             self.url = f"{scheme}://{self.host}{port}/services/collector/{endpoint}?auto_extract_timestamp=true"
+            if self.sourcetype == SourceType.RECORDS and self.index:
+                self.url = f"{self.url}&index={self.index}"
 
             self.headers = {
                 "Authorization": self.token,
@@ -280,7 +295,7 @@ def write(self, record: Record) -> None:
         if self.sourcetype == SourceType.RECORDS:
             rec = record_to_splunk_kv_line(record, self.tag)
         else:
-            rec = self.json_converter(self.packer, record, self.tag)
+            rec = self.json_converter(self.packer, record, self.tag, self.index)
 
         # Trail with a newline for line breaking.
         data = to_bytes(rec) + b"\n"

tests/adapter/test_splunk.py

@@ -440,3 +440,111 @@ def test_https_protocol_json_sourcetype(mock_httpx_package: MagicMock) -> None:
                 **BASE_FIELD_JSON_VALUES,
             )
         }
+
+
+def test_https_protocol_records_sourcetype_with_index(mock_httpx_package: MagicMock) -> None:
+    if "flow.record.adapter.splunk" in sys.modules:
+        del sys.modules["flow.record.adapter.splunk"]
+
+    from flow.record.adapter.splunk import Protocol, SourceType, SplunkWriter
+
+    with patch.object(
+        flow.record.adapter.splunk,
+        "HAS_HTTPX",
+        True,
+    ):
+        mock_httpx_package.Client.return_value.post.return_value.status_code = 200
+        https_writer = SplunkWriter("https://splunk:8088", token="password123", index="dissect")
+
+        assert https_writer.host == "splunk"
+        assert https_writer.protocol == Protocol.HTTPS
+        assert https_writer.sourcetype == SourceType.RECORDS
+        assert https_writer.verify is True
+        assert (
+            https_writer.url == "https://splunk:8088/services/collector/raw?auto_extract_timestamp=true&index=dissect"
+        )
+
+        _, kwargs = mock_httpx_package.Client.call_args
+        assert kwargs["verify"] is True
+
+        given_headers = kwargs["headers"]
+        assert given_headers["Authorization"] == "Splunk password123"
+        assert "X-Splunk-Request-Channel" in given_headers
+
+        test_record_descriptor = RecordDescriptor(
+            "test/record",
+            [("string", "foo")],
+        )
+
+        test_record = test_record_descriptor(foo="bar")
+        https_writer.write(test_record)
+
+        mock_httpx_package.Client.return_value.post.assert_not_called()
+
+        https_writer.close()
+        mock_httpx_package.Client.return_value.post.assert_called_with(
+            "https://splunk:8088/services/collector/raw?auto_extract_timestamp=true&index=dissect",
+            data=ANY,
+        )
+        _, kwargs = mock_httpx_package.Client.return_value.post.call_args
+        sent_data = kwargs["data"]
+        assert sent_data.startswith(b'rdtype="test/record" rdtag=None foo="bar" ' + BASE_FIELDS_KV_SUFFIX.encode())
+        assert sent_data.endswith(b'"\n')
+
+
+def test_https_protocol_json_sourcetype_with_index(mock_httpx_package: MagicMock) -> None:
+    if "flow.record.adapter.splunk" in sys.modules:
+        del sys.modules["flow.record.adapter.splunk"]
+
+    from flow.record.adapter.splunk import SplunkWriter
+
+    with patch.object(
+        flow.record.adapter.splunk,
+        "HAS_HTTPX",
+        True,
+    ):
+        mock_httpx_package.Client.return_value.post.return_value.status_code = 200
+
+        https_writer = SplunkWriter("https://splunk:8088", token="password123", sourcetype="json", index="dissect")
+
+        test_record_descriptor = RecordDescriptor(
+            "test/record",
+            [("string", "foo")],
+        )
+
+        https_writer.write(test_record_descriptor(foo="bar"))
+        https_writer.write(test_record_descriptor(foo="baz"))
+        mock_httpx_package.Client.return_value.post.assert_not_called()
+
+        https_writer.close()
+        mock_httpx_package.Client.return_value.post.assert_called_with(
+            "https://splunk:8088/services/collector/event?auto_extract_timestamp=true",
+            data=ANY,
+        )
+
+        _, kwargs = mock_httpx_package.Client.return_value.post.call_args
+        sent_data = kwargs["data"]
+        first_record_json, _, second_record_json = sent_data.partition(b"\n")
+        assert json.loads(first_record_json) == {
+            "index": "dissect",
+            "event": dict(
+                {
+                    "rdtag": None,
+                    "rdtype": "test/record",
+                    "foo": "bar",
+                },
+                **BASE_FIELD_JSON_VALUES,
+            ),
+        }
+        print(second_record_json)
+        assert json.loads(second_record_json) == {
+            "index": "dissect",
+            "event": dict(
+                {
+                    "rdtag": None,
+                    "rdtype": "test/record",
+                    "foo": "baz",
+                },
+                **BASE_FIELD_JSON_VALUES,
+            ),
+        }