- improve script for pypi mirror scan

- bugfix unit test if AURA_MIRROR_PATH is defined
- smaller cosmetic code changes

Author: RootLUG

aura/analyzers/package_enrichment.py

@@ -10,7 +10,9 @@ def analyze(*, location: ScanLocation) -> AnalyzerReturnType:
     if not (pkg_name:=location.metadata.get("package_name")):
         return
 
-    pkg = package.PypiPackage.from_cached(pkg_name)
+    if not (pkg:=location.metadata.get("package_instance")):
+        pkg = package.PypiPackage.from_cached(pkg_name)
+
     pkg_score = pkg.score
 
     extra = {

aura/cache.py

@@ -9,7 +9,7 @@
 import concurrent.futures
 from abc import ABC, abstractmethod
 from pathlib import Path
-from typing import Optional, List, Generator, Iterable, BinaryIO, Tuple
+from typing import Optional, List, Generator, Iterable, BinaryIO, Tuple, Set
 
 import click
 import requests
@@ -29,20 +29,29 @@ def __init__(self, path: Path):
         self.path = path
         self.metadata = loads(path.read_text())
 
-        self.cls = CACHE_TYPES[self.metadata["type"]]
+        self.cls = CACHE_TYPES[self.type]
         self.item_path = path.parent / f"{self.cls.prefix}{self.metadata['id']}"
         self.item_stat = self.item_path.stat()
         # Used to avoid re-listing the cache content to find deleted items
         # Used for example in tests to assert which cache items were deleted
         self._deleted = False
 
     @classmethod
-    def iter_items(cls) -> Generator[CacheItem, None, None]:
+    def iter_items(cls, tags=None) -> Generator[CacheItem, None, None]:
+        tags = set(tags or ())
+
         for x in Cache.get_location().iterdir():
             if not x.name.endswith(".metadata.json"):
                 continue
 
-            yield cls(x)
+            obj = cls(x)
+
+            if tags:
+                if obj.type in tags or tags.intersection(obj.tags):
+                    yield obj
+            else:
+                yield obj
+
 
     @property
     def mtime(self) -> int:
@@ -59,6 +68,14 @@ def is_expired(self) -> bool:
         exp_threshold = get_expiration()
         return now > modified+exp_threshold
 
+    @property
+    def type(self) -> str:
+        return self.metadata["type"]
+
+    @property
+    def tags(self) -> Set[str]:
+        return set(self.metadata.get("tags", ()))
+
     def delete(self):
         self.item_path.unlink(missing_ok=True)
         self.path.unlink(missing_ok=True)
@@ -144,10 +161,6 @@ def get_location(cls) -> Optional[Path]:
 
         return cls.__location
 
-    @classmethod
-    def proxy_url(cls, *, url, fd, cache_id=None):
-        return FileDownloadCache.proxy(url=url, fd=fd, cache_id=cache_id)
-
     def save_metadata(self):
         self.metadata_location.write_text(dumps(self.metadata))
 

aura/uri_handlers/mirror.py

@@ -17,11 +17,7 @@ def __init__(self, uri: ParseResult):
 
         self.package_name = uri.hostname
         self.mirror_path = mirror.LocalMirror.get_mirror_path()  # Path(uri.path)
-
-        if self.package_name == "$all":
-            self.package = "$all"
-        else:
-            self.package = PypiPackage.from_cached(self.package_name)
+        self.package = PypiPackage.from_cached(self.package_name)
 
         self.opts.update(parse_qs(uri.query))
         self.comment = uri.fragment.lstrip("#")

files/scan_pypi_mirror.sh

@@ -1,39 +1,57 @@
 #!/bin/bash
 set -e
 
-if [ "$#" -ne 1 ]; then
-  echo "You must provide a path to the offline PyPI mirror web folder" >>2;
-  exit 1;
-fi;
-
-export AURA_MIRROR_PATH=$1;
 export AURA_ALL_MODULE_IMPORTS=true;
 export PYTHONWARNINGS=ignore;
-export TEMPDIR=$(dirname $(mktemp -u))
+export OUTDIR=${AURA_SCAN_DIR:=aura_mirror_scan}
+
+
+if [[ -z "${AURA_MIRROR_PATH}" ]]; then
+  echo "You must set the AURA_MIRROR_PATH env variable!" >>2;
+  exit 1
+fi;
+
+
+# Create directory structure
+[ -d $OUTDIR ] || mkdir $OUTDIR
+[ -d $OUTDIR/package_errors ] || mkdir -p $OUTDIR/package_errors
+[ -d $OUTDIR/package_results ] || mkdir -p $OUTDIR/package_results
 
 
 if [ ! -d "${AURA_MIRROR_PATH}/json" ]; then
   echo "JSON directory not found at ${AURA_MIRROR_PATH}. You probably have not provided a correct path to the web mirror directory" >>2;
   exit 1;
 fi
 
-if [ ! -f "aura_mirror_scan/package_cache" ]; then
-  ls $AURA_MIRROR_PATH/json >aura_mirror_scan/package_cache;
+if [ ! -f "$OUTDIR/package_cache" ]; then
+  if [ -f $AURA_MIRROR_PATH/pypi_package_list.txt]; then
+    cp $AURA_MIRROR_PATH/pypi_package_list.txt $OUTDIR/package_cache;
+  else
+    ls $AURA_MIRROR_PATH/json >$OUTDIR/package_cache;
 fi
 
 
-PKGS=$(cat aura_mirror_scan/package_cache)
+PKGS=$(cat $OUTDIR/package_cache)
 
 scan() {
-  AURA_LOG_LEVEL="ERROR" AURA_NO_PROGRESS=true aura scan -f json mirror://$1 -v 1> >(tee -a "aura_mirror_scan/$1.results.json" |jq .) 2> >(tee -a aura_mirror_scan/$1.errors.log >&2)
+  ERROR_FILE=$OUTDIR/package_errors/$1.errors.log
+  RESULTS_FILE=$OUTDIR/package_results/$1.results.json
+
+  AURA_LOG_LEVEL="ERROR" AURA_NO_PROGRESS=true aura scan -f json mirror://$1 -v 1> >(tee -a $RESULTS_FILE |jq .) 2> >(tee -a $ERROR_FILE >&2)
   if [ $? -ne 0 ]; then
-    echo $1 >>aura_mirror_scan/failed_packages.log
+    echo $1 >>$OUTDIR/failed_packages.log
   else
-    echo $1 >>aura_mirror_scan/processed_packages.log
+    echo $1 >>$OUTDIR/processed_packages.log
+  fi
+
+  if [ -s $RESULTS_FILE ]; then
+    echo "Removing empty $RESULTS_FILE"
+    rm $RESULTS_FILE
   fi
 
-  if [ -s aura_mirror_scan/$1.errors.log ]; then
-    rm aura_mirror_scan/$1.errors.log
+  if [ -s $ERROR_FILE ]; then
+    echo "Removing empty $ERROR_FILE"
+    rm $ERROR_FILE
   fi
 
 }
@@ -42,4 +60,4 @@ export -f scan
 
 echo "Starting Aura scan"
 
-echo $PKGS|tr ' \r' '\n'| parallel --memfree 5G -j30 --progress --resume --timeout 1200 --joblog ${TEMPDIR}/aura_pypi_scan_joblog --max-args 1 scan
+echo $PKGS|tr ' \r' '\n'| parallel --memfree 5G -j30 --progress --resume-failed --timeout 1200 --joblog $OUTDIR/joblog --max-args 1 scan

tests/conftest.py

@@ -20,6 +20,9 @@
 
 os.environ["AURA_NO_CACHE"] = "true"
 
+if "AURA_MIRROR_PATH" in os.environ:
+    os.unsetenv("AURA_MIRROR_PATH")
+
 
 # Definition used to replicate the PyPI mirror file system structure
 MIRROR_FILES = {

tests/test_cache.py

@@ -300,3 +300,32 @@ def test_ast_cache(ast_compile_mock, patterns_hash_mock, mock_cache):
 
     cache.CacheItem.cleanup()
     assert len(tuple(cache.CacheItem.iter_items())) == 0
+
+
+@responses.activate
+def test_cache_tag_filtering(mock_cache, fixtures):
+    def _cb(request):
+        return (200, {}, "Hello world")
+
+    responses.add_callback(method=responses.GET, url="http://example.com/tag_test", match_querystring=False, callback=_cb)
+
+    tags = {"tag1", "tag2", "tag3"}
+
+    for tag in tags:
+        url = f"http://example.com/tag_test?tag={tag}"
+        cache.URLCache.proxy(url=url, tags=[tag])
+
+
+    all_items = tuple(cache.CacheItem.iter_items())
+    assert len(all_items) == len(tags)
+
+    items = tuple(cache.CacheItem.iter_items(tags=["non-existing"]))
+    assert len(items) == 0
+
+    items_by_type = tuple(cache.CacheItem.iter_items(tags=["url"]))
+    assert len(items_by_type) == len(all_items)
+
+    for tag in tags:
+        items = tuple(cache.CacheItem.iter_items(tags=[tag]))
+        assert len(items) == 1
+        assert tag in items[0].tags