|
30 | 30 |
|
31 | 31 | import lib.utils.api as api |
32 | 32 | from lib.core.data import conf |
| 33 | +from lib.core.data import kb |
33 | 34 | from lib.core.convert import encodeBase64 |
34 | 35 | from lib.core.enums import CONTENT_STATUS, CONTENT_TYPE |
35 | 36 | from thirdparty.bottle.bottle import default_app |
@@ -308,6 +309,62 @@ def test_logrecorder_emit_stores_log(self): |
308 | 309 | self.assertEqual(rows[0][2], "hello world") |
309 | 310 |
|
310 | 311 |
|
| 312 | +class TestStdDbOutMultiTarget(unittest.TestCase): |
| 313 | + """ |
| 314 | + Regression coverage for issue #2503: a multi-target scan (crawl or '-m' bulk |
| 315 | + file) inside a single REST API task runs every target through one StdDbOut |
| 316 | + instance. StdDbOut.write() keys each target's data rows by kb.reportTaskId |
| 317 | + ('<taskid>-<n>'), so a later target's COMPLETE row cannot delete an earlier |
| 318 | + target's rows of the same content_type in _storeData() - mirroring the |
| 319 | + --report-json fix for #6123 (see test_report.py's TestMultiTargetTaskId). |
| 320 | + """ |
| 321 | + |
| 322 | + def setUp(self): |
| 323 | + self.db = api.Database(":memory:") |
| 324 | + self.db.connect("client") |
| 325 | + self.db.init() |
| 326 | + self._saved_stdout = sys.stdout |
| 327 | + self._saved_cursor = conf.get("databaseCursor") |
| 328 | + self._saved_partRun = kb.get("partRun") |
| 329 | + self._had_reportTaskId = "reportTaskId" in kb |
| 330 | + self._saved_reportTaskId = kb.get("reportTaskId") |
| 331 | + conf.databaseCursor = self.db |
| 332 | + kb.partRun = None |
| 333 | + |
| 334 | + def tearDown(self): |
| 335 | + sys.stdout = self._saved_stdout |
| 336 | + conf.databaseCursor = self._saved_cursor |
| 337 | + kb.partRun = self._saved_partRun |
| 338 | + if self._had_reportTaskId: |
| 339 | + kb.reportTaskId = self._saved_reportTaskId |
| 340 | + elif "reportTaskId" in kb: |
| 341 | + del kb.reportTaskId |
| 342 | + self.db.disconnect() |
| 343 | + |
| 344 | + def test_second_target_does_not_overwrite_first(self): |
| 345 | + std = api.StdDbOut("t1", messagetype="stdout") |
| 346 | + sys.stdout = self._saved_stdout # StdDbOut hijacked it; call write() directly |
| 347 | + |
| 348 | + kb.reportTaskId = 1 |
| 349 | + std.write({"url": "http://host1/?id=1"}, status=CONTENT_STATUS.COMPLETE, content_type=CONTENT_TYPE.TARGET) |
| 350 | + kb.reportTaskId = 2 |
| 351 | + std.write({"url": "http://host2/?id=1"}, status=CONTENT_STATUS.COMPLETE, content_type=CONTENT_TYPE.TARGET) |
| 352 | + |
| 353 | + rows = self.db.execute("SELECT taskid FROM data ORDER BY id") |
| 354 | + self.assertEqual([row[0] for row in rows], ["t1-1", "t1-2"]) # neither target's row was deleted |
| 355 | + |
| 356 | + def test_single_target_without_report_task_id_keeps_plain_taskid(self): |
| 357 | + if "reportTaskId" in kb: |
| 358 | + del kb.reportTaskId |
| 359 | + std = api.StdDbOut("t1", messagetype="stdout") |
| 360 | + sys.stdout = self._saved_stdout |
| 361 | + |
| 362 | + std.write({"url": "http://host1/?id=1"}, status=CONTENT_STATUS.COMPLETE, content_type=CONTENT_TYPE.TARGET) |
| 363 | + |
| 364 | + rows = self.db.execute("SELECT taskid FROM data") |
| 365 | + self.assertEqual([row[0] for row in rows], ["t1"]) |
| 366 | + |
| 367 | + |
311 | 368 | # --------------------------------------------------------------------------- |
312 | 369 | # HTTP routes (WSGI test client) |
313 | 370 | # --------------------------------------------------------------------------- |
@@ -485,6 +542,26 @@ def test_scan_kill_invalid_when_not_running(self): |
485 | 542 | self.assertFalse(parsed["success"]) |
486 | 543 |
|
487 | 544 |
|
| 545 | +class TestScanDataMultiTarget(_ApiServerCase): |
| 546 | + """scan_data() folds the per-target '<taskid>-<n>' rows (see StdDbOut.write) back into the task's flat result, in target order.""" |
| 547 | + |
| 548 | + def test_data_merges_per_target_rows(self): |
| 549 | + taskid = self._new_task() |
| 550 | + api._storeData(self.db, "%s-2" % taskid, {"url": "http://host2/?id=1"}, CONTENT_STATUS.COMPLETE, CONTENT_TYPE.TARGET) |
| 551 | + api._storeData(self.db, "%s-1" % taskid, {"url": "http://host1/?id=1"}, CONTENT_STATUS.COMPLETE, CONTENT_TYPE.TARGET) |
| 552 | + code, parsed, _ = _wsgi_call("GET", "/scan/%s/data" % taskid) |
| 553 | + self.assertEqual(code, 200) |
| 554 | + self.assertTrue(parsed["success"]) |
| 555 | + self.assertEqual([d["value"]["url"] for d in parsed["data"]], ["http://host1/?id=1", "http://host2/?id=1"]) |
| 556 | + |
| 557 | + def test_data_ignores_other_tasks_rows(self): |
| 558 | + taskid = self._new_task() |
| 559 | + api._storeData(self.db, "othertask-1", {"url": "http://other/?id=1"}, CONTENT_STATUS.COMPLETE, CONTENT_TYPE.TARGET) |
| 560 | + code, parsed, _ = _wsgi_call("GET", "/scan/%s/data" % taskid) |
| 561 | + self.assertTrue(parsed["success"]) |
| 562 | + self.assertEqual(parsed["data"], []) |
| 563 | + |
| 564 | + |
488 | 565 | class TestScanStart(_ApiServerCase): |
489 | 566 | """scan_start, with the subprocess-spawning seam (engine_start) monkeypatched.""" |
490 | 567 |
|
|
0 commit comments