- add aura prefetch command for native prefetching of files into cache

- remove old r2c integration that is no longer working/udpated

Author: RootLUG

aura/cli.py

@@ -231,23 +231,10 @@ def find_typosquatting(out, max_distance, limit=None):
     commands.generate_typosquatting(out=out, distance=max_distance, limit=limit)
 
 
-@cli.group("r2c")
-def r2c():
-    pass
-
-
-@r2c.command(name="generate_input")
-@click.argument("out_file", metavar="<OUTPUT FILE>", type=click.File("w"))
-def generate_input(out_file):
-    commands.generate_r2c_input(out_file)
-
-
-@r2c.command(name="scan")
-@click.option("--out", default="/analysis/output/output.json", type=click.File("w"))
-@click.option("--mode", default="generic")
-@click.argument("source", nargs=-1, type=click.Path())
-def run_r2c_analyzer(source, out, mode):
-    commands.r2c_scan(source=source, out_file=out, mode=mode)
+@cli.command()
+@click.argument("uris", nargs=-1, metavar="<URI 1>, <URI 2>, ...")
+def prefetch(uris):
+    commands.prefetch(*uris)
 
 
 @cli.command(name="check_requirement")

aura/commands.py

@@ -5,6 +5,7 @@
 import json
 import time
 import traceback
+from concurrent import futures
 from pathlib import Path
 from functools import partial
 from typing import Union, Optional, Tuple, Generator, List
@@ -223,94 +224,10 @@ def generate_typosquatting(out, distance=2, limit=None):
         out.write(json.dumps({"original": x, "typosquatting": y}) + "\n")
 
 
-def generate_r2c_input(out_file):
-    inputs = []
-
-    for pkg_name in PypiPackage.list_packages():
-        try:
-            pkg = PypiPackage.from_pypi(pkg_name)
-        except exceptions.NoSuchPackage:
-            continue
-        targets = []
-
-        input_definition = {
-            "metadata": {"package": pkg_name},
-            "input_type": "AuraInput",
-        }
-
-        for url in pkg.info["urls"]:
-            targets.append({"url": url["url"], "metadata": url})
-
-        input_definition["targets"] = json.dumps(targets)
-        inputs.append(input_definition)
-
-    out_file.write(
-        json.dumps(
-            {
-                "name": "aura",
-                "version": "0.0.1",
-                "description": "This is a set of all PyPI packages",
-                "inputs": inputs,
-            }
-        )
-    )
-
-
-def r2c_scan(source, out_file, mode="generic"):
-    out = {"results": [], "errors": []}
-
-    pkg_metadata = {}
-
-    metadata = {"format": "none"}
-
-    if mode == "pypi":
-        logger.info("R2C mode set to PyPI")
-        assert len(source) == 1
-        location = Path(source[0])
-
-        meta_loc = location / "metadata.json"
-        if meta_loc.is_file():
-            with open(location / "metadata.json", "r") as fd:
-                pkg_metadata = json.loads(fd.read())
-                metadata.update(
-                    {
-                        "package_type": pkg_metadata.get("packagetype"),
-                        "package_name": pkg_metadata.get("name"),
-                        "python_version": pkg_metadata.get("python_version"),
-                    }
-                )
-        source = [
-            os.fspath(x.absolute())
-            for x in location.iterdir()
-            if x.name != "metadata.json"
-        ]
-    else:
-        logger.info("R2C mode set to generic")
-
-    for src in source:
-        logger.info(f"Enumerating {src} with metadata: {metadata}")
-
-        try:
-            data = scan_uri(src, metadata=metadata)
-
-            for loc in data:
-                for hit in loc["hits"]:
-                    rhit = {"check_id": hit.pop("type"), "extra": hit}
-                    if "line_no" in hit:
-                        rhit["start"] = {"line": hit["line_no"]}
-                        rhit["path"] = os.path.relpath(hit["location"], source[0])
-
-                    out["results"].append(rhit)
-
-        except Exception as exc:
-            exc_tb = sys.exc_info()[-1]
-
-            out["errors"].append(
-                {
-                    "message": f"[{exc_tb.tb_lineno}] An exception occurred: {str(exc)}",
-                    "data": {"path": str(src)},
-                }
-            )
+def prefetch(*uris):
+    with futures.ThreadPoolExecutor() as executor:
+        for uri in uris:
+            handler = URIHandler.from_uri(uri)
+            executor.submit(lambda :list(handler.get_paths()))
 
