feat: add rate limit for nvd downloads (from @param211) (#1230)

* Add rate limit for NVD downloads

* Add license information

* Fix merge conflicts

* fix: move license to top of file (per recommendation from our licensing team)

Co-authored-by: param211 <[email protected]>

fixes #1081

Author: terriko

.github/workflows/pythonapp.yml

@@ -94,7 +94,7 @@ jobs:
       matrix:
         os: [ubuntu-latest]
         python: [3.6, 3.7]
-    timeout-minutes: 10
+    timeout-minutes: 20
     env:
       ACTIONS: 1
       LONG_TESTS: 0
@@ -232,7 +232,7 @@ jobs:
   windows_tests:
     name: Windows py3.8
     runs-on: windows-latest
-    timeout-minutes: 10
+    timeout-minutes: 20
     env:
       ACTIONS: 1
       LONG_TESTS: 0

cve_bin_tool/async_utils.py

@@ -1,8 +1,32 @@
 # Copyright (C) 2021 Intel Corporation
 # SPDX-License-Identifier: GPL-3.0-or-later
 
+# This file also includes the RateLimiter function with the following license:
+#
+#    Copyright 2018 Quentin Pradet
+#
+#    Permission is hereby granted, free of charge, to any person obtaining a
+# copy of this software and associated documentation files (the "Software"), to
+# deal in the Software without restriction, including without limitation the
+# rights to use, copy, modify, merge, publish, distribute, sublicense, and/or
+# sell copies of the Software, and to permit persons to whom the Software is
+# furnished to do so, subject to the following conditions:
+
+#    The above copyright notice and this permission notice shall be included in
+# all copies or substantial portions of the Software.
+
+#    THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+# SOFTWARE."
+
 # pylint: disable=too-many-arguments
+
 """ Utility classes for the CVE Binary Tool """
+
 import asyncio
 import glob
 import gzip
@@ -12,6 +36,7 @@
 import subprocess
 import sys
 import tempfile
+import time
 from functools import partial, wraps
 
 from cve_bin_tool.util import inpath
@@ -203,6 +228,46 @@ class GzipFile(FileIO):
     _open = async_wrap(gzip.GzipFile)
 
 
+class RateLimiter:
+    """Rate limits an HTTP client that would make get() and post() calls.
+    Calls are rate-limited by host.
+    https://quentin.pradet.me/blog/how-do-you-rate-limit-calls-with-aiohttp.html
+    This class is not thread-safe.
+
+    Copyright 2018 Quentin Pradet
+    See license at top of file.
+    """
+
+    RATE = 10
+    MAX_TOKENS = 10
+
+    def __init__(self, client):
+        self.client = client
+        self.tokens = self.MAX_TOKENS
+        self.updated_at = time.monotonic()
+
+    async def get(self, *args, **kwargs):
+        await self.wait_for_token()
+        return self.client.get(*args, **kwargs)
+
+    async def wait_for_token(self):
+        while self.tokens < 1:
+            self.add_new_tokens()
+            await asyncio.sleep(0.1)
+        self.tokens -= 1
+
+    def add_new_tokens(self):
+        now = time.monotonic()
+        time_since_update = now - self.updated_at
+        new_tokens = time_since_update * self.RATE
+        if self.tokens + new_tokens >= 1:
+            self.tokens = min(self.tokens + new_tokens, self.MAX_TOKENS)
+            self.updated_at = now
+
+    async def close(self):
+        await self.client.close()
+
+
 aio_rmdir = async_wrap(shutil.rmtree)
 aio_rmfile = async_wrap(os.remove)
 aio_unpack_archive = async_wrap(shutil.unpack_archive)

cve_bin_tool/cvedb.py

@@ -20,7 +20,7 @@
 from bs4 import BeautifulSoup
 from rich.progress import track
 
-from cve_bin_tool.async_utils import FileIO, GzipFile, run_coroutine
+from cve_bin_tool.async_utils import FileIO, GzipFile, RateLimiter, run_coroutine
 from cve_bin_tool.error_handler import (
     AttemptedToWriteOutsideCachedir,
     CVEDataForCurlVersionNotInCache,
@@ -103,7 +103,7 @@ def get_db_update_date(self):
         return os.path.getmtime(self.dbpath)
 
     async def getmeta(self, session, meta_url):
-        async with session.get(meta_url) as response:
+        async with await session.get(meta_url) as response:
             response.raise_for_status()
             return (
                 meta_url.replace(".meta", ".json.gz"),
@@ -117,7 +117,7 @@ async def getmeta(self, session, meta_url):
             )
 
     async def nist_scrape(self, session):
-        async with session.get(self.feed) as response:
+        async with await session.get(self.feed) as response:
             response.raise_for_status()
             page = await response.text()
             json_meta_links = self.META_REGEX.findall(page)
@@ -161,8 +161,7 @@ async def cache_update(self, session, url, sha, chunk_size=16 * 1024):
                 self.LOGGER.debug(f"Correct SHA for {filename}")
                 return
         self.LOGGER.debug(f"Updating CVE cache for {filename}")
-
-        async with session.get(url) as response:
+        async with await session.get(url) as response:
             # Raise better error message on ratelimit by NVD
             if response.status == 403:
                 with ErrorHandler(mode=self.error_mode, logger=self.LOGGER):
@@ -187,7 +186,7 @@ async def cache_update(self, session, url, sha, chunk_size=16 * 1024):
     @staticmethod
     async def get_curl_versions(session):
         regex = re.compile(r"vuln-(\d+.\d+.\d+)\.html")
-        async with session.get(
+        async with await session.get(
             "https://curl.haxx.se/docs/vulnerabilities.html"
         ) as response:
             response.raise_for_status()
@@ -196,7 +195,7 @@ async def get_curl_versions(session):
         return [match.group(1) for match in matches]
 
     async def download_curl_version(self, session, version):
-        async with session.get(
+        async with await session.get(
             f"https://curl.haxx.se/docs/vuln-{version}.html"
         ) as response:
             response.raise_for_status()
@@ -233,13 +232,16 @@ async def refresh(self):
             check_latest_version()
         if not self.session:
             connector = aiohttp.TCPConnector(limit_per_host=19)
-            self.session = aiohttp.ClientSession(connector=connector, trust_env=True)
+            self.session = RateLimiter(
+                aiohttp.ClientSession(connector=connector, trust_env=True)
+            )
         self.LOGGER.info("Downloading CVE data...")
         nvd_metadata, curl_metadata = await asyncio.gather(
-            self.nist_scrape(self.session), self.get_curl_versions(self.session)
+            asyncio.ensure_future(self.nist_scrape(self.session)),
+            asyncio.ensure_future(self.get_curl_versions(self.session)),
         )
         tasks = [
-            self.cache_update(self.session, url, meta["sha256"])
+            asyncio.ensure_future(self.cache_update(self.session, url, meta["sha256"]))
             for url, meta in nvd_metadata.items()
             if meta is not None
         ]
@@ -251,7 +253,9 @@ async def refresh(self):
         tasks.append(
             asyncio.gather(
                 *(
-                    self.download_curl_version(self.session, version)
+                    asyncio.ensure_future(
+                        self.download_curl_version(self.session, version)
+                    )
                     for version in curl_metadata
                 )
             )