refactor the aura update to work with the new dataset release process

update docs with env config options related to dataset
add integration status to the `aura info`

Author: RootLUG

aura/config.py

@@ -22,6 +22,8 @@
 except ImportError:
     import json
 
+from .exceptions import InvalidConfiguration
+
 
 CFG: Optional[dict] = None
 CFG_PATH = None
@@ -168,7 +170,7 @@ def get_file_location(location: str, base_path: Optional[str]=None) -> str:
             return str(pth)
 
     # TODO: use custom exception here so we can log as fatal and sys.exit(1)
-    raise ValueError(f"Can't find configuration file `{location}` using base path `{base_path}`")
+    raise InvalidConfiguration(f"Can't find configuration file `{location}` using base path `{base_path}`")
 
 
 def get_file_content(location: str, base_path: Optional[str]=None) -> str:
@@ -242,6 +244,11 @@ def get_pypi_stats_path() -> Path:
     return Path(get_file_location(pth, CFG_PATH))
 
 
+def get_reverse_dependencies_path() -> Path:
+    pth = os.environ.get("AURA_REVERSE_DEPENDENCIES", None) or CFG["aura"]["reverse_dependencies"]
+    return Path(get_file_location(pth, CFG_PATH))
+
+
 def iter_pypi_stats() -> Generator[dict, None, None]:
     pth = get_pypi_stats_path()
     with pth.open() as fd:

aura/data/aura_config.yaml

@@ -20,6 +20,8 @@ aura: &aura_config
   # This file is needed for typosquatting detections
   pypi_stats: &pypi_stats pypi_stats.json
 
+  reverse_dependencies: &reverse_dependencies reverse_dependencies.json
+
   # Threshold for package download after which the package is considered not legitimate
   pypi_download_threshold: 10000
 

aura/exceptions.py

@@ -25,6 +25,10 @@ class NoSuchRepository(InvalidLocation):
     pass
 
 
+class InvalidConfiguration(AuraException, ValueError):
+    pass
+
+
 class ASTNodeRewrite(AuraException):
     pass
 

aura/info.py

@@ -13,6 +13,7 @@
     jsonschema = None
 
 from . import __version__ as version
+from .exceptions import InvalidConfiguration
 from .uri_handlers.base import URIHandler
 from . import plugins
 from . import config
@@ -27,16 +28,35 @@ def get_analyzer_description(analyzer) -> str:
 
 
 def check_pypi_stats() -> dict:
-    if config.get_pypi_stats_path():  # Put into try except
-        return {
-            "enabled": True,
-            "description": "PyPI typosquatting protection enabled"
-        }
-    else:
-        return {
-            "enabled": False,
-            "description": "PyPI download stats not found, typosquatting protection is disabled. Run `aura fetch-pypi-stats` to download"
-        }
+    try:
+        if config.get_pypi_stats_path():  # Put into try except
+            return {
+                "enabled": True,
+                "description": "PyPI typosquatting protection enabled"
+            }
+    except InvalidConfiguration:
+        pass
+
+    return {
+        "enabled": False,
+        "description": "PyPI download stats not found, typosquatting protection is disabled. Run `aura update` to download"
+    }
+
+
+def check_reverse_dependencies() -> dict:
+    try:
+        if config.get_reverse_dependencies_path():
+            return {
+                "enabled": True,
+                "description": "Reverse dependencies dataset present. Package scoring feature is fully enabled"
+            }
+    except InvalidConfiguration:
+        pass
+
+    return {
+        "enabled": False,
+        "description": "Reverse dependencies dataset not found, package scoring may be affected. Run `aura update` to download"
+    }
 
 
 def check_git() -> dict:
@@ -90,7 +110,6 @@ def gather_aura_information() -> dict:
         }
 
     for k, v in analyzers["disabled"]:
-        #doc = get_analyzer_description(v)
         info["analyzers"][k] = {
             "enabled": False,
             "description": v
@@ -105,6 +124,7 @@ def gather_aura_information() -> dict:
         info["uri_handlers"][k] = {"enabled": True}
 
     info["integrations"]["pypi_stats"] = check_pypi_stats()
+    info["integrations"]["reverse_dependencies"] = check_reverse_dependencies()
     info["integrations"]["git"] = check_git()
 
     return info

aura/output/text.py

@@ -52,6 +52,10 @@
 }
 
 