-    pprint(out)
-    out_file.write(json.dumps(out, default=utils.json_encoder))
+        executor.shutdown(wait=True)

aura/uri_handlers/base.py

@@ -94,7 +94,7 @@ def exists(self) -> bool:
         return True
 
     @abstractmethod
-    def get_paths(self, metadata: Optional[dict]) -> Generator[ScanLocation, None, None]:
+    def get_paths(self, metadata: Optional[dict]=None) -> Generator[ScanLocation, None, None]:
         ...
 
     def get_diff_paths(self, other: URIHandler) -> Generator[Tuple[ScanLocation, ScanLocation], None, None]:
@@ -156,7 +156,6 @@ def __post_init__(self):
                 except PythonExecutorError:
                     pass
 
-
     def __compute_hashes(self):
         tl = tlsh.Tlsh()
         md5 = hashlib.md5()
@@ -186,7 +185,6 @@ def __compute_hashes(self):
         self.metadata["sha256"] = sha256.hexdigest()
         self.metadata["sha512"] = sha512.hexdigest()
 
-
     def __str__(self):
         return self.strip(self.str_location)
 

aura/uri_handlers/git.py

@@ -2,6 +2,7 @@
 import pathlib
 import shutil
 from urllib.parse import ParseResult
+from typing import Optional
 
 from .. import config
 from ..exceptions import PluginDisabled
@@ -41,7 +42,7 @@ def metadata(self):
         }
         return m
 
-    def get_paths(self, metadata: dict=None):
+    def get_paths(self, metadata: Optional[dict]=None):
         if self.opts.get('download_dir') is None:
             p = tempfile.mkdtemp(prefix="aura_git_repo_clone_")
             self.opts["download_dir"] = p

aura/uri_handlers/mirror.py

@@ -1,6 +1,6 @@
 # -*- coding: utf-8 -*-
-import typing
 from urllib.parse import urlparse, ParseResult, parse_qs
+from typing import Optional, Generator
 
 from .base import URIHandler, ScanLocation
 from .. import cache
@@ -30,7 +30,7 @@ def __init__(self, uri: ParseResult):
     def metadata(self):
         return {"package": self.package_name, "package_opts": self.opts}
 
-    def get_paths(self, metadata: dict=None, package=None) -> typing.Generator[ScanLocation, None, None]:
+    def get_paths(self, metadata: Optional[dict]=None, package=None) -> Generator[ScanLocation, None, None]:
         if package is None:
             package = self.package
 

aura/uri_handlers/pypi.py

@@ -4,7 +4,7 @@
 import tempfile
 import pathlib
 import urllib.parse
-from typing import Generator, Tuple
+from typing import Generator, Tuple, Optional
 
 from .base import URIHandler, PackageProvider, ScanLocation
 from ..exceptions import UnsupportedDiffLocation
@@ -53,7 +53,7 @@ def metadata(self):
         }
         return m
 
-    def get_paths(self, metadata: dict=None):
+    def get_paths(self, metadata: Optional[dict]=None):
         if self.opts.get("download_dir") is None:
             self.opts["download_dir"] = pathlib.Path(
                 tempfile.mkdtemp(prefix="aura_pypi_download_")

tests/test_cli.py

@@ -158,23 +158,6 @@ def test_ast_parser(fixtures):
     assert "adalaraoawa aoalalaeaH" in result.stdout
 
 
-def test_r2c_integration():
-    f_name = 'r2c_test_output.json'
-    runner = CliRunner()
-    with runner.isolated_filesystem():
-        result = runner.invoke(cli.cli, ['r2c', 'scan', '--out', f_name])
-
-        if result.exception:
-            raise result.exception
-
-        assert result.exit_code == 0
-        with open(f_name, 'r') as fd:
-            data = json.loads(fd.read())
-
-    assert 'results' in data
-    assert 'errors' in data
-
-
 def test_async_cleanup(fixtures):
     from aura.uri_handlers import base