load_splunk_data.py

#!/usr/bin/env python3
"""
Python port of load-splunk-data.sh — load files into Splunk via REST + oneshot CLI.
"""

from __future__ import annotations

import argparse
import base64
import csv
import io
import os
import re
import ssl
import subprocess
import sys
import time
import urllib.error
import urllib.parse
import urllib.request
from pathlib import Path
from typing import Mapping

VERSION = "1.0.2"

# URL escape codes used in Splunk searches (same as shell script)
DQT = "%22"
PCT = "%25"
DOL = "%24"
SLB = "%5C"

CFG_FILE = Path("./sample.cfg")
SETTINGS_FILE = Path("./settings.txt")

PROPS_DESIRED = [
    "TIME_FORMAT",
    "TIME_PREFIX",
    "MAX_TIMESTAMP_LOOKAHEAD",
    "SHOULD_LINEMERGE",
    "LINE_BREAKER",
    "TRUNCATE",
    "EVENT_BREAKER_ENABLE",
    "EVENT_BREAKER",
]


def show_help() -> None:
    print(
        f"""

VERSION: {VERSION}
SOURCE: https://github.com/tmuth/splunk-data-load

##################################################################
This program streamlines getting files into Splunk. It will:
1. Delete the specified INDEX and recreate it
2. Reload the input, fields, transforms, and props configs
3. Oneshot load all files in the specified directory using the defined sourcetype and INDEX
4. Count events and show field summary (when enabled in settings)

Pass the path to a configuration file (same format as the shell script: KEY=value lines).
{CFG_FILE.name} and {SETTINGS_FILE.name} are generated on first run if missing.
##################################################################

"""
    )


def gen_sample_cfg(path: Path) -> None:
    text = f"""
SPLUNK_HOST=localhost:8089
AUTH_TOKEN=# Token authentication is the preferred method over username/password and is documented here:
# https://docs.splunk.com/Documentation/Splunk/latest/Security/EnableTokenAuth
SPLUNK_USERNAME=admin
SPLUNK_PASS=welcome1
INDEX=sample_index
INDEX_TYPE=event
SOURCETYPE=sample_sourcetype
DIRECTORY=.
#DIRECTORY=/Volumes/GoogleDrive/My\\ Drive/Projects/splunking-json/Docker/data/fio
# Either set EXTENSION to something like json to load a number of files or set FILE_NAME to a
# specific file name to load a single file. Don't set both. Leave the unused variable empty
# with no spaces.
EXTENSION=
FILE_NAME=test.json
APP_NAME=tmuth-data-load
REPORT_FIELDS=* # first_name,ip_address,last_name
HOST_SEGMENT= # Space by default. Set to a number of the segement of filename for host if needed.
"""
    path.write_text(text.lstrip("\n"), encoding="utf-8")


def gen_settings(path: Path, sourcetype: str) -> None:
    text = f"""
# The following are global settings used to change how load_splunk_data.py runs
#
DEBUG_TIMESTAMP=F # T or F    : Searches _internl component=DateParser OR component=DateParserVerbose for timestamp errors
DEBUG=T # T or F              : Searches log_level=ERROR sourcetype=splunkd after each oneshot filtering on your sourcetype
SHOW_BTOOL=F # F or F         : Shows btool 'splunk btool check' and 'splunk btool props list {sourcetype} --debug'
SHOW_GREAT_8=T # T or F       : Checks props.conf for the 'Great 8' or 'gr8' settings
SHOW_WALKLEX=F # T or F       : Rolls hot-buckets to warm, then runs walklex to show indexed-fields
SHOW_EVENT_SUMMARY=T # T or F : Searches '...| stats count' and '| fieldsummary '
SHOW_INDEX_CONF=T # T or F    : Displays status of index delete/create and reload of .conf files
"""
    path.write_text(text.lstrip("\n"), encoding="utf-8")


def _strip_inline_comment(value: str) -> str:
    in_single = False
    in_double = False
    for i, ch in enumerate(value):
        if ch == "'" and not in_double:
            in_single = not in_single
        elif ch == '"' and not in_single:
            in_double = not in_double
        elif ch == "#" and not in_single and not in_double:
            return value[:i].rstrip()
    return value.rstrip()


