Merge pull request #339 from aiverify-foundation/apgw-update-schema-nested-filepaths

wanfei-imda · web-flow · commit 22f388eb12bd · 2024-09-30T14:04:35.000+08:00
Support nested artifact filepath and updated schema and misc fixes
diff --git a/aiverify-apigw/aiverify_apigw/lib/file_utils.py b/aiverify-apigw/aiverify_apigw/lib/file_utils.py
@@ -0,0 +1,46 @@
+from pathlib import PurePath
+import io
+import re
+import hashlib
+
+
+class InvalidFilename(Exception):
+    """Raised when the filename is invalid"""
+    pass
+
+
+def check_valid_filename(filename: str):
+    # return PurePath(filename).stem.isalnum
+    if not filename[0].isalnum():
+        return False
+    # filename must be alphanumeric plus charters ./
+    # sanitized = re.fullmatch(r'[a-zA-Z0-9_./]*', filename)
+    # return sanitized is not None
+    return re.search(r'\.\.', filename) is None
+
+
+def append_filename(filename: str, append_name: str) -> str:
+    fpath = PurePath(filename)
+    return fpath.stem + append_name + fpath.suffix
+
+
+def get_suffix(filename: str) -> str:
+    return PurePath(filename).suffix.lower()
+
+
+def get_stem(filename: str) -> str:
+    return PurePath(filename).stem
+
+
+def get_file_digest(contents: io.BytesIO):
+    contents.seek(0)
+    return hashlib.file_digest(contents, "sha256").digest()
+
+
+def sanitize_filename(filename: str) -> str:
+    if not filename[0].isalnum():
+        raise InvalidFilename("The first character of the filename must be alphanumeric.")
+
+    # Use regex to replace invalid characters
+    sanitized = re.sub(r'[^a-zA-Z0-9./]', '', filename)
+    return sanitized
diff --git a/aiverify-apigw/aiverify_apigw/lib/filestore.py b/aiverify-apigw/aiverify_apigw/lib/filestore.py
@@ -1,9 +1,9 @@
 import os
-import io
-import re
-import hashlib
 import shutil
-from pathlib import Path, PurePath
+from pathlib import Path
+import urllib.parse
+urllib.parse.uses_relative.append("s3")
+urllib.parse.uses_netloc.append("s3")
 from urllib.parse import urljoin
 
 from .s3 import MyS3
@@ -23,18 +23,13 @@ class FileStoreError(Exception):
     pass
 
 
-class InvalidFilename(Exception):
-    """Raised when the filename is invalid"""
-    pass
-
-
-def sanitize_filename(filename: str) -> str:
-    if not filename[0].isalnum():
-        raise InvalidFilename("The first character of the filename must be alphanumeric.")
-
-    # Use regex to replace invalid characters
-    sanitized = re.sub(r'[^a-zA-Z0-9.]', '', filename)
-    return sanitized
+def check_relative_to_base(base_path: Path | str, filepath: str) -> bool:
+    logger.debug(f"check_relative_to_base: {base_path} -> {filepath}")
+    if isinstance(base_path, Path):
+        return base_path.joinpath(filepath).resolve().is_relative_to(base_path)
+    else:
+        # for s3 must be full url
+        return urljoin(base_path, filepath).startswith(base_path)
 
 
 def get_base_data_dir() -> Path | str:
@@ -80,27 +75,6 @@ def is_s3(path: Path | str) -> bool:
         raise InvalidFileStore(f"Invalid path: {path}")
 
 
-def check_valid_filename(filename: str):
-    return PurePath(filename).stem.isalnum
-
-
-def append_filename(filename: str, append_name: str) -> str:
-    fpath = PurePath(filename)
-    return fpath.stem + append_name + fpath.suffix
-
-
-def get_suffix(filename: str) -> str:
-    return PurePath(filename).suffix.lower()
-
-
-def get_stem(filename: str) -> str:
-    return PurePath(filename).stem
-
-
-def get_file_digest(contents: io.BytesIO):
-    contents.seek(0)
-    return hashlib.file_digest(contents, "sha256").digest()
-
 # def absolute_plugin_base_path(source: str) -> Path | str:
 #     if isinstance(base_plugin_dir, Path):
 #         return base_plugin_dir.joinpath(source).resolve()
@@ -143,14 +117,28 @@ def get_plugin_component_folder(gid: str, component_type: str) -> Path | str:
         return urljoin(plugin_path, f"{component_type}/")
 
 
