|
2 | 2 | import unittest |
3 | 3 | import subprocess |
4 | 4 | import time |
| 5 | +import json |
| 6 | +import base64 |
5 | 7 | import requests |
6 | 8 | import atexit |
7 | 9 | from contextlib import contextmanager |
|
12 | 14 | Run it like so: |
13 | 15 |
|
14 | 16 | # all tests |
15 | | -python test_helm_deployment.py |
| 17 | +uv run python test_helm_deployment.py |
16 | 18 |
|
17 | 19 | # specific test |
18 | | -python -m unittest test_helm_deployment.TestHelmDeployment.test_swagger_ui |
| 20 | +uv run python -m unittest test_helm_deployment.TestHelmDeployment.test_swagger_ui |
19 | 21 |
|
20 | 22 | To bring up the test environment manually, run: |
21 | 23 |
|
@@ -264,6 +266,91 @@ def test_swagger_ui(self): |
264 | 266 | self.assertEqual(response.status_code, 200, "Swagger UI should return 200") |
265 | 267 | self.assertIn("swagger-ui", response.text.lower(), "Response should contain swagger-ui") |
266 | 268 |
|
| 269 | + def test_ingest_and_query_assets(self): |
| 270 | + """Test ingesting events from evilcorp.json.gz and querying assets via bbctl""" |
| 271 | + # Get the server pod name |
| 272 | + result = self.kubectl( |
| 273 | + "get", |
| 274 | + "pods", |
| 275 | + "-l", |
| 276 | + "app=bbot-server", |
| 277 | + "-o", |
| 278 | + "jsonpath={.items[0].metadata.name}", |
| 279 | + capture_output=True, |
| 280 | + ) |
| 281 | + pod_name = result.stdout.strip() |
| 282 | + self.assertTrue(pod_name, "Should find a server pod") |
| 283 | + print(f"Server pod: {pod_name}") |
| 284 | + |
| 285 | + # Get the API key from the kubernetes secret |
| 286 | + result = self.kubectl( |
| 287 | + "get", |
| 288 | + "secret", |
| 289 | + f"{self.release_name}-api-key", |
| 290 | + "-o", |
| 291 | + "jsonpath={.data.api-key}", |
| 292 | + capture_output=True, |
| 293 | + ) |
| 294 | + api_key = base64.b64decode(result.stdout).decode() |
| 295 | + self.assertTrue(api_key, "Should retrieve API key from secret") |
| 296 | + |
| 297 | + # Write a bbctl config file inside the pod |
| 298 | + server_url = f"http://{self.release_name}-server:8807/v1/" |
| 299 | + config_content = f'url: "{server_url}"\napi_keys:\n - "{api_key}"' |
| 300 | + self.kubectl( |
| 301 | + "exec", |
| 302 | + pod_name, |
| 303 | + "--", |
| 304 | + "sh", |
| 305 | + "-c", |
| 306 | + f"printf '%s\\n' '{config_content}' > /tmp/bbctl.yaml", |
| 307 | + ) |
| 308 | + bbctl = "bbctl --no-color --config /tmp/bbctl.yaml" |
| 309 | + |
| 310 | + # Copy test data into the pod |
| 311 | + self.kubectl("cp", "tests/evilcorp.json.gz", f"{pod_name}:/tmp/evilcorp.json.gz") |
| 312 | + |
| 313 | + # Ingest events using bbctl |
| 314 | + print("Ingesting events...") |
| 315 | + result = self.kubectl( |
| 316 | + "exec", |
| 317 | + pod_name, |
| 318 | + "--", |
| 319 | + "sh", |
| 320 | + "-c", |
| 321 | + f"gunzip -c /tmp/evilcorp.json.gz | {bbctl} event ingest", |
| 322 | + capture_output=True, |
| 323 | + timeout=120, |
| 324 | + ) |
| 325 | + print(f"Ingest stdout: {result.stdout}") |
| 326 | + print(f"Ingest stderr: {result.stderr}") |
| 327 | + self.assertEqual(result.returncode, 0, f"Event ingest failed: {result.stderr}") |
| 328 | + |
| 329 | + # Wait for the watchdog to process events into assets |
| 330 | + print("Waiting for watchdog to process events into assets...") |
| 331 | + assets = [] |
| 332 | + for i in range(30): |
| 333 | + result = self.kubectl( |
| 334 | + "exec", |
| 335 | + pod_name, |
| 336 | + "--", |
| 337 | + "sh", |
| 338 | + "-c", |
| 339 | + f"{bbctl} asset list --json", |
| 340 | + capture_output=True, |
| 341 | + check=False, |
| 342 | + timeout=30, |
| 343 | + ) |
| 344 | + if result.returncode == 0 and result.stdout.strip(): |
| 345 | + assets = [json.loads(line) for line in result.stdout.strip().splitlines() if line.strip()] |
| 346 | + if assets: |
| 347 | + break |
| 348 | + print(f" Attempt {i + 1}/30: {len(assets)} assets so far...") |
| 349 | + time.sleep(2) |
| 350 | + |
| 351 | + print(f"Found {len(assets)} assets") |
| 352 | + self.assertGreater(len(assets), 0, "Should have ingested some assets from evilcorp.json.gz") |
| 353 | + |
267 | 354 | def tearDown(self): |
268 | 355 | """If a test fails, dump cluster debug info""" |
269 | 356 | failed = False |
|
0 commit comments