def parse_env_file(path: Path) -> dict[str, str]:
    """Parse KEY=value lines like bash `source` (comments with #, not inside quotes)."""
    out: dict[str, str] = {}
    if not path.is_file():
        return out
    for raw_line in path.read_text(encoding="utf-8", errors="replace").splitlines():
        line = raw_line.strip()
        if not line or line.startswith("#"):
            continue
        if "=" not in raw_line:
            continue
        key, _, rest = raw_line.partition("=")
        key = key.strip()
        if not key or key.startswith("#"):
            continue
        val = _strip_inline_comment(rest).strip()
        out[key] = val
    return out


def load_config(cfg_path: Path) -> dict[str, str]:
    merged: dict[str, str] = {}
    merged.update(parse_env_file(cfg_path))
    merged.update(parse_env_file(SETTINGS_FILE))
    return merged


def truthy(cfg: Mapping[str, str], key: str) -> bool:
    return (cfg.get(key) or "").strip().upper() == "T"


def print_section(title: str) -> None:
    pad = title + " " + ("*" * 120)
    print(f"\n****** {pad[:100]}")
    if title.endswith("_END"):
        print("\n")


def splunk_bin(cfg: Mapping[str, str]) -> Path:
    home = os.environ.get("SPLUNK_HOME", "")
    if not home:
        print(
            "\n!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n"
            "The SPLUNK_HOME environment variable is unset. Please configure before running this program.\n"
            "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!\n"
        )
        sys.exit(1)
    return Path(home) / "bin" / "splunk"


def _ssl_context() -> ssl.SSLContext:
    ctx = ssl.create_default_context()
    ctx.check_hostname = False
    ctx.verify_mode = ssl.CERT_NONE
    return ctx


def _request(
    method: str,
    url: str,
    cfg: Mapping[str, str],
    data: dict[str, str] | None = None,
    extra_headers: list[tuple[str, str]] | None = None,
) -> tuple[int, bytes]:
    token = (cfg.get("AUTH_TOKEN") or "").strip()
    if token.startswith("#"):
        token = ""
    user = (cfg.get("SPLUNK_USERNAME") or "").strip()
    password = (cfg.get("SPLUNK_PASS") or "").strip()

    headers: list[tuple[str, str]] = list(extra_headers or [])
    if token:
        headers.append(("Authorization", f"Bearer {token}"))
    elif user:
        creds = base64.b64encode(f"{user}:{password}".encode()).decode()
        headers.append(("Authorization", f"Basic {creds}"))
    encoded: bytes | None = None
    if data is not None:
        encoded = urllib.parse.urlencode(data).encode()
        headers.append(("Content-Type", "application/x-www-form-urlencoded"))
    req = urllib.request.Request(url, data=encoded, method=method.upper())
    for k, v in headers:
        req.add_header(k, v)
    try:
        with urllib.request.urlopen(req, context=_ssl_context(), timeout=120) as resp:
            return resp.getcode(), resp.read()
    except urllib.error.HTTPError as e:
        body = e.read() if e.fp else b""
        return e.code, body


def _extract_sid(xml: str) -> str:
    m = re.search(r"<sid>([^<]+)</sid>", xml)
    return m.group(1) if m else ""


def _extract_dispatch_state(xml: str) -> str:
    m = re.search(r'<s:key name="dispatchState">([^<]*)</s:key>', xml)
    return m.group(1) if m else ""


def _format_csv_like_shell(csv_text: str) -> None:
    """Approximate: sed 's/,/ ,/g' | column -t | awk ... | sed underline."""
    s = csv_text.replace(",", " ,")
    reader = csv.reader(io.StringIO(s), delimiter=",", skipinitialspace=True)
    rows = list(reader)
    if not rows:
        print(csv_text)
        return
    col_widths = [0] * max(len(r) for r in rows)
    for r in rows:
        for i, cell in enumerate(r):
            if i < len(col_widths):
                col_widths[i] = max(col_widths[i], len(cell))

    def fmt_row(r: list[str]) -> str:
        parts = []
        for i, w in enumerate(col_widths):
            cell = r[i] if i < len(r) else ""
            parts.append(cell.ljust(w))
        return "  ".join(parts)

    header = fmt_row(rows[0])
    print(header)
    print(re.sub(r"\S", "-", header))
    for r in rows[1:]:
        print(fmt_row(r))