+plugin_ignore_patten = shutil.ignore_patterns(
+    ".venv",
+    "venv",
+    "output",
+    "node_modules",
+    "build",
+    "temp",
+    "__pycache__",
+    ".pytest_cache",
+    ".cache"
+    "*.pyc",
+)
+
+
 def save_plugin(gid: str, source_dir: Path):
     folder = get_plugin_folder(gid)
     logger.debug(f"Copy plugin folder from {source_dir} to {folder}")
     if isinstance(folder, Path):
         if folder.exists():
             shutil.rmtree(folder)
         folder.mkdir(parents=True, exist_ok=True)
-        shutil.copytree(source_dir, folder, dirs_exist_ok=True)
+        shutil.copytree(source_dir, folder, dirs_exist_ok=True, ignore=plugin_ignore_patten)
     elif s3 is not None:
         # folder is s3 prefix
         if s3.check_s3_prefix_exists(folder):
@@ -201,10 +189,17 @@ def get_artifacts_folder(test_result_id: str):
 
 
 def save_artifact(test_result_id: str, filename: str, data: bytes):
-    filename = sanitize_filename(filename)
+    # validate input
+    if not test_result_id.isalnum():
+        raise FileStoreError(f"Invalid test result id {test_result_id}")
     folder = get_artifacts_folder(test_result_id)
+    if not check_relative_to_base(folder, filename):
+        raise FileStoreError(f"Invalid filename {filename}")
     if isinstance(folder, Path):
         filepath = folder.joinpath(filename)
+        # check for nest dir
+        if not filepath.parent.exists():
+            filepath.parent.mkdir(parents=True, exist_ok=True)
         with open(filepath, "wb") as fp:
             fp.write(data)
         return filepath
@@ -215,10 +210,13 @@ def save_artifact(test_result_id: str, filename: str, data: bytes):
 
 
 def get_artifact(test_result_id: str, filename: str):
-    filename = sanitize_filename(filename)
+    if not test_result_id.isalnum():
+        raise FileStoreError(f"Invalid test result id {test_result_id}")
     folder = get_artifacts_folder(test_result_id)
+    if not check_relative_to_base(folder, filename):
+        raise FileStoreError(f"Invalid filename {filename}")
     if isinstance(folder, Path):
-        filepath = folder.joinpath(filename)
+        filepath = folder.joinpath(filename).resolve()
         with open(filepath, "rb") as fp:
             data = fp.read()
         return data
diff --git a/aiverify-apigw/aiverify_apigw/lib/plugin_store.py b/aiverify-apigw/aiverify_apigw/lib/plugin_store.py
@@ -1,6 +1,7 @@
 from pathlib import Path
 import json
 import tomllib
+from typing import List
 
 from .schemas_utils import read_and_validate, plugin_schema, algorithm_schema
 from ..lib.syntax_checker import validate_python_script
@@ -62,8 +63,9 @@ def delete_plugin(cls, gid: str):
     def scan_stock_plugins(cls):
         logger.info(f"Scanning stock plugins in folder {str(cls.stock_plugin_folder)}..")
         cls.delete_all_plugins()  # remove all current plugins first
+        plugins: List[PluginModel] = []
         for plugin_dir in cls.stock_plugin_folder.iterdir():
-            if not plugin_dir.is_dir():
+            if not plugin_dir.is_dir() or not plugin_dir.name[0].isalnum() or plugin_dir.name == "user_defined_files":
                 continue
             logger.debug(f"Scanning directory {plugin_dir.name}")
             try:
@@ -72,13 +74,23 @@ def scan_stock_plugins(cls):
                 logger.debug(f"Invalid plugin: {e}")
                 continue
             try:
-                cls.scan_plugin_directory(plugin_dir)
+                plugin = cls.scan_plugin_directory(plugin_dir)
+                if plugin:
+                    plugins.append(plugin)
             except Exception as e:
                 logger.warning(f"Error saving plugin in directory {plugin_dir.name}: {e}")
         with SessionLocal() as session:
-            stmt = select(func.count("*")).select_from(PluginModel)
-            count = session.scalar(stmt)
-            logger.info(f"Finished scanning stock plugins. {count} plugins found")
+            # stmt = select(func.count("*")).select_from(PluginModel)
+            # count = session.scalar(stmt)
+            stmt = select(PluginModel)
+            my_plugins = list(session.scalars(stmt))
+            logger.info(f"Finished scanning stock plugins. {len(my_plugins)} plugins found")
+            for plugin in my_plugins:
+                logger.info(f"Stock plugin: gid {plugin.gid}, version {plugin.version}, name {plugin.name}")
+                if plugin.algorithms and len(plugin.algorithms) > 0:
+                    logger.info(f"  Number of algoritms: {len(plugin.algorithms)}")
+                    for algo in plugin.algorithms:
+                        logger.info(f"  Algorithm: cid {algo.cid}, version {algo.version}, name {algo.name}")
 
     @classmethod
     def check_plugin_registry(cls):
@@ -167,7 +179,8 @@ def scan_plugin_directory(cls, folder: Path):
                             if "project" not in pyproject_data or "name" not in pyproject_data["project"]:
                                 logger.debug(f"Algorithm folder {algopath.name} has invalid pyproject.toml")
                                 continue
-                            project_name = pyproject_data["project"]["name"]
+                            # TODO: is this the best way to get algorithm folder?
+                            project_name = pyproject_data["project"]["name"].replace("-", "_")
                             sub_path = algopath.joinpath(project_name)
                             module_name = project_name
                             meta_path = sub_path.joinpath("algo.meta.json")
@@ -185,9 +198,9 @@ def scan_plugin_directory(cls, folder: Path):
                         script_path = algopath.joinpath(f"{meta.cid}.py")
                         if not script_path.exists():
                             script_path = algopath.joinpath(f"algo.py")
-                        if script_path.exists(): # if script exists
+                        if script_path.exists():  # if script exists
                             if not validate_python_script(script_path):
-                                logger.warning(f"algorithm {cid} script is not valid")
+                                logger.warning(f"algorithm {meta.cid} script is not valid")
                                 continue
 
                         # validate requirements.txt
@@ -221,7 +234,7 @@ def scan_plugin_directory(cls, folder: Path):
                             input_schema=json.dumps(input_schema).encode("utf-8"),
                             output_schema=json.dumps(output_schema).encode("utf-8"),
                             algo_dir=algopath.relative_to(folder).as_posix(),
-                            language="python", # fixed to python first. To support other languages in future
+                            language="python",  # fixed to python first. To support other languages in future
                             script=script_path.name,
                             module_name=module_name,
                         )
@@ -239,6 +252,7 @@ def scan_plugin_directory(cls, folder: Path):
             session.commit()
 
             fs_save_plugin(plugin_meta.gid, folder)
+            return plugin
 
     @classmethod
     def validate_plugin_directory(cls, folder: Path) -> PluginMeta:
@@ -288,7 +302,8 @@ def validate_plugin_directory(cls, folder: Path) -> PluginMeta:
                     if "project" not in pyproject_data or "name" not in pyproject_data["project"]:
                         logger.debug(f"Algorithm folder {algopath.name} has invalid pyproject.toml")
                         continue
-                    project_name = pyproject_data["project"]["name"]
+                    # TODO: is this the best way to get algorithm folder?
+                    project_name = pyproject_data["project"]["name"].replace("-", "_")
                     sub_path = algopath.joinpath(project_name)
                     meta_path = sub_path.joinpath("algo.meta.json")
                     if not meta_path.exists():
diff --git a/aiverify-apigw/aiverify_apigw/lib/utils.py b/aiverify-apigw/aiverify_apigw/lib/utils.py
@@ -1,5 +1,5 @@
 import mimetypes
-from .filestore import get_suffix
+from .file_utils import get_suffix
 
 
 def guess_mimetype_from_filename(filename: str):
diff --git a/aiverify-apigw/aiverify_apigw/routers/test_result_router.py b/aiverify-apigw/aiverify_apigw/routers/test_result_router.py
@@ -12,8 +12,9 @@
 from ..lib.logging import logger
 from ..lib.database import get_db_session
 from ..lib.constants import TestDatasetFileType, TestDatasetStatus, TestModelMode, TestModelStatus
-from ..lib.filestore import save_artifact, get_artifact, get_suffix
+from ..lib.filestore import save_artifact, get_artifact
 from ..lib.utils import guess_mimetype_from_filename
+from ..lib.file_utils import get_suffix, check_valid_filename
 from ..schemas import TestResult, TestResultOutput, TestResultUpdate
 from ..schemas.load_examples import test_result_examples
 from ..models import AlgorithmModel, TestModelModel, TestResultModel, TestDatasetModel, TestArtifactModel
