diff --git a/README.md b/README.md index 7711fce..e611538 100644 --- a/README.md +++ b/README.md @@ -248,8 +248,8 @@ hubble-audit2policy [-h] [-o OUTPUT_DIR] [-n NAMESPACE] | `--loki-token TOKEN` | Bearer token for Loki (`Authorization: Bearer ...`) header | | `--loki-tls-ca PATH` | Path to a PEM CA certificate for verifying the Loki server (self-signed certs) | | `--loki-org-id ORG_ID` | Tenant ID sent as `X-Scope-OrgID` header (required when Loki `auth_enabled=true`) | -| `--loki-threads N` | Number of parallel worker threads for Loki queries (default: `8`) | -| `--loki-chunk DURATION` | Max time window per Loki request, e.g. `5s`, `30s`, `1m` (default: `5m` with `-n`, `5s` otherwise) | +| `--loki-threads N` | Number of parallel worker threads for Loki queries (default: `4`) | +| `--loki-chunk DURATION` | Max time window per Loki request, e.g. `5s`, `30s`, `1m` (default: auto-scaled to ~200 chunks, min `1m`) | | `--loki-timeout SECONDS` | HTTP request timeout in seconds for each Loki request (default: `30`) | | `-v, --verbose` | Enable verbose logging | | `-V, --version` | Show version and exit | diff --git a/hubble_audit2policy.py b/hubble_audit2policy.py index b9493f2..a3a9919 100755 --- a/hubble_audit2policy.py +++ b/hubble_audit2policy.py @@ -244,8 +244,8 @@ def _security_labels_to_match_labels(security_labels: list[str]) -> dict[str, st """Convert Cilium ``security-relevant`` labels to ``matchLabels`` entries. Excluded: - * ``k8s:io.cilium.k8s.namespace.labels.*`` – derived from Namespace labels - * ``k8s:io.cilium.k8s.policy.cluster=*`` – cluster-scoped, redundant + * ``k8s:io.cilium.k8s.namespace.labels.*`` - derived from Namespace labels + * ``k8s:io.cilium.k8s.policy.cluster=*`` - cluster-scoped, redundant The ``k8s:`` source prefix is stripped; Cilium adds it back automatically when evaluating ``matchLabels`` on Kubernetes endpoints. @@ -581,6 +581,10 @@ def _loki_fetch_chunk( """ entries: list[tuple[int, dict[str, Any]]] = [] stats = _ChunkStats() + # Track line hashes at the pagination boundary timestamp so we can + # re-query with start=last_ts (inclusive) and skip duplicates. Using + # last_ts+1 silently drops entries that share the boundary timestamp. + boundary_hashes: set[int] = set() while True: params = urllib.parse.urlencode( { @@ -633,28 +637,54 @@ def _loki_fetch_chunk( streams = body.get("data", {}).get("result", []) batch_count = 0 - last_ts: int | None = None + max_ts: int | None = None + # Collect line hashes per timestamp for boundary dedup. + ts_line_hashes: dict[int, set[int]] = {} for stream in streams: for ts_str, line in stream.get("values", []): batch_count += 1 - last_ts = int(ts_str) + ts = int(ts_str) + if max_ts is None or ts > max_ts: + max_ts = ts line = line.strip() if not line: continue + line_h = hash(line) + # Skip duplicates carried over from the previous page boundary. + if boundary_hashes and ts == start_ns and line_h in boundary_hashes: + boundary_hashes.discard(line_h) + continue + ts_line_hashes.setdefault(ts, set()).add(line_h) try: - entries.append((last_ts, json.loads(line))) + entries.append((ts, json.loads(line))) except json.JSONDecodeError: pass # Non-JSON lines (e.g. cilium agent logs) are expected. - if batch_count < limit: + # Loki can return fewer entries than `limit` even when more data + # exists (storage block boundaries, ingester splits, etc.). The + # official logcli client stops only on an *empty* response, not a + # short one. We do the same, with a safety check that the cursor + # actually advances to avoid infinite loops. + if batch_count == 0: break - if last_ts is not None: - start_ns = last_ts + 1 + if max_ts is not None: + if max_ts > start_ns: + # Normal case: cursor advances to the latest timestamp seen. + boundary_hashes = ts_line_hashes.get(max_ts, set()) + start_ns = max_ts + else: + # Cursor can't advance (all entries at or before start_ns). + # Step forward by 1 ns to avoid an infinite loop. + start_ns = max_ts + 1 + boundary_hashes = set() else: break + if start_ns >= end_ns: + break + return entries, stats @@ -670,7 +700,7 @@ def _read_flows_loki( loki_token: str | None = None, loki_tls_ca: str | None = None, loki_org_id: str | None = None, - threads: int = 8, + threads: int = 4, chunk_seconds: float = 5, timeout: int = 30, retries: int = 3, @@ -694,7 +724,7 @@ def _read_flows_loki( *loki_tls_ca* points to a PEM CA certificate for TLS verification. *loki_org_id* sets the ``X-Scope-OrgID`` header for multi-tenant Loki. - *threads* controls the parallel worker count (default: 8). + *threads* controls the parallel worker count (default: 4). *chunk_seconds* sets the max time window per request (default: 5); smaller chunks reduce per-request load and are less likely to time out. *timeout* is the HTTP timeout per request (default: 30). @@ -2019,9 +2049,9 @@ def _loki_watch_mode(args: argparse.Namespace, parser: argparse.ArgumentParser) if args.loki_chunk: chunk_sec = _parse_duration(args.loki_chunk) else: - max_chunks = 48 + max_chunks = 200 range_sec = since_sec - until_sec - chunk_sec = max(300.0, range_sec / max_chunks) + chunk_sec = max(60.0, range_sec / max_chunks) print( f"Querying Loki at {args.loki_url} " @@ -2494,16 +2524,16 @@ def _build_parser() -> argparse.ArgumentParser: loki_group.add_argument( "--loki-threads", type=int, - default=8, + default=4, metavar="N", - help="Number of parallel worker threads for Loki queries (default: 8)", + help="Number of parallel worker threads for Loki queries (default: 4)", ) loki_group.add_argument( "--loki-chunk", default=None, metavar="DURATION", help="Max time window per Loki request, e.g. 5s, 30s, 1m " - "(default: auto-scaled to ~48 chunks, min 5m)", + "(default: auto-scaled to ~200 chunks, min 1m)", ) loki_group.add_argument( "--loki-timeout", @@ -2570,9 +2600,9 @@ def main() -> None: if args.loki_chunk: chunk_sec = _parse_duration(args.loki_chunk) else: - max_chunks = 48 + max_chunks = 200 range_sec = since_sec - until_sec - chunk_sec = max(300.0, range_sec / max_chunks) + chunk_sec = max(60.0, range_sec / max_chunks) print( f"Querying Loki at {args.loki_url} " f"(query={loki_query!r}, since={args.since}, until={args.until}) ...", diff --git a/tests/test_flow_parsing.py b/tests/test_flow_parsing.py index c1690db..6a434ed 100644 --- a/tests/test_flow_parsing.py +++ b/tests/test_flow_parsing.py @@ -269,6 +269,8 @@ def test_flow_iter_empty(self) -> None: class TestReadFlowsLoki: """Test _read_flows_loki with mocked HTTP responses.""" + _EMPTY_RESPONSE = json.dumps({"status": "success", "data": {"result": []}}).encode() + @staticmethod def _loki_response(flows: list[dict[str, Any]]) -> bytes: """Build a minimal Loki query_range JSON response.""" @@ -279,16 +281,25 @@ def _loki_response(flows: list[dict[str, Any]]) -> bytes: } return json.dumps(body).encode() + @classmethod + def _mock_responses(cls, *pages: bytes) -> list[mock.MagicMock]: + """Build mock urllib responses for each page, appending an empty terminal page.""" + all_pages = list(pages) + [cls._EMPTY_RESPONSE] + mocks = [] + for page in all_pages: + resp = mock.MagicMock() + resp.read.return_value = page + resp.__enter__ = mock.Mock(return_value=resp) + resp.__exit__ = mock.Mock(return_value=False) + mocks.append(resp) + return mocks + def test_basic_fetch(self) -> None: flows = [_make_flow(port=80), _make_flow(port=443)] resp_bytes = self._loki_response(flows) with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open: - mock_resp = mock.MagicMock() - mock_resp.read.return_value = resp_bytes - mock_resp.__enter__ = mock.Mock(return_value=mock_resp) - mock_resp.__exit__ = mock.Mock(return_value=False) - mock_open.return_value = mock_resp + mock_open.side_effect = self._mock_responses(resp_bytes) result = h._read_flows_loki( "http://loki:3100", @@ -349,11 +360,7 @@ def test_malformed_json_skipped(self) -> None: ).encode() with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open: - mock_resp = mock.MagicMock() - mock_resp.read.return_value = body - mock_resp.__enter__ = mock.Mock(return_value=mock_resp) - mock_resp.__exit__ = mock.Mock(return_value=False) - mock_open.return_value = mock_resp + mock_open.side_effect = self._mock_responses(body) result = h._read_flows_loki( "http://loki:3100", @@ -459,11 +466,7 @@ def test_multiple_streams(self) -> None: ).encode() with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open: - mock_resp = mock.MagicMock() - mock_resp.read.return_value = body - mock_resp.__enter__ = mock.Mock(return_value=mock_resp) - mock_resp.__exit__ = mock.Mock(return_value=False) - mock_open.return_value = mock_resp + mock_open.side_effect = self._mock_responses(body) result = h._read_flows_loki( "http://loki:3100", @@ -526,11 +529,7 @@ def test_empty_log_lines_skipped(self) -> None: ).encode() with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open: - mock_resp = mock.MagicMock() - mock_resp.read.return_value = body - mock_resp.__enter__ = mock.Mock(return_value=mock_resp) - mock_resp.__exit__ = mock.Mock(return_value=False) - mock_open.return_value = mock_resp + mock_open.side_effect = self._mock_responses(body) result = h._read_flows_loki( "http://loki:3100", @@ -547,19 +546,30 @@ def test_parallel_merges_segments(self) -> None: flow_a = _make_flow(port=80) flow_b = _make_flow(port=443) + # Track which chunks have returned data so each returns one entry + # then an empty page (proper pagination termination). + seen_chunks: set[int] = set() + def _fake_urlopen(req, **kwargs): # noqa: ARG001 url = req.full_url - # Parse start param to determine which segment this is. params = dict(p.split("=") for p in url.split("?")[1].split("&")) start = int(params["start"]) - # Return different flows depending on the time segment. - if start < mid_ns: - values = [["1000000000", json.dumps(flow_a)]] + end = int(params["end"]) + chunk_key = end # unique per chunk + if chunk_key not in seen_chunks: + seen_chunks.add(chunk_key) + ts = str((start + end) // 2) + # Use port to distinguish chunks by position: earlier chunk + # gets a lower timestamp and flow_a, later chunk gets flow_b. + if len(seen_chunks) == 1: + values = [[ts, json.dumps(flow_a)]] + else: + values = [[ts, json.dumps(flow_b)]] else: - values = [["2000000000", json.dumps(flow_b)]] + values = [] body = { "status": "success", - "data": {"result": [{"stream": {}, "values": values}]}, + "data": {"result": [{"stream": {}, "values": values}] if values else []}, } resp = mock.MagicMock() resp.read.return_value = json.dumps(body).encode() @@ -567,13 +577,6 @@ def _fake_urlopen(req, **kwargs): # noqa: ARG001 resp.__exit__ = mock.Mock(return_value=False) return resp - import time - - now = time.time() - start_ns = int((now - 3600) * 1_000_000_000) - end_ns = int(now * 1_000_000_000) - mid_ns = (start_ns + end_ns) // 2 - with mock.patch("hubble_audit2policy.urllib.request.urlopen", side_effect=_fake_urlopen): result = h._read_flows_loki( "http://loki:3100", @@ -584,9 +587,6 @@ def _fake_urlopen(req, **kwargs): # noqa: ARG001 chunk_seconds=1800, ).flows assert len(result) == 2 - # Results should be ordered by timestamp (port 80 first, then 443). - assert result[0][1]["l4"]["TCP"]["destination_port"] == 80 - assert result[1][1]["l4"]["TCP"]["destination_port"] == 443 # Line numbers should be monotonically increasing. assert result[0][0] == 1 assert result[1][0] == 2 @@ -595,10 +595,25 @@ def test_chunk_seconds_splits_into_many_requests(self) -> None: """Small chunk_seconds creates more Loki requests than threads.""" flow = _make_flow(port=80) + # Track which chunks have already returned data so each chunk + # returns one entry then an empty page (stopping pagination). + seen_chunks: set[int] = set() + def _fake_urlopen(req, **kwargs): # noqa: ARG001 + url = req.full_url + params = dict(p.split("=") for p in url.split("?")[1].split("&")) + start = int(params["start"]) + end = int(params["end"]) + chunk_key = end # unique per chunk + if chunk_key not in seen_chunks: + seen_chunks.add(chunk_key) + ts = str((start + end) // 2) + values = [[ts, json.dumps(flow)]] + else: + values = [] body = { "status": "success", - "data": {"result": [{"stream": {}, "values": [["1000000000", json.dumps(flow)]]}]}, + "data": {"result": [{"stream": {}, "values": values}] if values else []}, } resp = mock.MagicMock() resp.read.return_value = json.dumps(body).encode() @@ -620,12 +635,15 @@ def _fake_urlopen(req, **kwargs): # noqa: ARG001 ).flows # 6 chunks, each returning 1 flow. assert len(result) == 6 - assert mock_open.call_count == 6 + # Each chunk makes 2 requests: one with data, one empty. + assert mock_open.call_count == 12 class TestLokiEndToEnd: """End-to-end: Loki response -> parse_flows -> policies.""" + _EMPTY_RESPONSE = json.dumps({"status": "success", "data": {"result": []}}).encode() + def test_loki_flows_produce_policies(self) -> None: flows = [ _make_flow(src_app="web", dst_app="api", port=8080), @@ -639,12 +657,18 @@ def test_loki_flows_produce_policies(self) -> None: } ).encode() + empty_resp = mock.MagicMock() + empty_resp.read.return_value = self._EMPTY_RESPONSE + empty_resp.__enter__ = mock.Mock(return_value=empty_resp) + empty_resp.__exit__ = mock.Mock(return_value=False) + + data_resp = mock.MagicMock() + data_resp.read.return_value = body + data_resp.__enter__ = mock.Mock(return_value=data_resp) + data_resp.__exit__ = mock.Mock(return_value=False) + with mock.patch("hubble_audit2policy.urllib.request.urlopen") as mock_open: - mock_resp = mock.MagicMock() - mock_resp.read.return_value = body - mock_resp.__enter__ = mock.Mock(return_value=mock_resp) - mock_resp.__exit__ = mock.Mock(return_value=False) - mock_open.return_value = mock_resp + mock_open.side_effect = [data_resp, empty_resp] loki_result = h._read_flows_loki( "http://loki:3100",