+OK = '\u2713'
+NOK = '\u2717'
+
+
 
 class PrettyReport:
     ANSI_RE = re.compile(r"""
@@ -311,10 +315,14 @@ class TextInfoOutput(InfoOutputBase):
     def protocol(cls) -> str:
         return "text"
 
-    def output_info_data(self, data):
-        OK = '\u2713'
-        NOK = '\u2717'
+    def get_feature_status(self, fmt, name, status):
+        enabled = status.get("enabled", True)
+        mark = OK if enabled else NOK
+        description = status.get("description", "Description N/A")
+        s = {"fg": ("bright_green" if enabled else "bright_red")}
+        return style(fmt.format(mark=mark, status=status, name=name, enabled=enabled, description=description), **s)
 
+    def output_info_data(self, data):
         out = PrettyReport()
 
         # Left hand side of the table contains logo and basic project information
@@ -343,22 +351,15 @@ def output_info_data(self, data):
         rhs_lines.append("Installed analyzers:")
 
         for name, i in data["analyzers"].items():
-            if i["enabled"]:
-                mark = OK
-                s = {"fg": "bright_green"}
-            else:
-                mark = NOK
-                s = {"fg": "bright_red"}
+            rhs_lines.append(self.get_feature_status(fmt=" {mark} {name}: {description}", name=name, status=i))
 
-            rhs_lines.append(style(f" {mark} {name}: {i['description']}", **s))
+        rhs_lines.append("Integrations:")
+        for name, i in data["integrations"].items():
+            rhs_lines.append(self.get_feature_status(fmt=" {mark} {name}: {description}", name=name, status=i))
 
         rhs_lines.append("Installed URI handlers:")
         for name, i in data["uri_handlers"].items():
-            enabled = i.get("enabled", True)
-            mark = OK if enabled else NOK
-            s = {"fg": ("bright_green" if enabled else "bright_red") }
-
-            rhs_lines.append(style(f" {mark} `{name}://` - {i.get('description', 'Description N/A')}", **s))
+            rhs_lines.append(self.get_feature_status(fmt=" {mark} `{name}://` - {description}", name=name, status=i))
 
         rhs_size = max(len(x) for x in rhs_lines)
 

aura/update.py

@@ -1,14 +1,15 @@
 import os
 import shutil
+import tarfile
+import tempfile
 from pathlib import Path
 
 from click import secho
 
 from . import utils
 
 
-STATS_CDN_URL = "https://cdn.sourcecode.ai/aura/pypi_stats.json"
-REVERSE_CDN_URL = "https://cdn.sourcecode.ai/aura/reverse_dependencies.json"
+DATASET_CDN_URL = "https://cdn.sourcecode.ai/aura/aura_dataset.tgz"
 
 
 def backup_file(file_path):
@@ -20,28 +21,19 @@ def backup_file(file_path):
         shutil.copyfile(file_path, f"{file_path}.bak")
 
 
-def update_pypi_stats(outfile=None):
-    if outfile is None:
-        outfile = Path("pypi_stats.json")
+def update_dataset():
+    cwd = Path.cwd()
 
-    backup_file(outfile)
-    secho("Downloading latest pypi download stats dataset")
+    secho("Downloading latest aura dataset files")
+    with tempfile.NamedTemporaryFile(prefix="aura_dataset_update_", suffix=".tgz", mode="wb") as fd:
+        utils.download_file(DATASET_CDN_URL, fd)
 
-    fd = outfile.open("wb")
-    utils.download_file(STATS_CDN_URL, fd)
+        archive = tarfile.open(fd.name, "r:*")
 
-
-def update_reverse_dependencies(outfile=None):
-    if outfile is None:
-        outfile = Path("reverse_dependencies.json")
-
-    backup_file(outfile)
-    secho("Downloading latest reverse PyPI dependencies dataset")
-
-    fd = outfile.open("wb")
-    utils.download_file(REVERSE_CDN_URL, fd)
+        for f in ("reverse_dependencies.json", "pypi_stats.json"):
+            backup_file(cwd/f)
+            archive.extract(f, path=cwd)
 
 
 def update_all():
-    update_pypi_stats()
-    update_reverse_dependencies()
+    update_dataset()

docs/source/configuration.rst

@@ -115,14 +115,16 @@ Environment config options
 
 The following environment variable configuration options can be used to configure the aura behaviour:
 
-======================= =============================================================
-Environment variable    Explanation
-======================= =============================================================
-AURA_CFG                Overwrite the path to the main configuration file
-AURA_SIGNATURES         Overwrite the path to the configuration file for signatures/patterns
-AURA_MIRROR_PATH        Location to the local pypi mirror repository
-AURA_LOG_LEVEL          Output log level
-AURA_NO_BLOBS           Disable extraction of data blobs for further analysis
-AURA_NO_PROGRESS        Disable cli progress bar, useful when redirecting stderr and stdout
-AURA_DEBUG_LINES        List of line numbers separated by ``,``. Aura will then call ``breakpoint()`` when traversing AST tree and it visits a node located on those specific line numbers
-======================= =============================================================
+=========================== =============================================================
+Environment variable        Explanation
+=========================== =============================================================
+AURA_CFG                    Overwrite the path to the main configuration file
+AURA_SIGNATURES             Overwrite the path to the configuration file for signatures/patterns
+AURA_MIRROR_PATH            Location to the local pypi mirror repository
+AURA_PYPI_STATS             Overwrite the path to the aura `pypi_stats` dataset
+AURA_REVERSE_DEPENDENCIES   Overwrite the path to the aura `reverse_dependencies` dataset
+AURA_LOG_LEVEL              Output log level
+AURA_NO_BLOBS               Disable extraction of data blobs for further analysis
+AURA_NO_PROGRESS            Disable cli progress bar, useful when redirecting stderr and stdout
+AURA_DEBUG_LINES            List of line numbers separated by ``,``. Aura will then call ``breakpoint()`` when traversing AST tree and it visits a node located on those specific line numbers
+=========================== =============================================================