def splunk_search_polling(cfg: Mapping[str, str], search: str, search_level: str = "verbose") -> None:
    host = (cfg.get("SPLUNK_HOST") or "").strip()
    url = f"https://{host}/services/search/jobs"
    _, body = _request(
        "POST",
        url,
        cfg,
        data={"search": search, "adhoc_search_level": search_level},
    )
    xml = body.decode("utf-8", errors="replace")
    sid = _extract_sid(xml)
    print()
    print(f"SID: {sid}")
    if not sid:
        return

    status_url = f"https://{host}/services/search/jobs/{sid}"
    search_status = ""
    counter = 30
    wait_seconds = 1
    while counter > 0:
        _, out = _request("POST", status_url, cfg, data={"adhoc_search_level": search_level})
        status_xml = out.decode("utf-8", errors="replace")
        status = _extract_dispatch_state(status_xml)
        if status == "DONE":
            search_status = "DONE"
            break
        counter -= 1
        time.sleep(wait_seconds)

    if search_status == "DONE":
        results_url = f"https://{host}/services/search/jobs/{sid}/results"
        results_url_q = f"{results_url}?{urllib.parse.urlencode({'output_mode': 'csv', 'adhoc_search_level': search_level})}"
        _, res_body = _request("GET", results_url_q, cfg)
        _format_csv_like_shell(res_body.decode("utf-8", errors="replace"))


def config_reload(cfg: Mapping[str, str], config_name: str) -> None:
    host = (cfg.get("SPLUNK_HOST") or "").strip()
    url = f"https://{host}/servicesNS/-/-/admin/{config_name}/_reload"
    if truthy(cfg, "SHOW_INDEX_CONF"):
        # Mimic curl --write-out for status
        code, _ = _request("POST", url, cfg)
        print(f"{config_name} reload, http-status: {code}")
    else:
        _request("POST", url, cfg)


def cli_auth_args(cfg: Mapping[str, str]) -> list[str]:
    token = (cfg.get("AUTH_TOKEN") or "").strip()
    if token.startswith("#"):
        token = ""
    if token:
        return ["-token", token]
    user = (cfg.get("SPLUNK_USERNAME") or "").strip()
    password = (cfg.get("SPLUNK_PASS") or "").strip()
    return ["-auth", f"{user}:{password}"]


def run_splunk(cfg: Mapping[str, str], args: list[str], *, check: bool = False) -> subprocess.CompletedProcess[str]:
    bin_path = splunk_bin(cfg)
    cmd = [str(bin_path), *args]
    return subprocess.run(
        cmd,
        check=check,
        text=True,
        capture_output=True,
    )


def file_matches(path: Path, file_name: str, extension: str) -> bool:
    fn = (file_name or "").strip()
    ext = (extension or "").strip()
    name = path.name
    if fn:
        return name == fn
    if ext:
        return name.endswith(f".{ext}")
    return False


def iter_load_files(directory: str) -> list[Path]:
    base = Path(directory)
    if not base.is_absolute():
        base = Path(".") / base
    if not base.exists():
        return []
    out: list[Path] = []
    for p in sorted(base.rglob("*")):
        if p.is_file():
            out.append(p)
    return out


