Merge pull request #2778 from blacklanternsecurity/download-api-key-cycle

Add api_download() helper

Author: TheTechromancer

bbot/modules/base.py

@@ -4,8 +4,8 @@
 from sys import exc_info
 from contextlib import suppress
 
-from ..errors import ValidationError
 from ..core.helpers.misc import get_size  # noqa
+from ..errors import ValidationError, WebError
 from ..core.helpers.async_helpers import TaskCounter, ShuffleQueue
 
 
@@ -1249,6 +1249,24 @@ async def api_request(self, *args, **kwargs):
 
         return r
 
+    async def api_download(self, url, **kwargs):
+        """
+        A wrapper around the `download()` web helper that incorporates API key cycling.
+        """
+        error = None
+        raise_error = kwargs.pop("raise_error", False)
+        for _ in range(self.api_retries):
+            new_url, kwargs = self.prepare_api_request(url, kwargs)
+            if "raise_error" not in kwargs:
+                kwargs["raise_error"] = True
+            try:
+                return await self.helpers.download(new_url, **kwargs)
+            except WebError as e:
+                error = e
+                self.cycle_api_key()
+        if raise_error:
+            raise error
+
     def _get_retry_after(self, r):
         # try to get retry_after from headers first
         headers = getattr(r, "headers", {})

bbot/test/test_step_1/test_web.py

@@ -498,3 +498,49 @@ def echo_cookies_handler(request):
     assert r1 is not None, "Request to self-signed SSL server went through even with ssl_verify=True"
     assert "bar" in r1.text
     await scan1._cleanup()
+
+
+@pytest.mark.asyncio
+async def test_api_download_api_key_cycle(bbot_scanner, bbot_httpserver):
+    from werkzeug.wrappers import Response
+    from bbot.modules.base import BaseModule
+
+    endpoint = "/api_download_cycle_one_test"
+    url = bbot_httpserver.url_for(endpoint)
+
+    seen_auth = []
+    n_request = 0
+
+    # First key should trigger 500, second key should succeed with 200
+    def handler(request):
+        nonlocal n_request
+        n_request += 1
+        auth = request.headers.get("Authorization", "")
+        seen_auth.append(auth)
+        if auth == "Bearer k1":
+            if n_request == 1:
+                return Response("ok_k1", status=200)
+            return Response("fail_k1", status=500)
+        elif auth == "Bearer k2":
+            return Response("ok_k2", status=200)
+        return Response("unexpected_key", status=400)
+
+    bbot_httpserver.expect_request(uri=endpoint).respond_with_handler(handler)
+
+    scan = bbot_scanner("127.0.0.1")
+    module = BaseModule(scan)
+    module.api_key = ["k1", "k2"]
+
+    filename = await module.api_download(url)
+    assert filename is not None
+    with open(filename) as f:
+        assert f.read() == "ok_k1"
+
+    assert seen_auth == ["Bearer k1"]
+
+    filename = await module.api_download(url)
+
+    # verify the requests occurred in expected order with expected API keys
+    assert seen_auth == ["Bearer k1", "Bearer k1", "Bearer k2"]
+
+    await scan._cleanup()