@@ -37,10 +38,10 @@ async def _save_test_result(session: Session, test_result: TestResult, artifact_
     )
     algorithm = session.scalar(stmt)
     if algorithm is None:
-        raise HTTPException(status_code=400, detail=f"Algorithm {test_result.gid} not found")
+        raise HTTPException(status_code=400, detail=f"Algorithm not found: gid: {test_result.gid}, cid: {test_result.cid}")
 
     now = datetime.now()
-    test_arguments = test_result.test_arguments
+    test_arguments = test_result.testArguments
 
     # validate output
     output_schema = json.loads(algorithm.output_schema.decode("utf-8"))
@@ -133,8 +134,8 @@ async def _save_test_result(session: Session, test_result: TestResult, artifact_
         test_dataset=test_dataset,
         ground_truth_dataset=ground_truth_dataset,
         ground_truth=test_arguments.groundTruth,
-        start_time=test_result.start_time,
-        time_taken=test_result.time_taken,
+        start_time=test_result.startTime,
+        time_taken=test_result.timeTaken,
         algo_arguments=json.dumps(test_arguments.algorithmArgs).encode('utf-8'),
         output=json.dumps(test_result.output).encode('utf-8'),
         created_at=now,
@@ -152,6 +153,9 @@ async def _save_test_result(session: Session, test_result: TestResult, artifact_
             if filename not in artifact_set:
                 logger.warn(f"Unable to find artifact filename {filename} in uploaded files, skipping")
                 continue
+            if not check_valid_filename(filename):
+                logger.warn(f"Invalid artifact filename {filename}, skipping")
+                raise HTTPException(status_code=400, detail=f"Invalid artifact filename in result output: {filename}")
             artifact_file = artifact_set[filename]
             data = artifact_file.file.read()
             save_artifact(test_result_id, filename, data)
@@ -271,12 +275,12 @@ async def upload_zip_file(
                     for zip_info in zip_infos:
                         if zip_info.filename != result_filename and zip_info.filename != foldername:
                             with zip_ref.open(zip_info) as artifact_file:
-                                filename = zip_info.filename[len(foldername):] # remove foldername
+                                filename = zip_info.filename[len(foldername):]  # remove foldername
                                 artifact_set[filename] = UploadFile(
                                     filename=filename,
                                     file=io.BytesIO(artifact_file.read())
-                            )
-                                
+                                )
+
                 logger.debug(f"artifact_set: {artifact_set}")
 
                 if isinstance(results_data, dict):
@@ -311,7 +315,7 @@ async def upload_zip_file(
         raise HTTPException(status_code=500, detail="Internal server error")
 
 
-@router.get("/{test_result_id}/artifacts/{filename}")
+@router.get("/{test_result_id}/artifacts/{filename:path}")
 async def get_test_result_artifact(
     test_result_id: str,
     filename: str,
@@ -320,6 +324,7 @@ async def get_test_result_artifact(
     """
     Endpoint to retrieve an artifact file by test_result_id and filename.
     """
+    logger.debug(f"get_test_result_artifact: {test_result_id}, {filename}")
     try:
         stmt = (
             select(TestArtifactModel)
diff --git a/aiverify-apigw/aiverify_apigw/schemas/examples/test_result_examples.json b/aiverify-apigw/aiverify_apigw/schemas/examples/test_result_examples.json
@@ -3,9 +3,9 @@
     "gid": "aiverify.stock.fairness_metrics_toolbox_for_classification",
     "version": "0.9.0",
     "cid": "fairness_metrics_toolbox_for_classification",
-    "start_time": "2024-07-24T09:20:24.822881",
-    "time_taken": 0,
-    "test_arguments": {
+    "startTime": "2024-07-24T09:20:24.822881",
+    "timeTaken": 0,
+    "testArguments": {
       "testDataset": "file:///examples/data/sample_bc_credit_data.sav",
       "modelFile": "file:///examples/model/sample_bc_credit_sklearn_linear.LogisticRegression.sav",
       "groundTruthDataset": "file:///examples/data/sample_bc_credit_data.sav",
diff --git a/aiverify-apigw/aiverify_apigw/schemas/test_result.py b/aiverify-apigw/aiverify_apigw/schemas/test_result.py
diff --git a/aiverify-apigw/tests/routers/test_test_result_router.py b/aiverify-apigw/tests/routers/test_test_result_router.py
diff --git a/common/schemas/aiverify.algorithm.schema.json b/common/schemas/aiverify.algorithm.schema.json
