Preventing exceptions, handling too many requests.

Author: nicholasks

safety/cli.py

@@ -7,7 +7,7 @@
 from safety.formatter import report, license_report
 import itertools
 from safety.util import read_requirements, read_vulnerabilities, get_proxy_dict, get_packages_licenses
-from safety.errors import DatabaseFetchError, DatabaseFileNotFoundError, InvalidKeyError
+from safety.errors import DatabaseFetchError, DatabaseFileNotFoundError, InvalidKeyError, TooManyRequestsError
 
 try:
     from json.decoder import JSONDecodeError
@@ -151,7 +151,26 @@ def license(key, db, cache, files, proxyprotocol, proxyhost, proxyport):
         ]  
 
     proxy_dictionary = get_proxy_dict(proxyprotocol, proxyhost, proxyport)
-    licenses_db = safety.get_licenses(key, db, cache, proxy_dictionary)
+    try:
+        licenses_db = safety.get_licenses(key, db, cache, proxy_dictionary)
+    except InvalidKeyError:
+        click.secho("Your API Key '{key}' is invalid. See {link}".format(
+            key=key, link='https://goo.gl/O7Y1rS'),
+            fg="red",
+            file=sys.stderr)
+        sys.exit(-1)
+    except DatabaseFileNotFoundError:
+        click.secho("Unable to load licenses database from {db}".format(db=db), fg="red", file=sys.stderr)
+        sys.exit(-1)
+    except TooManyRequestsError:
+        click.secho("Unable to load licenses database (Too many requests, please wait before another request)",
+            fg="red",
+            file=sys.stderr
+        )
+        sys.exit(-1)
+    except DatabaseFetchError:
+        click.secho("Unable to load licenses database", fg="red", file=sys.stderr)
+        sys.exit(-1)
     filtered_packages_licenses = get_packages_licenses(packages, licenses_db)
     output_report = license_report(packages=packages, licenses=filtered_packages_licenses)
     click.secho(output_report, nl=True)

safety/errors.py

@@ -8,3 +8,7 @@ class DatabaseFileNotFoundError(DatabaseFetchError):
 
 class InvalidKeyError(DatabaseFetchError):
     pass
+
+
+class TooManyRequestsError(DatabaseFetchError):
+    pass

safety/safety.py

@@ -11,7 +11,7 @@
 from .constants import (API_MIRRORS, CACHE_FILE, CACHE_LICENSES_VALID_SECONDS,
                         CACHE_VALID_SECONDS, OPEN_MIRRORS, REQUEST_TIMEOUT)
 from .errors import (DatabaseFetchError, DatabaseFileNotFoundError,
-                     InvalidKeyError)
+                     InvalidKeyError, TooManyRequestsError)
 from .util import RequirementFile
 
 
@@ -95,6 +95,8 @@ def fetch_database_url(mirror, db_name, key, cached, proxy):
         return data
     elif r.status_code == 403:
         raise InvalidKeyError()
+    elif r.status_code == 429:
+        raise TooManyRequestsError()
 
 
 def fetch_database_file(path, db_name):