def main() -> None:
    parser = argparse.ArgumentParser(description="Load Splunk data (Python port of load-splunk-data.sh).")
    parser.add_argument(
        "config",
        nargs="?",
        help=f"Path to config file (KEY=value). If omitted, prints help and may generate {CFG_FILE} / {SETTINGS_FILE}.",
    )
    args = parser.parse_args()

    if not args.config:
        show_help()
        if not CFG_FILE.is_file():
            print(f"{CFG_FILE} does not exist. Generating now...")
            gen_sample_cfg(CFG_FILE)
        if not SETTINGS_FILE.is_file():
            print(f"{SETTINGS_FILE} does not exist. Generating now...")
            gen_settings(SETTINGS_FILE, sourcetype="")
        sys.exit(0)

    cfg_path = Path(args.config).expanduser()
    print(str(cfg_path))
    if not SETTINGS_FILE.is_file():
        print(f"{SETTINGS_FILE} does not exist. Generating now...")
        gen_settings(SETTINGS_FILE, sourcetype="")
        sys.exit(0)

    cfg = load_config(cfg_path)

    index = (cfg.get("INDEX") or "").strip()
    if index == "main" or index.startswith("_"):
        print(f"Index: {index}")
        print("Invalid index name. Choose an index name that can be deleted and recreated.")
        sys.exit(1)

    splunk_bin(cfg)

    host = (cfg.get("SPLUNK_HOST") or "").strip()
    app = (cfg.get("APP_NAME") or "").strip()
    sourcetype = (cfg.get("SOURCETYPE") or "").strip()
    directory = (cfg.get("DIRECTORY") or ".").strip()
    file_name = (cfg.get("FILE_NAME") or "").strip()
    extension = (cfg.get("EXTENSION") or "").strip()
    index_type = (cfg.get("INDEX_TYPE") or "event").strip()
    index_data_type = index_type or "event"
    report_fields = (cfg.get("REPORT_FIELDS") or "*").strip()
    host_segment = (cfg.get("HOST_SEGMENT") or "").strip()

    dqt = DQT
    pct = PCT
    slb = SLB

    if truthy(cfg, "SHOW_INDEX_CONF"):
        print_section("INDEX_CONF_BEGIN")

    del_url = f"https://{host}/servicesNS/nobody/{app}/data/indexes/{urllib.parse.quote(index)}"
    code_del, _ = _request("DELETE", del_url, cfg)
    if truthy(cfg, "SHOW_INDEX_CONF"):
        print(f"delete index http-status: {code_del}")

    create_url = f"https://{host}/servicesNS/nobody/{app}/data/indexes"
    code_create, _ = _request(
        "POST",
        create_url,
        cfg,
        data={"name": index, "datatype": index_data_type},
    )
    if truthy(cfg, "SHOW_INDEX_CONF"):
        print(f"create index http-status: {code_create}")

    for c in (
        "conf-inputs",
        "conf-fields",
        "conf-transforms",
        "transforms-reload",
        "conf-props",
    ):
        config_reload(cfg, c)
    print()
    if truthy(cfg, "SHOW_INDEX_CONF"):
        print_section("INDEX_CONF_END")

    if truthy(cfg, "SHOW_GREAT_8"):
        print("\n")
        print_section("GREAT_8_BEGIN")
        print(f"Checking the Great 8 Settings in props.conf for sourcetype {sourcetype} in app {app}")
        proc = run_splunk(cfg, ["btool", "props", "list", sourcetype, "--debug", f"--app={app}"])
        props_existing = (proc.stdout or "") + (proc.stderr or "")
        first_line = props_existing.splitlines()[0] if props_existing else ""
        print(first_line)
        for key in PROPS_DESIRED:
            matches = sum(1 for line in props_existing.splitlines() if key in line)
            if matches == 0:
                print(f"Great 8 setting missing: {key}")
        print_section("GREAT_8_END")
        print("\n")

    print_section("ONESHOT_BEGIN")
    auth_tail = cli_auth_args(cfg)
    splunk = splunk_bin(cfg)

    for path in iter_load_files(directory):
        if not file_matches(path, file_name, extension):
            continue
        i = str(path)

        oneshot_args = [
            str(splunk),
            "add",
            "oneshot",
            "-source",
            i,
            "-index",
            index,
            "-sourcetype",
            sourcetype,
        ]
        if host_segment:
            print(f"host segment: -host_segment {host_segment}")
            oneshot_args.extend(["-host_segment", host_segment])
        oneshot_args.extend(auth_tail)

        subprocess.run(oneshot_args, check=False)

        if truthy(cfg, "DEBUG"):
            print("DEBUG waiting a few seconds so errors will be logged")
            time.sleep(3)
            print("\n\nErrors:\n")
            esc_path = f"{dqt}{pct}{i}{dqt}"
            splunk_search_polling(
                cfg,
                "search index=_* OR index=* log_level=ERROR sourcetype=splunkd earliest=-1m "
                f"| where LIKE(data_source,{esc_path}) "
                f"| eval time=strftime(_time, {dqt}{pct}I:{pct}M:{pct}S:{pct}p{dqt}) "
                "| table time,event_message",
            )
            print("\n")

    print_section("ONESHOT_END")

    if truthy(cfg, "SHOW_BTOOL"):
        print_section("BTOOL_BEGIN")
        print("btool check for errors in sourcetype:")
        proc = run_splunk(cfg, ["btool", "check"])
        for line in (proc.stdout or "").splitlines():
            if sourcetype in line:
                print(line)
        print("\n")
        print(f"btool debug of props.conf for sourcetype {sourcetype}:")
        subprocess.run([str(splunk), "btool", "props", "list", sourcetype, "--debug"], check=False)
        print_section("BTOOL_END")

    print("Waiting a few seconds so some of the files will be indexed...")
    time.sleep(3)

    if truthy(cfg, "SHOW_EVENT_SUMMARY") and index_data_type == "event":
        print_section("EVENT_SUMMARY_BEGIN")
        print("\n\nEvent Count:")
        splunk_search_polling(cfg, f"search index={index} sourcetype={sourcetype} | stats count")
        print("\n\nField Summary:\n")
        splunk_search_polling(
            cfg,
            f"search index={index} sourcetype={sourcetype} | fieldsummary | fields field,count",
        )
        print_section("EVENT_SUMMARY_END")

    if truthy(cfg, "SHOW_WALKLEX") and index_data_type == "event":
        print_section("WALKLEX_BEGIN")
        print(f"Rolling hot buckets to warm for index {index}")
        proc = subprocess.run(
            [str(splunk), "_internal", "call", f"/data/indexes/{index}/roll-hot-buckets", *auth_tail],
            text=True,
            capture_output=True,
        )
        for line in (proc.stdout or "").splitlines():
            if "HTTP" in line:
                print(line)
        print("\n")
        search_string = (
            f" |  walklex index={dqt}{index}{dqt} type=field | search NOT field={dqt} *{dqt}  "
            f"| where  NOT LIKE(field,{dqt}date_{pct}{dqt}) "
            "| search NOT field IN ("
            f"{dqt}source{dqt},{dqt}sourcetype{dqt},{dqt}punct{dqt},{dqt}linecount{dqt},"
            f"{dqt}timeendpos{dqt},{dqt}timestartpos{dqt},{dqt}_indextime{dqt},{dqt}snc_io_parser{dqt}) "
            "| rename field as indexed_field "
            "| stats sum(distinct_values) by indexed_field"
        )
        splunk_search_polling(cfg, search_string)
        print_section("WALKLEX_END")

    if truthy(cfg, "DEBUG_TIMESTAMP"):
        print_section("DEBUG_TIMESTAMP_BEGIN")
        splunk_search_polling(
            cfg,
            "search index=_internal sourcetype=splunkd (component=DateParser OR component=DateParserVerbose) earliest=-2m "
            "| transaction component _time log_level | sort _time | table _time,component,log_level,event_message ",
        )
        print("\n\n_time vs _raw:\n")
        splunk_search_polling(
            cfg,
            f"search index={index} sourcetype={sourcetype} earliest=1 | sort - _time | head 2 | table _time,_raw",
        )
        print_section("DEBUG_TIMESTAMP_END")

    if index_data_type == "event":
        print_section("EVENT_SEARCH_BEGIN")
        print(f"search index={index} sourcetype={sourcetype} earliest=-5y | sort - _time | head 20 ")
        splunk_search_polling(
            cfg,
            "search index="
            + index
            + " sourcetype="
            + sourcetype
            + " | sort - _time | head 20 | fields - _raw,index,timestamp,eventtype,punct,splunk_server,splunk_server_group,_bkt,_cd,tag,_sourcetype,_si,_indextime,source,_eventtype_color,linecount,"
            + f"{dqt}tag::eventtype{dqt},date,date_hour,date_mday,date_minute,date_month,date_second,date_wday,date_year,date_zone,_kv,_serial | fields "
            + report_fields
            + " | table *",
        )
        print_section("EVENT_SEARCH_END")
        print("\n")
    else:
        print_section("METRIC_SUMMARY_BEGIN")
        print("\n\nMetric Names:\n")
        search_string = f" | mcatalog values(metric_name) as metric_name  WHERE index={index} | mvexpand metric_name | table metric_name "
        print(search_string)
        splunk_search_polling(cfg, search_string)

        print("\n\nMetric Dimensions:\n")
        search_string = f" | mcatalog values(_dims) AS dimensions  WHERE index={index} | mvexpand dimensions | table dimensions "
        print(search_string)
        splunk_search_polling(cfg, search_string)

        print("\n\nMetric Value Summary by Metric Name:\n")
        search_string = (
            " | mcatalog values(metric_name) as metric_name WHERE index=drive_metrics "
            "| mvexpand metric_name "
            f"| map search={dqt} | mstats avg(_value) as avg_value prestats=false WHERE metric_name={slb}{dqt}{DOL}metric_name{DOL}{slb}{dqt} AND index={slb}{dqt}drive_metrics{slb}{dqt} by metric_name span=1m {dqt} "
            "| stats avg(avg_value) as avg_value,min(_time) as min_time,max(_time) as max_time by metric_name "
            f"| eval min_time=strftime(min_time,{dqt}{pct}m/{pct}d/{pct}y {pct}H:{pct}M{dqt}),max_time=strftime(max_time,{dqt}{pct}m/{pct}d/{pct}y {pct}H:{pct}M{dqt}),avg_value=round(avg_value,1) "
            "| table metric_name,avg_value,min_time,max_time "
        )
        print(search_string)
        splunk_search_polling(cfg, search_string)

        print_section("METRIC_SUMMARY_END")
        print("\n")


if __name__ == "__main__":
    main()