diff --git a/AGENTS.md b/AGENTS.md index e6675c47..aa0e1b27 100644 --- a/AGENTS.md +++ b/AGENTS.md @@ -54,7 +54,7 @@ To skip DNS lookups during testing, set `GITHUB_ACTIONS=true`. - `parsedmarc/types.py` — TypedDict definitions for all report types (`AggregateReport`, `FailureReport`, `SMTPTLSReport`, `ParsingResults`). Legacy alias `ForensicReport = FailureReport` preserved. - `parsedmarc/utils.py` — IP/DNS/GeoIP enrichment, base64 decoding, compression handling - `parsedmarc/mail/` — Polymorphic mail connections: `IMAPConnection`, `GmailConnection`, `MSGraphConnection`, `MaildirConnection` -- `parsedmarc/{elastic,opensearch,splunk,kafkaclient,loganalytics,syslog,s3,webhook,gelf}.py` — Output integrations +- `parsedmarc/{elastic,opensearch,splunk,kafkaclient,loganalytics,syslog,s3,webhook,gelf,loki}.py` — Output integrations ### Report type system diff --git a/CHANGELOG.md b/CHANGELOG.md index 9b5be7d2..d55f190e 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -2,6 +2,11 @@ ## Unreleased +### New features + +- **Grafana Loki output.** A new `[loki]` section pushes every aggregate, failure, and SMTP TLS report record as one JSON log line to the Loki push API (`/loki/api/v1/push`), using the same flat records as the syslog and GELF outputs plus `report_type` and `domain`. Streams are labeled `job=parsedmarc` and `host=`, overridable and extendable with `labels`, plus the record fields named in `label_fields` (default `report_type, domain`). Entries are timestamped at ingest by default, or at the begin of the reported period with `timestamp = report_begin`. Basic authentication, a bearer token, and a tenant ID (`X-Scope-OrgID`) are supported, and large reports are split into pushes of `batch_size` lines. Redirects are not followed, so a 3xx response counts as a failed push. A failed push raises `LokiError`, so it counts as a failed save and holds the batch's mailbox messages back for a retry. Needs only `httpx`, so it is part of the base install. +- **Grafana dashboard for the Loki output.** `dashboards/grafana/Grafana-DMARC_Reports-Loki.json` mirrors the PostgreSQL dashboard's layout with LogQL queries against the Loki output's lines, filtered by the `domain` stream label. The dashboard dev stack now runs Loki, seeds it with the sample reports, and imports the dashboard. + ### Security - **S3 object keys and metadata are now sanitized, and each key ends in a hash of the fields that identify the report** ([GHSA-7wqr-f6g9-cm29](https://github.com/domainaware/parsedmarc/security/advisories/GHSA-7wqr-f6g9-cm29), affects 11.1.0 and earlier). `S3Client.save_report_to_s3()` put the sender-controlled report ID into the object key as-is, and sent `org_name` and the other metadata values as-is. A report ID such as `../../x` added path segments to the key, an ID could carry `\`, `%`, and the other characters AWS lists as "characters to avoid", or control characters, which AWS says "might require special handling", and a long ID could push the key past the 1,024-byte limit ([AWS, "Naming Amazon S3 objects"](https://docs.aws.amazon.com/AmazonS3/latest/userguide/object-keys.html)). botocore rejects non-ASCII metadata with `ParamValidationError`, so a report from an organization with a non-ASCII name was never saved, and its batch was held back and retried. A bare CR or LF in a metadata value made botocore raise `HTTPClientError: Invalid header value` before sending, which held the batch back the same way, while CR or LF followed by a space or tab (an obs-fold) was sent in the `x-amz-meta-*` header; the 2 KB metadata limit ([AWS, "Working with object metadata"](https://docs.aws.amazon.com/AmazonS3/latest/userguide/UsingMetadata.html)) was not enforced. The key also identified a report only by its type, begin day, and report ID, so anyone who could send reports to a monitored mailbox could replace another reporter's object for that day by reusing its report ID. The report ID is now percent-encoded into a single key segment and shortened so the whole key fits in 1,024 bytes. The encoded characters are `/`, every AWS "character to avoid" (`` \ { } ^ [ ] " < > ~ # | ` `` and `%`), control characters, and non-ASCII characters, so `%` appears in a key only as the escape character. The punctuation real report IDs use, such as `:`, `@`, `.`, and `_`, is kept. Each metadata value is percent-encoded (`%`, control characters, and non-ASCII characters) so `urllib.parse.unquote()` reverses it, and is cut to 256 characters, which keeps the total under 2 KB. Each key now ends in `_.json`, where the hash is the first 16 hex characters of a SHA-256 hash of the report type and the fields that identify the report (for aggregate reports: organization name, report ID, policy domain, begin date, and end date; for SMTP TLS reports: organization name, report ID, begin date, end date, and each policy domain). Enrichment such as reverse DNS is not hashed, so a retried report still overwrites its own object. This narrows the overwrite but does not close it: these fields are not secret, and a forged report whose reporter name, report ID, policy domain, and dates all match a real one still replaces it. Turning on S3 Versioning on the bucket keeps both. Compatibility: every report's key changes. Objects saved by earlier versions keep their old keys and are not rewritten, and a batch that is retried across the upgrade is saved again under its new key. Reported by helixnebulatech. diff --git a/README.md b/README.md index 35a9a86f..d7d8c57b 100644 --- a/README.md +++ b/README.md @@ -44,7 +44,7 @@ Please consider [sponsoring my work](https://github.com/sponsors/seanthegeek) if PostgreSQL, for use with premade dashboards - Optionally send the results to Apache Kafka, Amazon S3, Azure Log Analytics (Microsoft Sentinel), a Graylog (GELF) endpoint, a syslog server, - or an HTTP webhook + Grafana Loki, or an HTTP webhook ## Python Compatibility diff --git a/dashboard-dev-bootstrap.sh b/dashboard-dev-bootstrap.sh index 9b76a6dd..95a3ca99 100755 --- a/dashboard-dev-bootstrap.sh +++ b/dashboard-dev-bootstrap.sh @@ -174,6 +174,7 @@ wait_for "OpenSearch Dashboards" \ wait_for "Grafana" curl -sf http://localhost:3000/api/health wait_for "PostgreSQL" \ "${COMPOSE[@]}" exec -T postgresql pg_isready -U "$PG_USER" -d "$PG_DB" +wait_for "Loki" curl -sf http://localhost:3100/ready # Splunk's HEC port is healthy once management API is up too. wait_for "Splunk HEC" curl -ksf https://localhost:8088/services/collector/health # Splunkd management API (used for dashboard imports) lives inside the container. @@ -276,9 +277,10 @@ elif [ -n "$SCHED_SHARING" ]; then fi # --------------------------------------------------------------------------- -# 4. Seed sample data via parsedmarc -> ES, OS, Splunk HEC +# 4. Seed sample data via parsedmarc -> ES, OS, Splunk HEC, PostgreSQL, Loki # Skipped when ES already has aggregate docs from a prior run. Set -# RESEED=1 to wipe ES/OS/Splunk parsedmarc data first and re-seed. +# RESEED=1 to wipe parsedmarc data from every backend first and re-seed +# (Loki is wiped on RESEED=1 even when ES is empty). # --------------------------------------------------------------------------- log "Seeding sample data with parsedmarc-dev.ini" ES_AGG_COUNT=$(curl -sf 'http://localhost:9200/dmarc_aggregate*/_count' 2>/dev/null \ @@ -326,6 +328,15 @@ else -c 'DROP SCHEMA public CASCADE; CREATE SCHEMA public;' >/dev/null 2>&1 || true fi + if [ "${RESEED:-0}" = "1" ]; then + # Loki: the container keeps its data inside itself, so recreating it + # is a clean wipe. Done whenever RESEED=1, even if ES is empty (e.g. + # the ES container was recreated): Loki does not dedupe, so seeding on + # top of leftover logs would append a duplicate copy. + "${COMPOSE[@]}" up -d --force-recreate loki >/dev/null 2>&1 || true + wait_for "Loki (after wipe)" curl -sf http://localhost:3100/ready + fi + # Resolve a Python environment for the seed and make sure parsedmarc plus # every output integration ([all]) and the PostgreSQL extra (psycopg) are # installed in it, so the same run can populate Elasticsearch, OpenSearch, @@ -429,7 +440,10 @@ else echo " ${PARSEDMARC_BIN} — the seed run will fail to initialize its" echo " outputs. Install them with: pip install -e '.[all,postgresql]'" fi - env "${pg_seed_env[@]}" \ + # The Loki output is part of the base install, so it is always seeded. + # Lines keep the default ingest timestamp: Loki rejects samples older + # than a week, and the sample reports are years old. + env "${pg_seed_env[@]}" PARSEDMARC_LOKI_URL=http://localhost:3100 \ "$PARSEDMARC_BIN" -t 2.0 --dns-retries 1 -c parsedmarc-dev.ini "${SAMPLE_FILES[@]}" || true fi @@ -563,6 +577,28 @@ EOF echo " created datasource 'PostgreSQL'" fi +# Loki datasource for the Loki DMARC dashboard. Created with uid dmarc-loki +# when missing; an existing datasource named Loki keeps its own uid, which the +# dashboard import below resolves its ${DS_LOKI} input to. +LOKI_DS_UID=dmarc-loki +loki_ds_resp=$(curl -sS -u "${GRAFANA_USER}:${GRAFANA_PASSWORD}" \ + -w $'\n%{http_code}' \ + "http://localhost:3000/api/datasources/name/Loki") +loki_ds_code="${loki_ds_resp##*$'\n'}" +if [ "$loki_ds_code" = "200" ]; then + LOKI_DS_UID=$(printf '%s' "${loki_ds_resp%$'\n'*}" \ + | python3 -c 'import json, sys; print(json.load(sys.stdin)["uid"])') + echo " datasource 'Loki' already exists (uid ${LOKI_DS_UID}) — skipping" +else + curl -sS -u "${GRAFANA_USER}:${GRAFANA_PASSWORD}" \ + -H 'Content-Type: application/json' \ + -X POST "http://localhost:3000/api/datasources" \ + -d '{"name": "Loki", "uid": "dmarc-loki", "type": "loki", "url": "http://loki:3100", "access": "proxy", "isDefault": false}' \ + | sed 's/^/ /' + echo + echo " created datasource 'Loki'" +fi + log "Importing Grafana dashboard" GF_BODY=$(python3 -c ' import json, sys @@ -597,6 +633,23 @@ curl -sS -u "${GRAFANA_USER}:${GRAFANA_PASSWORD}" \ -X POST "http://localhost:3000/api/dashboards/db" \ -d "$GF_PG_BODY" | sed 's/^/ /' +log "Importing Grafana Loki dashboard" +GF_LOKI_BODY=$(LOKI_DS_UID="$LOKI_DS_UID" python3 -c ' +import json, os +with open("dashboards/grafana/Grafana-DMARC_Reports-Loki.json") as f: + text = f.read() +text = text.replace("${DS_LOKI}", os.environ["LOKI_DS_UID"]) +d = json.loads(text) +d.pop("__inputs", None) +d.pop("__requires", None) +d["id"] = None +print(json.dumps({"dashboard": d, "overwrite": True, "folderUid": ""})) +') +curl -sS -u "${GRAFANA_USER}:${GRAFANA_PASSWORD}" \ + -H 'Content-Type: application/json' \ + -X POST "http://localhost:3000/api/dashboards/db" \ + -d "$GF_LOKI_BODY" | sed 's/^/ /' + log "Importing Splunk dashboard views into the DMARC app" splunk_import_view() { local name="$1" @@ -626,4 +679,5 @@ cat <` + (in a container, the FQDN is the container's hostname) + - `label_fields` - str: Comma-separated record fields used as stream + labels (Default: `report_type, domain`; empty for none). Keep these to + fields with few distinct values — every combination is its own Loki + stream, so fields like `source_ip_address` belong in the line only. + Records without a value for a field get no such label. + - `timestamp` - str: `ingest` (the time of the push) or `report_begin` + (the begin of the reported period, or the arrival date of a failure + report) (Default: `ingest`). Loki rejects entries older than its + `reject_old_samples_max_age` (one week by default) and, within a + stream, entries far older than the newest one, so `report_begin` + loses late or out-of-order reports unless Loki is configured for them. + The reported period is in the line either way (`begin_date`, + `end_date`). + - `username` - str: Basic authentication user + - `password` - str: Basic authentication password + - `token` - str: Bearer token, instead of basic authentication + - `tenant_id` - str: Tenant ID, sent as the `X-Scope-OrgID` header + - `skip_certificate_verification` - bool: Skip TLS certificate + verification (not recommended) + - `timeout` - float: Timeout in seconds for push requests (Default: `60`) + - `batch_size` - int: Maximum log lines per push request + (Default: `1000`) + :::{warning} It is **strongly recommended** to **not** use the `nameservers` setting. By default, `parsedmarc` uses @@ -1070,7 +1117,7 @@ an already-saved report — are unaffected, and S3 is idempotent because each report is written to an object key built from its type, date, report ID, and a hash of the fields that identify the report, so a retry overwrites the same object. Kafka, -Splunk HEC, syslog, GELF, webhooks, Azure Log Analytics, and the +Splunk HEC, syslog, GELF, webhooks, Loki, Azure Log Analytics, and the `--output` JSON/CSV files all append unconditionally, so each retry adds another copy of every report in the batch. That is why the default retry cap is deliberately low: at most three deliveries per report before its @@ -1086,8 +1133,8 @@ syslog and GELF outputs send through Python logging handlers, which swallow delivery errors by design — so an unreachable webhook, syslog, or GELF endpoint is *not* treated as a failed save and does not hold a batch's messages back. Failures in Elasticsearch, OpenSearch, Splunk HEC, -Kafka, S3, PostgreSQL, Azure Log Analytics, and the `--output` files are -all detected and do. +Kafka, S3, PostgreSQL, Loki, Azure Log Analytics, and the `--output` files +are all detected and do. ::: :::{note} @@ -1233,6 +1280,7 @@ For sections with underscores in the name, the full section name is used: | `log_analytics` | `PARSEDMARC_LOG_ANALYTICS_` | | `gelf` | `PARSEDMARC_GELF_` | | `webhook` | `PARSEDMARC_WEBHOOK_` | +| `loki` | `PARSEDMARC_LOKI_` | ## Using parsedmarc as a library diff --git a/parsedmarc/cli.py b/parsedmarc/cli.py index 5ca01d89..9f602ed3 100644 --- a/parsedmarc/cli.py +++ b/parsedmarc/cli.py @@ -36,6 +36,7 @@ from parsedmarc import ( email_results_via_msgraph, get_dmarc_reports_from_mailbox, get_dmarc_reports_from_mbox, + loki, postgres, save_output, splunk, @@ -526,6 +527,7 @@ _KNOWN_SECTIONS = frozenset( "log_analytics", "gelf", "webhook", + "loki", } ) @@ -1516,6 +1518,51 @@ def _parse_config(config: ConfigParser, opts): if "timeout" in webhook_config: opts.webhook_timeout = webhook_config.getint("timeout") + if "loki" in config.sections(): + loki_config = config["loki"] + if "url" in loki_config: + opts.loki_url = loki_config["url"] + else: + raise ConfigurationError("url setting missing from the loki config section") + if "labels" in loki_config: + opts.loki_labels = {} + for pair in _str_to_list(loki_config["labels"]): + name, sep, value = pair.partition("=") + if not sep or not name.strip(): + raise ConfigurationError( + f"Invalid loki label '{pair}', expected name=value" + ) + opts.loki_labels[name.strip()] = value.strip() + if "label_fields" in loki_config: + opts.loki_label_fields = [ + field.strip() + for field in _str_to_list(loki_config["label_fields"]) + if field.strip() + ] + if "timestamp" in loki_config: + opts.loki_timestamp = loki_config["timestamp"] + if opts.loki_timestamp not in loki.TIMESTAMP_MODES: + raise ConfigurationError( + f"Invalid loki timestamp '{opts.loki_timestamp}', must be one " + f"of: {', '.join(loki.TIMESTAMP_MODES)}" + ) + if "username" in loki_config: + opts.loki_username = loki_config["username"] + if "password" in loki_config: + opts.loki_password = loki_config["password"] + if "token" in loki_config: + opts.loki_token = loki_config["token"] + if "tenant_id" in loki_config: + opts.loki_tenant_id = loki_config["tenant_id"] + if "skip_certificate_verification" in loki_config: + opts.loki_skip_certificate_verification = loki_config.getboolean( + "skip_certificate_verification" + ) + if "timeout" in loki_config: + opts.loki_timeout = loki_config.getfloat("timeout") + if "batch_size" in loki_config: + opts.loki_batch_size = loki_config.getint("batch_size") + return index_prefix_domain_map @@ -2045,6 +2092,25 @@ def _build_output_clients(opts, clients, index_prefix_domain_map=None): except Exception as e: raise RuntimeError(f"Webhook: {e}") from e + try: + if opts.loki_url: + logger.debug("Initializing Loki client: url=%s", opts.loki_url) + clients["loki_client"] = loki.LokiClient( + opts.loki_url, + labels=opts.loki_labels, + label_fields=opts.loki_label_fields, + timestamp=opts.loki_timestamp, + username=opts.loki_username, + password=opts.loki_password, + token=opts.loki_token, + tenant_id=opts.loki_tenant_id, + verify=not opts.loki_skip_certificate_verification, + timeout=opts.loki_timeout, + batch_size=opts.loki_batch_size, + ) + except Exception as e: + raise RuntimeError(f"Loki: {e}") from e + # The Log Analytics client is built per batch in process_reports(), # under the same opts.la_dce guard. Checking it here means a missing # extra is reported at startup -- and again on a SIGHUP reload -- @@ -2459,6 +2525,7 @@ def _main(): hec_client = clients.get("hec_client") gelf_client = clients.get("gelf_client") webhook_client = clients.get("webhook_client") + loki_client = clients.get("loki_client") pg_client = clients.get("postgresql_client") kafka_aggregate_topic = opts.kafka_aggregate_topic @@ -2550,6 +2617,12 @@ def _main(): except Exception as error_: log_output_error("Webhook", error_.__str__()) + try: + if loki_client: + loki_client.save_aggregate_report_to_loki(report) + except Exception as error_: + log_output_error("Loki", error_.__str__()) + if hec_client: try: aggregate_reports_ = reports_["aggregate_reports"] @@ -2643,6 +2716,12 @@ def _main(): except Exception as error_: log_output_error("Webhook", error_.__str__()) + try: + if loki_client: + loki_client.save_failure_report_to_loki(report) + except Exception as error_: + log_output_error("Loki", error_.__str__()) + if hec_client: try: failure_reports_ = reports_["failure_reports"] @@ -2736,6 +2815,12 @@ def _main(): except Exception as error_: log_output_error("Webhook", error_.__str__()) + try: + if loki_client: + loki_client.save_smtp_tls_report_to_loki(report) + except Exception as error_: + log_output_error("Loki", error_.__str__()) + if hec_client: try: smtp_tls_reports_ = reports_["smtp_tls_reports"] @@ -3054,6 +3139,17 @@ def _main(): webhook_failure_url=None, webhook_smtp_tls_url=None, webhook_timeout=60, + loki_url=None, + loki_labels=None, + loki_label_fields=None, + loki_timestamp="ingest", + loki_username=None, + loki_password=None, + loki_token=None, + loki_tenant_id=None, + loki_skip_certificate_verification=False, + loki_timeout=60, + loki_batch_size=1000, normalize_timespan_threshold_hours=24.0, postgresql_host=None, postgresql_port=5432, diff --git a/parsedmarc/loki.py b/parsedmarc/loki.py new file mode 100644 index 00000000..4c92f449 --- /dev/null +++ b/parsedmarc/loki.py @@ -0,0 +1,212 @@ +# -*- coding: utf-8 -*- + +from __future__ import annotations + +import json +import re +import socket +import time +from collections.abc import Callable +from typing import Any + +import httpx + +from parsedmarc import ( + parsed_aggregate_reports_to_csv_rows, + parsed_failure_reports_to_csv_rows, + parsed_smtp_tls_reports_to_csv_rows, +) +from parsedmarc.constants import USER_AGENT +from parsedmarc.log import logger +from parsedmarc.types import AggregateReport, FailureReport, SMTPTLSReport +from parsedmarc.utils import human_timestamp_to_unix_timestamp + +PUSH_PATH = "/loki/api/v1/push" +TIMESTAMP_MODES = ("ingest", "report_begin") +DEFAULT_LABEL_FIELDS = ["report_type", "domain"] + +_LABEL_NAME = re.compile(r"^[a-zA-Z_][a-zA-Z0-9_]*$") + +# Per report type: the row field holding the domain the report is about, +# and the row field with its begin time (and whether that value is UTC +# without an offset; aggregate rows carry the record's interval_begin, +# which parsedmarc builds as naive UTC wall-clock time). +_REPORT_TYPES: dict[str, tuple[str, str, bool]] = { + "aggregate": ("domain", "begin_date", True), + "failure": ("reported_domain", "arrival_date_utc", True), + "smtp_tls": ("policy_domain", "begin_date", False), +} + + +class LokiError(RuntimeError): + """Raised when pushing to Loki fails""" + + +class LokiClient(object): + """A client for the Grafana Loki push API""" + + def __init__( + self, + url: str, + *, + labels: dict[str, str] | None = None, + label_fields: list[str] | None = None, + timestamp: str = "ingest", + username: str | None = None, + password: str | None = None, + token: str | None = None, + tenant_id: str | None = None, + verify: bool = True, + timeout: float = 60, + batch_size: int = 1000, + ): + """ + Initializes the LokiClient + + Every report record becomes one JSON log line. Streams are labeled + with ``labels`` plus the record fields named in ``label_fields``; + keep those to fields with few distinct values. + + Args: + url (str): Loki base URL (``/loki/api/v1/push`` is appended when + the URL has no path) or the full push URL + labels (dict): Static stream labels, merged over the defaults + ``job=parsedmarc`` and ``host=`` + label_fields (list): Record fields used as stream labels + (Default: ``report_type``, ``domain``) + timestamp (str): ``ingest`` (time of the push) or + ``report_begin`` (begin of the reported period; Loki rejects + entries older than its ``reject_old_samples_max_age``) + username (str): Basic auth user + password (str): Basic auth password + token (str): Bearer token, instead of basic auth + tenant_id (str): Sent as ``X-Scope-OrgID`` + verify (bool): Verify TLS certificates + timeout (float): Request timeout in seconds + batch_size (int): Maximum log lines per push request + """ + if timestamp not in TIMESTAMP_MODES: + raise ValueError( + f"Invalid Loki timestamp '{timestamp}'. " + f"Must be one of: {', '.join(TIMESTAMP_MODES)}" + ) + if token and (username or password): + raise ValueError("Loki token and username/password are exclusive") + + self.labels = {"job": "parsedmarc", "host": socket.getfqdn()} + self.labels.update(labels or {}) + self.label_fields = list( + DEFAULT_LABEL_FIELDS if label_fields is None else label_fields + ) + for name in list(self.labels) + self.label_fields: + if not _LABEL_NAME.match(name): + raise ValueError(f"Invalid Loki label name '{name}'") + + parsed = httpx.URL(url) + # copy_with(path=...) keeps any query string, userinfo and port intact + self.url = ( + str(parsed.copy_with(path=PUSH_PATH)) if parsed.path in ("", "/") else url + ) + self.timestamp = timestamp + self.timeout = timeout + self.batch_size = max(1, int(batch_size)) + + headers = {"User-Agent": USER_AGENT, "Content-Type": "application/json"} + if token: + headers["Authorization"] = f"Bearer {token}" + if tenant_id: + headers["X-Scope-OrgID"] = tenant_id + auth = httpx.BasicAuth(username or "", password or "") if username else None + # Redirects are not followed: httpx re-sends a 301/302 POST as a GET, + # which the target can answer 200 without storing anything. + self.session = httpx.Client( + headers=headers, auth=auth, verify=verify, follow_redirects=False + ) + + def save_aggregate_report_to_loki( + self, aggregate_reports: AggregateReport | list[AggregateReport] + ): + self._save("aggregate", parsed_aggregate_reports_to_csv_rows, aggregate_reports) + + def save_failure_report_to_loki( + self, failure_reports: FailureReport | list[FailureReport] + ): + self._save("failure", parsed_failure_reports_to_csv_rows, failure_reports) + + def save_smtp_tls_report_to_loki( + self, smtp_tls_reports: SMTPTLSReport | list[SMTPTLSReport] + ): + self._save("smtp_tls", parsed_smtp_tls_reports_to_csv_rows, smtp_tls_reports) + + def _save(self, report_type: str, to_rows: Callable, reports: Any): + rows = to_rows(reports) + if not rows: + return + logger.debug(f"Saving {len(rows)} {report_type} records to Loki") + domain_field, begin_field, begin_is_utc = _REPORT_TYPES[report_type] + now = time.time_ns() + entries = [] + for i, row in enumerate(rows): + line = dict(row) + line["report_type"] = report_type + line.setdefault("domain", row.get(domain_field)) + ts = now + if self.timestamp == "report_begin": + try: + begin = human_timestamp_to_unix_timestamp( + str(row[begin_field]), assume_utc=begin_is_utc + ) + ts = begin * 10**9 + except Exception: + logger.debug(f"Unparsable {begin_field}, using ingest time") + entries.append((self._stream_labels(line), ts + i, line)) + + # Loki only accepts out-of-order entries within half of max_chunk_age + # (so 1h by default) of the newest line in a stream, else it rejects + # them with too_far_behind. Aggregate rows come source-major then per + # day, so sort the whole list (not per push, which would still go + # backwards across batches). + entries.sort(key=lambda entry: entry[1]) + + for start in range(0, len(entries), self.batch_size): + self._push(entries[start : start + self.batch_size]) + + def _stream_labels(self, line: dict[str, Any]) -> tuple[tuple[str, str], ...]: + labels = dict(self.labels) + for field in self.label_fields: + value = line.get(field) + if value is not None and value != "": + labels[field] = str(value) + return tuple(labels.items()) + + def _push(self, entries: list): + streams: dict[tuple, list[list[str]]] = {} + for labels, ts, line in entries: + streams.setdefault(labels, []).append( + [str(ts), json.dumps(line, ensure_ascii=False, default=str)] + ) + body = { + "streams": [ + {"stream": dict(labels), "values": values} + for labels, values in streams.items() + ] + } + try: + response = self.session.post( + self.url, content=json.dumps(body), timeout=self.timeout + ) + except httpx.HTTPError as error_: + raise LokiError(f"Loki push to {self.url} failed: {error_}") from error_ + # Loki answers 204 when it stored the push; 200 is also accepted. + # Other 2xx codes are not success: with block_ingestion_until set it + # returns a custom 260 and discards the logs + # (block_ingestion_status_code). + if response.status_code not in (200, 204): + raise LokiError( + f"Loki push to {self.url} failed with HTTP " + f"{response.status_code}: {response.text[:500]}" + ) + + def close(self): + """Close the underlying HTTP session.""" + self.session.close() diff --git a/tests/test_cli.py b/tests/test_cli.py index b6a4af52..66500f8a 100644 --- a/tests/test_cli.py +++ b/tests/test_cli.py @@ -8963,6 +8963,17 @@ def _output_client_opts(**overrides) -> SimpleNamespace: "kafka_skip_certificate_verification", "kafka_username", "la_dce", + "loki_batch_size", + "loki_label_fields", + "loki_labels", + "loki_password", + "loki_skip_certificate_verification", + "loki_tenant_id", + "loki_timeout", + "loki_timestamp", + "loki_token", + "loki_url", + "loki_username", "opensearch_api_key", "opensearch_auth_type", "opensearch_aws_region", @@ -9263,5 +9274,191 @@ class TestOptionalIntegrationExtras(unittest.TestCase): ) +class TestParseConfigLoki(unittest.TestCase): + def test_loki_complete(self): + from parsedmarc.cli import _parse_config + + cp = _config_with( + "loki", + { + "url": "https://loki.example.com", + "labels": "host=mail2.example.com, env = prod", + "label_fields": "report_type, domain, org_name", + "timestamp": "report_begin", + "username": "u", + "password": "p", + "tenant_id": "fosdem", + "skip_certificate_verification": "true", + "timeout": "12.5", + "batch_size": "500", + }, + ) + opts = _opts() + _parse_config(cp, opts) + self.assertEqual(opts.loki_url, "https://loki.example.com") + self.assertEqual(opts.loki_labels, {"host": "mail2.example.com", "env": "prod"}) + self.assertEqual(opts.loki_label_fields, ["report_type", "domain", "org_name"]) + self.assertEqual(opts.loki_timestamp, "report_begin") + self.assertEqual(opts.loki_username, "u") + self.assertEqual(opts.loki_password, "p") + self.assertEqual(opts.loki_tenant_id, "fosdem") + self.assertTrue(opts.loki_skip_certificate_verification) + self.assertEqual(opts.loki_timeout, 12.5) + self.assertEqual(opts.loki_batch_size, 500) + + def test_loki_token(self): + from parsedmarc.cli import _parse_config + + cp = _config_with("loki", {"url": "http://l:3100", "token": "t"}) + opts = _opts() + _parse_config(cp, opts) + self.assertEqual(opts.loki_token, "t") + + def test_loki_empty_label_fields_means_no_dynamic_labels(self): + from parsedmarc.cli import _parse_config + + cp = _config_with("loki", {"url": "http://l:3100", "label_fields": ""}) + opts = _opts() + _parse_config(cp, opts) + self.assertEqual(opts.loki_label_fields, []) + + def test_loki_label_fields_strips_whitespace_around_each_field(self): + """Padding around commas must not leak into label names (it made + LokiClient raise ``Invalid Loki label name 'report_type '``).""" + from parsedmarc.cli import _parse_config + + cp = _config_with( + "loki", + {"url": "http://l:3100", "label_fields": "report_type , domain ,org_name "}, + ) + opts = _opts() + _parse_config(cp, opts) + self.assertEqual(opts.loki_label_fields, ["report_type", "domain", "org_name"]) + + def test_loki_missing_url_raises(self): + from parsedmarc.cli import ConfigurationError, _parse_config + + cp = _config_with("loki", {"labels": "a=b"}) + with self.assertRaises(ConfigurationError): + _parse_config(cp, _opts()) + + def test_loki_malformed_labels_raise(self): + from parsedmarc.cli import ConfigurationError, _parse_config + + cp = _config_with("loki", {"url": "http://l:3100", "labels": "host"}) + with self.assertRaises(ConfigurationError): + _parse_config(cp, _opts()) + + def test_loki_invalid_timestamp_raises(self): + from parsedmarc.cli import ConfigurationError, _parse_config + + cp = _config_with("loki", {"url": "http://l:3100", "timestamp": "never"}) + with self.assertRaises(ConfigurationError): + _parse_config(cp, _opts()) + + def test_loki_env_section_resolves(self): + from parsedmarc.cli import _resolve_section_key + + self.assertEqual(_resolve_section_key("LOKI_URL"), ("loki", "url")) + self.assertEqual(_resolve_section_key("LOKI_PASSWORD"), ("loki", "password")) + + +class TestLokiCliWiring(unittest.TestCase): + """End-to-end: a [loki] config reaches LokiClient, every report type + reaches its save method, and a failed push counts as an output error + (so mailbox messages are held back instead of archived).""" + + def _run_main(self, reports, extra_general="", save_side_effect=None): + config = f"""[general] +save_aggregate = true +save_failure = true +save_smtp_tls = true +silent = true +{extra_general} + +[imap] +host = imap.example.com +user = test-user +password = test-password + +[loki] +url = http://loki.example.com:3100 +labels = host=mail2.example.com +timestamp = report_begin +tenant_id = fosdem +""" + with tempfile.NamedTemporaryFile( + "w", suffix=".ini", delete=False + ) as config_file: + config_file.write(config) + config_path = config_file.name + self.addCleanup(lambda: os.path.exists(config_path) and os.remove(config_path)) + + with ( + patch("parsedmarc.cli.loki.LokiClient") as mock_client_cls, + patch( + "parsedmarc.cli.get_dmarc_reports_from_mailbox", + side_effect=_fetch_invoking_save_callback(reports), + ), + patch("parsedmarc.cli.IMAPConnection", return_value=object()), + patch.object(sys, "argv", ["parsedmarc", "-c", config_path]), + ): + client = mock_client_cls.return_value + if save_side_effect is not None: + for m in ( + "save_aggregate_report_to_loki", + "save_failure_report_to_loki", + "save_smtp_tls_report_to_loki", + ): + getattr(client, m).side_effect = save_side_effect + parsedmarc.cli._main() + return mock_client_cls + + def _reports(self): + return { + "aggregate_reports": [{"policy_published": {"domain": "example.com"}}], + "failure_reports": [{"reported_domain": "example.com"}], + "smtp_tls_reports": [{"organization_name": "Org", "policies": []}], + } + + def test_client_constructed_from_config(self): + cls = self._run_main(self._reports()) + cls.assert_called_once() + self.assertEqual(cls.call_args.args, ("http://loki.example.com:3100",)) + kwargs = cls.call_args.kwargs + self.assertEqual(kwargs["labels"], {"host": "mail2.example.com"}) + self.assertIsNone(kwargs["label_fields"]) + self.assertEqual(kwargs["timestamp"], "report_begin") + self.assertEqual(kwargs["tenant_id"], "fosdem") + self.assertTrue(kwargs["verify"]) + self.assertEqual(kwargs["timeout"], 60) + self.assertEqual(kwargs["batch_size"], 1000) + + def test_all_report_types_saved_and_client_closed(self): + reports = self._reports() + client = self._run_main(reports).return_value + client.save_aggregate_report_to_loki.assert_called_once_with( + reports["aggregate_reports"][0] + ) + client.save_failure_report_to_loki.assert_called_once_with( + reports["failure_reports"][0] + ) + client.save_smtp_tls_report_to_loki.assert_called_once_with( + reports["smtp_tls_reports"][0] + ) + client.close.assert_called() + + def test_loki_error_is_an_output_error(self): + from parsedmarc.loki import LokiError + + with self.assertRaises(SystemExit) as ctx: + self._run_main( + self._reports(), + extra_general="fail_on_output_error = true", + save_side_effect=LokiError("HTTP 503"), + ) + self.assertEqual(ctx.exception.code, 1) + + if __name__ == "__main__": unittest.main(verbosity=2) diff --git a/tests/test_grafana_loki_dashboard.py b/tests/test_grafana_loki_dashboard.py new file mode 100644 index 00000000..444b79c2 --- /dev/null +++ b/tests/test_grafana_loki_dashboard.py @@ -0,0 +1,347 @@ +"""Tests for the Grafana dashboard for the Loki output""" + +import json +import os +import re +import subprocess +import tempfile +import unittest +from typing import Any, cast +from unittest.mock import patch + +import parsedmarc +from parsedmarc.loki import LokiClient + +ROOT = os.path.join(os.path.dirname(__file__), "..") +GRAFANA = os.path.join(ROOT, "dashboards", "grafana") +SAMPLES = os.path.join(ROOT, "samples") +LOKI_DASHBOARD = os.path.join(GRAFANA, "Grafana-DMARC_Reports-Loki.json") +PG_DASHBOARD = os.path.join(GRAFANA, "Grafana-DMARC_Reports-PostgreSQL.json") +DATASOURCE = {"type": "loki", "uid": "${DS_LOKI}"} + +SELECTOR = re.compile(r"\{([^{}]*report_type=\"(\w+)\"[^{}]*)\}") +JSON_STAGE = re.compile(r"\|\s*json\b([^|\[)]*)") +LABEL_FORMAT = re.compile(r"label_format\s+(\w+)=") + + +def _load(path: str) -> dict[str, Any]: + with open(path) as f: + return json.load(f) + + +def _panels(dashboard: dict[str, Any]): + for panel in dashboard["panels"]: + yield panel + yield from panel.get("panels", []) + + +def _targets(dashboard: dict[str, Any]): + for panel in _panels(dashboard): + for target in panel.get("targets", []): + yield panel, target + + +def _lines(report_type: str, path: str) -> list[dict[str, Any]]: + """Log lines the Loki output writes for one sample report.""" + parsed = parsedmarc.parse_report_file(path, offline=True) + client = LokiClient("http://loki:3100") + pushed: list = [] + with patch.object(client, "_push", side_effect=pushed.extend): + save = getattr(client, f"save_{report_type}_report_to_loki") + save(cast(Any, parsed["report"])) + return [line for _, _, line in pushed] + + +def _sample_fields() -> dict[str, set[str]]: + samples = { + "aggregate": os.path.join( + SAMPLES, "aggregate", "!example.com!1538204542!1538463818.xml" + ), + "failure": os.path.join(SAMPLES, "failure", "dmarc_ruf_report_linkedin.eml"), + "smtp_tls": os.path.join(SAMPLES, "smtp_tls", "google.com_smtp_tls_report.eml"), + } + fields = {} + for report_type, path in samples.items(): + fields[report_type] = {k for line in _lines(report_type, path) for k in line} + # Columns the CSV rows only carry when the report has them. + fields["aggregate"] |= {"org_extra_contact_info", "fo", "np"} + fields["failure"] |= {"sample_headers_only"} + fields["smtp_tls"] |= { + "result_type", + "sending_mta_ip", + "receiving_ip", + "receiving_mx_hostname", + "failure_reason_code", + "additional_info_uri", + "policy_strings", + "mx_host_patterns", + } + return fields + + +class TestGrafanaLokiDashboard(unittest.TestCase): + @classmethod + def setUpClass(cls): + cls.dashboard = _load(LOKI_DASHBOARD) + cls.postgres = _load(PG_DASHBOARD) + + def test_opens_on_the_last_seven_days(self): + self.assertEqual(self.dashboard["time"], {"from": "now-7d", "to": "now"}) + + def test_declares_the_loki_datasource_input(self): + [entry] = self.dashboard["__inputs"] + self.assertEqual(entry["name"], "DS_LOKI") + self.assertEqual(entry["pluginId"], "loki") + + def test_same_layout_as_the_postgresql_dashboard(self): + def layout(d): + return [(p["type"], p.get("title"), p["gridPos"]) for p in _panels(d)] + + self.assertEqual(layout(self.dashboard), layout(self.postgres)) + + def test_every_query_uses_the_loki_input(self): + for panel, target in _targets(self.dashboard): + with self.subTest(panel=panel["title"]): + self.assertEqual(panel["datasource"], DATASOURCE) + self.assertEqual(target["datasource"], DATASOURCE) + + def test_every_query_selects_parsedmarc_streams(self): + for panel, target in _targets(self.dashboard): + with self.subTest(panel=panel["title"], ref=target["refId"]): + match = SELECTOR.search(target["expr"]) + self.assertIsNotNone(match, target["expr"]) + assert match is not None + self.assertIn('job="parsedmarc"', match.group(1)) + self.assertIn(match.group(2), ("aggregate", "failure", "smtp_tls")) + + def test_aggregate_queries_follow_the_domain_variable(self): + for panel, target in _targets(self.dashboard): + match = SELECTOR.search(target["expr"]) + if match and match.group(2) == "aggregate": + with self.subTest(panel=panel["title"]): + self.assertIn('domain=~"$domain"', match.group(1)) + + def test_json_stages_name_their_fields(self): + """A bare `| json` turns every field into a label and hits max_query_series.""" + for panel, target in _targets(self.dashboard): + for stage in JSON_STAGE.findall(target["expr"]): + with self.subTest(panel=panel["title"], ref=target["refId"]): + self.assertTrue(stage.strip(), target["expr"]) + + def test_fields_exist_in_the_lines_the_loki_output_writes(self): + fields = _sample_fields() + for panel, target in _targets(self.dashboard): + expr = target["expr"] + match = SELECTOR.search(expr) + assert match is not None + known = fields[match.group(2)] | set(LABEL_FORMAT.findall(expr)) + used = { + f.strip() + for stage in JSON_STAGE.findall(expr) + for f in stage.split(",") + } + with self.subTest(panel=panel["title"], ref=target["refId"]): + self.assertEqual(used - known, set()) + + def test_extracted_columns_exist_in_the_log_lines(self): + fields = _sample_fields() + for panel in _panels(self.dashboard): + for t in panel.get("transformations", []): + if t["id"] != "extractFields": + continue + match = SELECTOR.search(panel["targets"][0]["expr"]) + assert match is not None + report_type = match.group(2) + paths = {p["path"] for p in t["options"]["jsonPaths"]} + with self.subTest(panel=panel["title"]): + self.assertEqual(paths - fields[report_type], set()) + + def test_range_queries_keep_their_buckets_inside_the_time_range(self): + """A fixed 1d step stamps the current bucket after `now`, so the chart is empty.""" + for panel, target in _targets(self.dashboard): + if target.get("queryType") == "range" and "unwrap" in target["expr"]: + with self.subTest(panel=panel["title"], ref=target["refId"]): + self.assertNotIn("step", target) + self.assertNotIn("$interval", target["expr"]) + + def test_report_dedup_groups_by_reporter_too(self): + """report_id is only unique per reporter, so every `by (...)` that + groups on it must group on org_name as well.""" + for panel, target in _targets(self.dashboard): + for clause in re.findall(r"\bby \(([^)]*)\)", target["expr"]): + names = [n.strip() for n in clause.split(",")] + if "report_id" in names: + with self.subTest(panel=panel["title"], ref=target["refId"]): + self.assertIn("org_name", names) + + def test_fixed_lookback_series_have_no_additive_legend_calcs(self): + """A fixed range such as [1d] with the automatic step yields + overlapping rolling totals, so summing them counts each value many times.""" + for panel in _panels(self.dashboard): + fixed = False + for target in panel.get("targets", []): + if target.get("queryType") != "range": + continue + for dur in re.findall(r"\[([^\]]+)\]", target["expr"]): + if dur not in ("$__interval", "$__auto"): + fixed = True + if fixed and panel.get("type") == "timeseries": + calcs = panel["options"]["legend"].get("calcs", []) + with self.subTest(panel=panel["title"]): + self.assertFalse({"sum", "count"} & set(calcs)) + + def test_pies_use_one_slice_per_row(self): + for panel in _panels(self.dashboard): + if panel["type"] == "piechart": + with self.subTest(panel=panel["title"]): + self.assertTrue(panel["options"]["reduceOptions"]["values"]) + + def test_table_counts_are_named_after_what_they_count(self): + """Grafana shows the value field as "Value #A" unless a display name says otherwise.""" + for panel in _panels(self.dashboard): + if panel["type"] != "table": + continue + names = set() + for t in panel.get("transformations", []): + if t["id"] == "organize": + names.update( + dst + for src, dst in t["options"].get("renameByName", {}).items() + if src.startswith("Value") + ) + if not names: + continue + shown = { + prop["value"] + for o in panel.get("fieldConfig", {}).get("overrides", []) + for prop in o["properties"] + if prop["id"] == "displayName" + } + with self.subTest(panel=panel["title"]): + self.assertEqual(names - shown, set()) + + def test_summary_tables_combine_one_frame_per_series(self): + for panel in _panels(self.dashboard): + ids = [t["id"] for t in panel.get("transformations", [])] + if "labelsToFields" in ids: + with self.subTest(panel=panel["title"]): + self.assertTrue({"merge", "joinByField"} & set(ids)) + + def test_tls_session_counts_only_sum_policy_rows(self): + """Failure-detail rows repeat failed_session_count; summing them double counts.""" + for panel, target in _targets(self.dashboard): + if re.search(r"unwrap\s+\w+_session_count", target["expr"]): + with self.subTest(panel=panel["title"], ref=target["refId"]): + self.assertIn('policy_type!=""', target["expr"]) + + +BOOTSTRAP = os.path.join(ROOT, "dashboard-dev-bootstrap.sh") + +# Stands in for the Grafana HTTP API: the Loki datasource lookup answers from +# FAKE_LOKI_UID (404 when unset) and every request body is appended to +# FAKE_POSTS, one per line. +FAKE_CURL = r"""#!/bin/bash +url="" data="" write="" +while [ $# -gt 0 ]; do + case "$1" in + -w) write="$2"; shift ;; + -d) data="$2"; shift ;; + -u|-H|-X|-o) shift ;; + -*) ;; + *) url="$1" ;; + esac + shift +done +if [ -n "$data" ]; then + printf '%s\n' "$data" >> "$FAKE_POSTS" +fi +if [[ "$url" == */api/datasources/name/Loki ]]; then + if [ -n "${FAKE_LOKI_UID:-}" ]; then + printf '{"name": "Loki", "uid": "%s"}' "$FAKE_LOKI_UID" + code=200 + else + printf '{"message": "Data source not found"}' + code=404 + fi + [ -n "$write" ] && printf '\n%s' "$code" +fi +exit 0 +""" + + +class TestBootstrapLokiDatasourceUid(unittest.TestCase): + """The dev bootstrap must bind the imported dashboard to whichever uid the + 'Loki' datasource really has, not assume the uid it creates.""" + + def _import(self, existing_uid: str | None) -> tuple[str, list[str]]: + """Run the bootstrap's Loki datasource and dashboard import steps + against a fake curl. Returns the datasource uid the imported + dashboard uses and the POSTed request bodies.""" + with open(BOOTSTRAP) as f: + script = f.read() + start = script.index("# Loki datasource for the Loki DMARC dashboard") + end = script.index('log "Importing Splunk dashboard views') + with tempfile.TemporaryDirectory() as tmp: + curl = os.path.join(tmp, "curl") + with open(curl, "w") as f: + f.write(FAKE_CURL) + os.chmod(curl, 0o755) + posts = os.path.join(tmp, "posts") + env = { + **os.environ, + "PATH": tmp + os.pathsep + os.environ["PATH"], + "FAKE_POSTS": posts, + "GRAFANA_USER": "admin", + "GRAFANA_PASSWORD": "pw", + } + env.pop("FAKE_LOKI_UID", None) + if existing_uid: + env["FAKE_LOKI_UID"] = existing_uid + prelude = "set -euo pipefail\nlog() { :; }\n" + subprocess.run( + ["bash", "-c", prelude + script[start:end]], + cwd=ROOT, + env=env, + check=True, + capture_output=True, + text=True, + ) + with open(posts) as f: + bodies = f.read().splitlines() + imported = [b for b in bodies if '"dashboard"' in b] + loki = [ + b + for b in imported + if json.loads(b)["dashboard"]["uid"] == self._loki_dashboard_uid() + ] + self.assertEqual(len(loki), 1) + uids = { + t["datasource"]["uid"] + for _, t in _targets(json.loads(loki[0])["dashboard"]) + } + self.assertEqual(len(uids), 1) + return uids.pop(), bodies + + @staticmethod + def _created(bodies: list[str]) -> list[dict[str, Any]]: + """Loki datasources the script POSTed (dashboard imports excluded).""" + parsed = [json.loads(b) for b in bodies] + return [d for d in parsed if d.get("type") == "loki"] + + def _loki_dashboard_uid(self) -> str: + return _load(LOKI_DASHBOARD)["uid"] + + def test_existing_datasource_uid_is_used(self): + uid, bodies = self._import("custom-loki-uid") + self.assertEqual(uid, "custom-loki-uid") + # An existing datasource is not created again. + self.assertEqual(self._created(bodies), []) + + def test_created_datasource_uses_dmarc_loki(self): + uid, bodies = self._import(None) + self.assertEqual(uid, "dmarc-loki") + self.assertEqual([d["uid"] for d in self._created(bodies)], ["dmarc-loki"]) + + +if __name__ == "__main__": + unittest.main(verbosity=2) diff --git a/tests/test_loki.py b/tests/test_loki.py new file mode 100644 index 00000000..d50f016c --- /dev/null +++ b/tests/test_loki.py @@ -0,0 +1,472 @@ +"""Tests for parsedmarc.loki""" + +import json +import os +import time +import unittest +from typing import cast +from unittest.mock import MagicMock, patch + +import httpx + +import parsedmarc +from parsedmarc.loki import LokiClient, LokiError +from parsedmarc.types import AggregateReport, FailureReport, SMTPTLSReport +from parsedmarc.utils import human_timestamp_to_unix_timestamp + +# Placeholders: the row conversion is patched in most tests. +AGGREGATE = cast(AggregateReport, {}) +FAILURE = cast(FailureReport, {}) +SMTP_TLS = cast(SMTPTLSReport, {}) + +SAMPLES = os.path.join(os.path.dirname(__file__), "..", "samples") + +AGGREGATE_ROWS = [ + { + "org_name": "google.com", + "report_id": "r1", + "begin_date": "2026-10-20 00:00:00", + "end_date": "2026-10-20 23:59:59", + "domain": "example.org", + "source_ip_address": "192.0.2.1", + "count": 3, + "dmarc_aligned": True, + }, + { + "org_name": "google.com", + "report_id": "r1", + "begin_date": "2026-10-20 00:00:00", + "end_date": "2026-10-20 23:59:59", + "domain": "lists.example.org", + "source_ip_address": "192.0.2.2", + "count": 1, + "dmarc_aligned": False, + }, +] +FAILURE_ROWS = [ + { + "reported_domain": "example.org", + "arrival_date_utc": "2026-10-20 09:20:27", + "source_ip_address": "192.0.2.3", + } +] +SMTP_TLS_ROWS = [ + { + "organization_name": "Reporter", + "begin_date": "2026-10-20T00:00:00Z", + "policy_domain": "example.org", + "failed_session_count": 1, + } +] + + +def _ok_response(status_code=204): + response = MagicMock() + response.status_code = status_code + response.text = "" + return response + + +def _client(**kwargs): + kwargs.setdefault("labels", {"host": "mail.example.org"}) + client = LokiClient("http://loki.example.org:3100", **kwargs) + client.session = MagicMock() + client.session.post.return_value = _ok_response() + return client + + +def _session(client: LokiClient) -> MagicMock: + return cast(MagicMock, client.session) + + +def _pushes(client): + """The JSON bodies of every push request, in order.""" + return [json.loads(c.kwargs["content"]) for c in client.session.post.call_args_list] + + +def _entries(client): + """(labels, timestamp_ns, line dict) of every pushed entry.""" + out = [] + for body in _pushes(client): + for stream in body["streams"]: + for ts, line in stream["values"]: + out.append((stream["stream"], int(ts), json.loads(line))) + return out + + +class TestLokiClientInit(unittest.TestCase): + def test_push_path_appended_to_base_url(self): + for url in ("http://loki:3100", "http://loki:3100/"): + client = LokiClient(url) + self.assertEqual(client.url, "http://loki:3100/loki/api/v1/push") + + def test_push_path_appended_before_query_string(self): + for url in ( + "https://gw.example.org?route=logs", + "https://gw.example.org/?route=logs", + ): + client = LokiClient(url) + self.assertEqual( + client.url, "https://gw.example.org/loki/api/v1/push?route=logs" + ) + + def test_explicit_push_path_kept(self): + client = LokiClient("https://gw.example.org/custom/push") + self.assertEqual(client.url, "https://gw.example.org/custom/push") + + def test_default_labels_job_and_fqdn_host(self): + with patch("parsedmarc.loki.socket.getfqdn", return_value="mx1.example.org"): + client = LokiClient("http://loki:3100") + self.assertEqual( + client.labels, {"job": "parsedmarc", "host": "mx1.example.org"} + ) + + def test_configured_labels_override_and_extend_defaults(self): + client = LokiClient("http://loki:3100", labels={"host": "mail2", "env": "prod"}) + self.assertEqual(client.labels["host"], "mail2") + self.assertEqual(client.labels["env"], "prod") + self.assertEqual(client.labels["job"], "parsedmarc") + + def test_default_label_fields(self): + self.assertEqual( + LokiClient("http://loki:3100").label_fields, ["report_type", "domain"] + ) + + def test_invalid_label_names_rejected(self): + with self.assertRaises(ValueError): + LokiClient("http://loki:3100", labels={"bad-name": "x"}) + with self.assertRaises(ValueError): + LokiClient("http://loki:3100", label_fields=["1field"]) + + def test_invalid_timestamp_mode_rejected(self): + with self.assertRaises(ValueError): + LokiClient("http://loki:3100", timestamp="report_end") + + def test_token_and_basic_auth_are_exclusive(self): + with self.assertRaises(ValueError): + LokiClient("http://loki:3100", username="u", password="p", token="t") + + def test_headers(self): + client = LokiClient("http://loki:3100", token="secret", tenant_id="fosdem") + headers = client.session.headers + self.assertEqual(headers["Content-Type"], "application/json") + self.assertIn("parsedmarc", headers["User-Agent"]) + self.assertEqual(headers["Authorization"], "Bearer secret") + self.assertEqual(headers["X-Scope-OrgID"], "fosdem") + + def test_no_auth_headers_by_default(self): + headers = LokiClient("http://loki:3100").session.headers + self.assertNotIn("Authorization", headers) + self.assertNotIn("X-Scope-OrgID", headers) + + def test_basic_auth(self): + client = LokiClient("http://loki:3100", username="u", password="p") + self.assertIsInstance(client.session.auth, httpx.BasicAuth) + + +@patch( + "parsedmarc.loki.parsed_aggregate_reports_to_csv_rows", return_value=AGGREGATE_ROWS +) +class TestLokiAggregate(unittest.TestCase): + def test_one_line_per_record_with_report_type(self, _rows): + client = _client() + client.save_aggregate_report_to_loki(AGGREGATE) + lines = [line for _, _, line in _entries(client)] + self.assertEqual(len(lines), 2) + for line, row in zip(lines, AGGREGATE_ROWS): + self.assertEqual(line, {**row, "report_type": "aggregate"}) + + def test_streams_split_by_label_fields(self, _rows): + client = _client() + client.save_aggregate_report_to_loki(AGGREGATE) + streams = sorted( + (s["stream"] for s in _pushes(client)[0]["streams"]), + key=lambda s: s["domain"], + ) + self.assertEqual( + streams, + [ + { + "job": "parsedmarc", + "host": "mail.example.org", + "report_type": "aggregate", + "domain": "example.org", + }, + { + "job": "parsedmarc", + "host": "mail.example.org", + "report_type": "aggregate", + "domain": "lists.example.org", + }, + ], + ) + + def test_custom_label_fields(self, _rows): + client = _client(label_fields=["org_name"]) + client.save_aggregate_report_to_loki(AGGREGATE) + streams = _pushes(client)[0]["streams"] + self.assertEqual(len(streams), 1) + self.assertEqual(streams[0]["stream"]["org_name"], "google.com") + self.assertNotIn("domain", streams[0]["stream"]) + + def test_empty_or_missing_label_values_are_left_out(self, rows): + rows.return_value = [{"domain": "", "org_name": None}] + client = _client(label_fields=["domain", "org_name", "nope"]) + client.save_aggregate_report_to_loki(AGGREGATE) + labels = _pushes(client)[0]["streams"][0]["stream"] + self.assertEqual(set(labels), {"job", "host"}) + + def test_ingest_timestamps_are_now_and_unique(self, _rows): + client = _client() + with patch("parsedmarc.loki.time.time_ns", return_value=1_000_000_000_000): + client.save_aggregate_report_to_loki(AGGREGATE) + stamps = sorted(ts for _, ts, _ in _entries(client)) + self.assertEqual(stamps, [1_000_000_000_000, 1_000_000_000_001]) + + def test_report_begin_timestamps(self, _rows): + client = _client(timestamp="report_begin") + client.save_aggregate_report_to_loki(AGGREGATE) + begin = ( + human_timestamp_to_unix_timestamp("2026-10-20 00:00:00", assume_utc=True) + * 10**9 + ) + stamps = sorted(ts for _, ts, _ in _entries(client)) + self.assertEqual(stamps, [begin, begin + 1]) + + def test_report_begin_falls_back_to_ingest_time(self, rows): + rows.return_value = [{"domain": "example.org", "begin_date": "garbage"}] + client = _client(timestamp="report_begin") + with patch("parsedmarc.loki.time.time_ns", return_value=42): + client.save_aggregate_report_to_loki(AGGREGATE) + self.assertEqual([ts for _, ts, _ in _entries(client)], [42]) + + def test_timestamps_are_strings(self, _rows): + client = _client() + client.save_aggregate_report_to_loki(AGGREGATE) + for stream in _pushes(client)[0]["streams"]: + for ts, line in stream["values"]: + self.assertIsInstance(ts, str) + self.assertIsInstance(line, str) + + def test_large_reports_split_into_batches(self, rows): + rows.return_value = [ + dict(AGGREGATE_ROWS[0], report_id=str(i)) for i in range(5) + ] + client = _client(batch_size=2) + client.save_aggregate_report_to_loki(AGGREGATE) + sizes = [ + sum(len(s["values"]) for s in body["streams"]) for body in _pushes(client) + ] + self.assertEqual(sizes, [2, 2, 1]) + self.assertEqual(len({ts for _, ts, _ in _entries(client)}), 5) + + def test_nothing_sent_without_rows(self, rows): + rows.return_value = [] + client = _client() + client.save_aggregate_report_to_loki([]) + _session(client).post.assert_not_called() + + def test_posts_to_push_url_with_timeout(self, _rows): + client = _client(timeout=12) + client.save_aggregate_report_to_loki(AGGREGATE) + call = _session(client).post.call_args + self.assertEqual(call.args[0], "http://loki.example.org:3100/loki/api/v1/push") + self.assertEqual(call.kwargs["timeout"], 12) + + +class TestLokiOrdering(unittest.TestCase): + def test_report_begin_entries_sent_in_timestamp_order(self): + """Loki rejects entries more than max_chunk_age/2 behind the newest + line of a stream (too_far_behind), so entries must be pushed oldest + first even though normalized multi-day aggregate rows arrive + source-major: all days of one source, then back to day 1 for the + next source.""" + rows = [ + {"source_ip_address": ip, "begin_date": f"2024-01-{day:02d} 00:00:00"} + for ip in ("192.0.2.1", "192.0.2.2") + for day in (1, 2, 3) + ] + client = _client(timestamp="report_begin", batch_size=2) + with patch( + "parsedmarc.loki.parsed_aggregate_reports_to_csv_rows", return_value=rows + ): + client.save_aggregate_report_to_loki(cast(AggregateReport, {})) + stamps = [ts for _, ts, _ in _entries(client)] + self.assertEqual(len(stamps), 6) + self.assertEqual(len(_session(client).post.call_args_list), 3) + self.assertEqual(stamps, sorted(stamps)) + + +class TestLokiOtherReportTypes(unittest.TestCase): + @patch( + "parsedmarc.loki.parsed_failure_reports_to_csv_rows", return_value=FAILURE_ROWS + ) + def test_failure_domain_label_and_utc_begin(self, _rows): + client = _client(timestamp="report_begin") + client.save_failure_report_to_loki(FAILURE) + [(labels, ts, line)] = _entries(client) + self.assertEqual(labels["report_type"], "failure") + self.assertEqual(labels["domain"], "example.org") + self.assertEqual(line["domain"], "example.org") + expected = human_timestamp_to_unix_timestamp( + "2026-10-20 09:20:27", assume_utc=True + ) + self.assertEqual(ts, expected * 10**9) + + @patch( + "parsedmarc.loki.parsed_smtp_tls_reports_to_csv_rows", + return_value=SMTP_TLS_ROWS, + ) + def test_smtp_tls_domain_label_and_begin(self, _rows): + client = _client(timestamp="report_begin") + client.save_smtp_tls_report_to_loki(SMTP_TLS) + [(labels, ts, line)] = _entries(client) + self.assertEqual(labels["report_type"], "smtp_tls") + self.assertEqual(labels["domain"], "example.org") + self.assertEqual(line["policy_domain"], "example.org") + self.assertEqual( + ts, human_timestamp_to_unix_timestamp("2026-10-20T00:00:00Z") * 10**9 + ) + + @patch( + "parsedmarc.loki.parsed_aggregate_reports_to_csv_rows", + return_value=[{"domain": "a.example", "report_type": "x"}], + ) + def test_existing_fields_are_not_overwritten_except_report_type(self, _rows): + client = _client() + client.save_aggregate_report_to_loki(AGGREGATE) + [(_, _, line)] = _entries(client) + self.assertEqual(line["domain"], "a.example") + self.assertEqual(line["report_type"], "aggregate") + + +@patch( + "parsedmarc.loki.parsed_aggregate_reports_to_csv_rows", return_value=AGGREGATE_ROWS +) +class TestLokiErrors(unittest.TestCase): + """A failed push must raise: the CLI then keeps the mailbox messages + for a retry instead of archiving reports that were never stored.""" + + def test_error_status_raises_with_status_and_body(self, _rows): + client = _client() + response = _ok_response(400) + response.text = "entry too far behind" + _session(client).post.return_value = response + with self.assertRaises(LokiError) as ctx: + client.save_aggregate_report_to_loki(AGGREGATE) + self.assertIn("400", str(ctx.exception)) + self.assertIn("entry too far behind", str(ctx.exception)) + + def test_server_error_raises(self, _rows): + client = _client() + _session(client).post.return_value = _ok_response(503) + with self.assertRaises(LokiError): + client.save_aggregate_report_to_loki(AGGREGATE) + + def test_network_error_raises(self, _rows): + client = _client() + _session(client).post.side_effect = httpx.ConnectError("refused") + with self.assertRaises(LokiError) as ctx: + client.save_aggregate_report_to_loki(AGGREGATE) + self.assertIn("refused", str(ctx.exception)) + + def test_blocked_ingestion_status_raises(self, _rows): + """Loki returns 260 (block_ingestion_status_code default) and drops + the logs when ingestion is blocked; that must not read as success.""" + client = _client() + response = _ok_response(260) + response.text = "ingestion blocked" + _session(client).post.return_value = response + with self.assertRaises(LokiError) as ctx: + client.save_aggregate_report_to_loki(AGGREGATE) + self.assertIn("260", str(ctx.exception)) + + def test_redirect_raises_and_is_not_followed(self, _rows): + """httpx turns a redirected POST into a GET, which the target would + answer 200 without storing anything; a 3xx must raise instead.""" + requested = [] + + def handler(request: httpx.Request) -> httpx.Response: + requested.append(str(request.url)) + if request.url.host == "loki.example.org": + return httpx.Response( + 301, headers={"Location": "http://elsewhere.example.org/ok"} + ) + return httpx.Response(200) + + real_client = httpx.Client + + def client_factory(**kwargs): + return real_client(transport=httpx.MockTransport(handler), **kwargs) + + with patch("parsedmarc.loki.httpx.Client", side_effect=client_factory): + client = LokiClient("http://loki.example.org:3100") + with self.assertRaises(LokiError) as ctx: + client.save_aggregate_report_to_loki(AGGREGATE) + self.assertIn("301", str(ctx.exception)) + self.assertEqual(requested, ["http://loki.example.org:3100/loki/api/v1/push"]) + + def test_success_statuses(self, _rows): + for status in (200, 204): + client = _client() + _session(client).post.return_value = _ok_response(status) + client.save_aggregate_report_to_loki(AGGREGATE) + + +class TestLokiRealReports(unittest.TestCase): + """End to end through parsedmarc's own row conversion.""" + + @unittest.skipUnless(hasattr(time, "tzset"), "needs time.tzset") + def test_report_begin_is_utc_in_a_non_utc_timezone(self): + """The aggregate row's interval_begin is naive UTC, so the pushed + timestamp must equal the XML epoch whatever the local zone + (it was 4h late under America/New_York when read as local time).""" + old_tz = os.environ.get("TZ") + + def restore(): + if old_tz is None: + os.environ.pop("TZ", None) + else: + os.environ["TZ"] = old_tz + time.tzset() + + self.addCleanup(restore) + os.environ["TZ"] = "America/New_York" + time.tzset() + path = os.path.join( + SAMPLES, "aggregate", "!example.com!1538204542!1538463818.xml" + ) + report = parsedmarc.parse_report_file(path, offline=True)["report"] + client = _client(timestamp="report_begin") + client.save_aggregate_report_to_loki(cast(AggregateReport, report)) + # of the sample is 1538413632; a one-day report is one record + # bucket, so every pushed entry starts at that epoch + stamps = {ts for _, ts, _ in _entries(client)} + self.assertEqual(stamps, {1538413632 * 10**9}) + + def test_aggregate_sample(self): + sample = os.path.join( + SAMPLES, "aggregate", "!example.com!1538204542!1538463818.xml" + ) + report = cast( + AggregateReport, + parsedmarc.parse_report_file(sample, offline=True)["report"], + ) + client = _client() + client.save_aggregate_report_to_loki(report) + [(labels, _, line)] = _entries(client) + self.assertEqual(labels["domain"], "example.com") + self.assertEqual(line["source_ip_address"], "12.20.127.122") + self.assertEqual(line["report_type"], "aggregate") + + +class TestLokiClose(unittest.TestCase): + def test_close_closes_session(self): + client = _client() + session = _session(client) + client.close() + session.close.assert_called_once() + + +if __name__ == "__main__": + unittest.main(verbosity=2)