- improve caching mechanism to ensure zero network transfers when content is already cached from mirror

- extend cache tests

Author: RootLUG

aura/analyzers/python/visitor.py

@@ -4,13 +4,11 @@
 from __future__ import annotations
 
 import os
-import copy
 import time
 from functools import partial, wraps
 from collections import deque, OrderedDict
-from pathlib import Path
 from warnings import warn
-from typing import Optional, Tuple, Union
+from typing import Optional, Tuple, Union, Dict
 
 import pkg_resources
 
@@ -45,7 +43,7 @@ def wrapper(*args, **kwargs):
     return wrapper
 
 
-def get_ast_tree(location: Union[ScanLocation, bytes], metadata=None):
+def get_ast_tree(location: Union[ScanLocation, bytes], metadata=None) -> dict:
     if type(location) == bytes:
         kwargs = {
             "command": [INSPECTOR_PATH, "-"],
@@ -103,7 +101,7 @@ def __init__(self, *, location: ScanLocation):
         self.max_queue_size = int(config.get_settings("aura.max-ast-queue-size", 10000))
 
     @classmethod
-    def from_visitor(cls, visitor: Visitor):
+    def from_visitor(cls, visitor: Visitor) -> Visitor:
         obj = cls(location=visitor.location)
         obj.tree = visitor.tree
         obj.hits = visitor.hits
@@ -138,7 +136,7 @@ def run_stages(cls, *, location: ScanLocation, stages: Optional[Tuple[str, ...]]
         return v
 
     @classmethod
-    def get_visitors(cls):
+    def get_visitors(cls) -> Dict[str, Visitor]:
         global VISITORS
         if VISITORS is None:
             VISITORS = {

aura/cache.py

@@ -66,9 +66,7 @@ def proxy_url(cls, *, url, fd, cache_id=None):
 
     @classmethod
     def proxy_mirror(cls, *, src: Path, cache_id=None):
-        if not src.exists():
-            return None
-        elif cls.get_location() is None:
+        if cls.get_location() is None:  # Caching is disabled
             return src
 
         if cache_id is None:
@@ -81,6 +79,9 @@ def proxy_mirror(cls, *, src: Path, cache_id=None):
             logger.debug(f"Retrieving mirror file path {cache_id} from cache")
             return cache_pth
 
+        if not src.exists():
+            return src
+
         try:
             shutil.copyfile(src, cache_pth, follow_symlinks=True)
             return cache_pth
@@ -90,18 +91,19 @@ def proxy_mirror(cls, *, src: Path, cache_id=None):
 
     @classmethod
     def proxy_mirror_json(cls, *, src: Path):
-        if not src.exists():
-            return src
-        elif cls.get_location() is None:
+        if cls.get_location() is None:  # Caching is disabled
             return src
 
         cache_id = f"mirrorjson_{src.name}"
-        cache_path = cls.get_location()/cache_id
+        cache_path = cls.get_location() / cache_id
 
         if cache_path.exists():
             logger.debug(f"Retrieving package mirror JSON {cache_id} from cache")
             return cache_path
 
+        if not src.exists():
+            return src
+
         try:
             shutil.copyfile(src=src, dst=cache_path, follow_symlinks=True)
             return cache_path

aura/config.py

@@ -16,7 +16,7 @@
 import tqdm
 import pkg_resources
 import ruamel.yaml
-from ruamel.yaml import YAML, composer
+from ruamel.yaml import composer
 try:
     import rapidjson as json
 except ImportError:
@@ -33,8 +33,8 @@
 # This is used to trigger breakpoint during AST traversing of specific lines
 DEBUG_LINES = set()
 DEFAULT_AST_STAGES = ("convert", "rewrite", "ast_pattern_matching", "taint_analysis", "readonly")
-AST_PATTERNS_CACHE = None
-PROGRESSBAR_DISABLED = ("AURA_NO_PROGRESS" in os.environ)
+AST_PATTERNS_CACHE: Optional[tuple] = None
+PROGRESSBAR_DISABLED: bool = ("AURA_NO_PROGRESS" in os.environ)
 
 DEFAULT_CFG_PATH = "aura.data.aura_config.yaml"
 DEFAULT_SIGNATURE_PATH = "aura.data.signatures.yaml"
@@ -185,7 +185,7 @@ def get_file_content(location: str, base_path: Optional[str]=None) -> str:
             return fd.read()
 
 
-def parse_config(pth, default_pth):
+def parse_config(pth, default_pth) -> dict:
     logger.debug(f"Aura configuration located at {pth}")
 
     content = get_file_content(pth)
@@ -282,7 +282,7 @@ def get_ast_stages() -> typing.Tuple[str,...]:
     return [x for x in cfg_value if x]
 
 
-def get_ast_patterns():
+def get_ast_patterns() -> tuple:
     global AST_PATTERNS_CACHE
     from .pattern_matching import ASTPattern
 
@@ -295,7 +295,5 @@ def get_ast_patterns():
     return AST_PATTERNS_CACHE
 
 
-
-
 CFG_PATH = find_configuration()
 load_config()

aura/mirror.py

@@ -12,7 +12,7 @@
 from .config import CFG
 
 
-class LocalMirror(object):
+class LocalMirror:
     @classmethod
     def get_mirror_path(cls) -> Path:
         env_path = os.environ.get('AURA_MIRROR_PATH', None)
@@ -24,23 +24,20 @@ def get_mirror_path(cls) -> Path:
     def list_packages(self) -> typing.Generator[Path, None, None]:
         yield from (self.get_mirror_path() / "json").iterdir()
 
-    def get_json(self, package_name):
-        if package_name is None:
-            raise NoSuchPackage(f"Could not find package '{package_name}' json at the mirror location")
-
+    def get_json(self, package_name) -> dict:
+        assert package_name
         json_path = self.get_mirror_path() / "json" / package_name
+        target = cache.Cache.proxy_mirror_json(src=json_path)
 
-        if not json_path.is_file():
+        if not target.is_file():
             json_path = self.get_mirror_path() / "json" / canonicalize_name(package_name)
-            if not json_path.exists():
+            target = cache.Cache.proxy_mirror_json(src=json_path)
+            if not target.exists():
                 raise NoSuchPackage(package_name)
 
-        target = cache.Cache.proxy_mirror_json(src=json_path)
-
-        with open(target, "r") as fd:
-            return json.loads(fd.read())
+        return json.loads(target.read_text())
 
-    def url2local(self, url):
+    def url2local(self, url: typing.Union[ParseResult, str]) -> Path:
         if not isinstance(url, ParseResult):
             url = urlparse(url)
 

aura/pattern_matching.py

@@ -44,7 +44,7 @@ def match_node(self, context: nodes.Context) -> bool:
         ...
 
     @property
-    def message(self):
+    def message(self) -> str:
         """
         return a message identifying the match
         """
@@ -121,7 +121,7 @@ def match_node(self, node: nodes.NodeType) -> bool:
         return False
 
     @abstractmethod
-    def match_string(self, value: str):
+    def match_string(self, value: str) -> bool:
         ...
 
 

files/prefetch_mirror.sh

@@ -12,6 +12,7 @@
 
 from click.testing import CliRunner, Result
 import responses
+import tqdm
 import pytest
 
 
@@ -370,3 +371,10 @@ def reset_plugins():
 
     ReadOnlyAnalyzer.hooks = []
     plugins.PLUGIN_CACHE = {}
+
+
+@pytest.fixture(scope="module")
+def mock_tqdm_log_write():
+
+    with mock.patch.object(tqdm.tqdm, "write") as m:
+        yield m

tests/conftest.py

@@ -1,16 +1,67 @@
 import json
+import uuid
 from unittest import mock
 
+import pytest
+
 from aura import cache
+from aura import mirror
+from aura import exceptions
+
+
+
+@mock.patch("aura.cache.Cache.get_location")
+def test_cache_mock_location(cache_mock, tmp_path):
+    cache_mock.return_value = tmp_path
+    assert cache.Cache.get_location() == tmp_path
+
+
+@mock.patch("aura.cache.Cache.get_location")
+@pytest.mark.parametrize("filename,content,cache_id,call", (
+        ("testjson_file", "json_content", "mirrorjson_testjson_file", cache.Cache.proxy_mirror_json),
+        ("testpkg_file", "pkg_content", "mirror_testpkg_file", cache.Cache.proxy_mirror)
+))
+def test_proxy_mirror_json(cache_mock, tmp_path, filename, content, cache_id, call):
+    f = tmp_path / filename
+    cache_path = tmp_path / "cache"
+    cache_path.mkdir()
+    cache_file = cache_path/cache_id
+    cache_mock.return_value = cache_path
+
+    assert f.exists() is False
+    out = call(src=f)
+    assert out == f
+    assert cache_file.exists() is False
+    assert len(list(cache_path.iterdir())) == 0
+
+    f.write_text(content)
+    assert f.exists() is True
+    out = call(src=f)
+    assert out != f
+    assert out == cache_file
+    assert len(list(cache_path.iterdir())) == 1
+    assert out.read_text() == content
 
+    # Make sure the cache does not attempt to do any kind of file access if the cache entry exists
+    m = mock.MagicMock(spec_set=("name",), side_effect=ValueError("Call prohibited"))
+    m.name = filename
+    out = call(src=m)
+    assert out == cache_file
 
-def test_mirror_cache(fixtures, simulate_mirror, tmp_path):
+    # Original path should be returned if cache is disabled
+    cache_mock.return_value = None
+    out = call(src=f)
+    assert out == f
+
+
+@mock.patch("aura.cache.Cache.get_location")
+def test_mirror_cache(cache_mock, fixtures, simulate_mirror, tmp_path):
     cache_content = list(tmp_path.iterdir())
     assert len(cache_content) == 0
 
-    with mock.patch.object(cache.Cache, 'get_location', return_value=tmp_path) as m:
-        assert cache.Cache.get_location() == tmp_path
-        out = fixtures.get_cli_output(['scan', '--download-only', 'mirror://wheel', '-f', 'json'])
+    cache_mock.return_value = tmp_path
+    assert cache.Cache.get_location() == tmp_path
+    out = fixtures.get_cli_output(['scan', '--download-only', 'mirror://wheel', '-f', 'json'])
 
     parsed_output = json.loads(out.stdout)
     assert len(parsed_output["detections"]) == 0
@@ -19,3 +70,31 @@ def test_mirror_cache(fixtures, simulate_mirror, tmp_path):
     assert len(cache_content) > 0
     assert "mirror_wheel-0.34.2.tar.gz" in cache_content, cache_content
     assert "mirror_wheel-0.34.2-py2.py3-none-any.whl" in cache_content
+
+
+@mock.patch("aura.cache.Cache.get_location")
+@mock.patch("aura.mirror.LocalMirror.get_mirror_path")
+def test_mirror_cache_no_remote_access(mirror_mock, cache_mock, fixtures, tmp_path):
+    """
+    Test that if the content is fully cached, the mirror uri handler does not attempt to access the mirror but rather retrieves **all** content from cache only
+    This is mainly to test correctness of prefetching the data for global PyPI scan to ensure no further network calls are made
+    """
+    pkg = str(uuid.uuid4())
+    pkg_content = {"id": pkg}
+    mirror_path = tmp_path / "mirror"
+    cache_path = tmp_path / "cache"
+    cache_path.mkdir()
+    cache_mock.return_value = cache_path
+    mirror_mock.return_value = mirror_path
+    m = mirror.LocalMirror()
+
+    assert cache.Cache.get_location() == cache_path
+    assert m.get_mirror_path() == mirror_path
+    assert mirror_path.exists() == False
+
+    with pytest.raises(exceptions.NoSuchPackage):
+        m.get_json(pkg)
+
+    (cache_path/f"mirrorjson_{pkg}").write_text(json.dumps(pkg_content))
+    out = m.get_json(pkg)
+    assert out == pkg_content