Update src/aibom_generator/api.py

src/aibom_generator/api.py

@@ -13,6 +13,10 @@ from datasets import Dataset, load_dataset, concatenate_datasets
 import os
 import logging
 from urllib.parse import urlparse
+import re # Import regex module
+import html # Import html module for escaping
+from huggingface_hub import HfApi
+from huggingface_hub.utils import RepositoryNotFoundError # For specific error handling
 
 # Configure logging
 logging.basicConfig(level=logging.INFO)
@@ -45,21 +49,51 @@ class StatusResponse(BaseModel):
     version: str
     generator_version: str
 
-# --- Add Model ID Normalization Helper --- 
+
+# --- Model ID Validation and Normalization Helpers --- 
+# Regex for valid Hugging Face ID parts (alphanumeric, -, _, .)
+# Allows owner/model format
+HF_ID_REGEX = re.compile(r"^[a-zA-Z0-9\.\-\_]+/[a-zA-Z0-9\.\-\_]+$")
+
+def is_valid_hf_input(input_str: str) -> bool:
+    """Checks if the input is a valid Hugging Face model ID or URL."""
+    if not input_str or len(input_str) > 200: # Basic length check
+        return False
+        
+    if input_str.startswith(("http://", "https://") ):
+        try:
+            parsed = urlparse(input_str)
+            # Check domain and path structure
+            if parsed.netloc == "huggingface.co":
+                path_parts = parsed.path.strip("/").split("/")
+                # Must have at least owner/model, can have more like /tree/main
+                if len(path_parts) >= 2 and path_parts[0] and path_parts[1]:
+                     # Check characters in the relevant parts
+                     if re.match(r"^[a-zA-Z0-9\.\-\_]+$", path_parts[0]) and \
+                        re.match(r"^[a-zA-Z0-9\.\-\_]+$", path_parts[1]):
+                         return True
+            return False # Not a valid HF URL format
+        except Exception:
+            return False # URL parsing failed
+    else:
+        # Assume owner/model format, check with regex
+        return bool(HF_ID_REGEX.match(input_str))
+
 def _normalise_model_id(raw_id: str) -> str:
     """
-    Accept either 'owner/model' or a full URL like
+    Accept either validated 'owner/model' or a validated full URL like
     'https://huggingface.co/owner/model'. Return 'owner/model'.
+    Assumes input has already been validated by is_valid_hf_input.
     """
     if raw_id.startswith(("http://", "https://") ):
         path = urlparse(raw_id).path.lstrip("/")
-        # path can contain extra segments (e.g. /commit/...), keep first two
         parts = path.split("/")
-        if len(parts) >= 2:
-            return "/".join(parts[:2])
-        return path # Fallback if path doesn't have owner/model
-    return raw_id
-# --- End Model ID Normalization Helper ---
+        # We know from validation that parts[0] and parts[1] exist
+        return f"{parts[0]}/{parts[1]}"
+    return raw_id # Already in owner/model format
+
+# --- End Model ID Helpers ---
+
 
 # --- Add Counter Helper Functions ---
 def log_sbom_generation(model_id: str):
@@ -370,9 +404,45 @@ async def generate_form(
     use_best_practices: bool = Form(True)
 ):
     sbom_count = get_sbom_count() # Get count early for context
-    # --- Normalize the model ID for display and filename --- 
-    normalized_model_id = _normalise_model_id(model_id) 
-    # --- End Normalization ---
+    
+    
+    # --- Input Sanitization --- 
+    sanitized_model_id = html.escape(model_id)
+    
+    # --- Input Format Validation --- 
+    if not is_valid_hf_input(sanitized_model_id):
+        error_message = "Invalid input format. Please provide a valid Hugging Face model ID (e.g., 'owner/model') or a full model URL (e.g., 'https://huggingface.co/owner/model') ."
+        logger.warning(f"Invalid model input format received: {model_id}") # Log original input
+        # Try to display sanitized input in error message
+        return templates.TemplateResponse(
+            "error.html", {"request": request, "error": error_message, "sbom_count": sbom_count, "model_id": sanitized_model_id}
+        )
+        
+    # --- Normalize the SANITIZED and VALIDATED model ID --- 
+    normalized_model_id = _normalise_model_id(sanitized_model_id)
+    
+    # --- Check if the ID corresponds to an actual HF Model --- 
+    try:
+        hf_api = HfApi()
+        logger.info(f"Attempting to fetch model info for: {normalized_model_id}")
+        model_info = hf_api.model_info(normalized_model_id)
+        logger.info(f"Successfully fetched model info for: {normalized_model_id}")
+    except RepositoryNotFoundError:
+        error_message = f"Error: The provided ID \"{normalized_model_id}\" could not be found on Hugging Face or does not correspond to a model repository."
+        logger.warning(f"Repository not found for ID: {normalized_model_id}")
+        return templates.TemplateResponse(
+            "error.html", {"request": request, "error": error_message, "sbom_count": sbom_count, "model_id": normalized_model_id}
+        )
+    except Exception as api_err: # Catch other potential API errors
+        error_message = f"Error verifying model ID with Hugging Face API: {str(api_err)}"
+        logger.error(f"HF API error for {normalized_model_id}: {str(api_err)}")
+        return templates.TemplateResponse(
+            "error.html", {"request": request, "error": error_message, "sbom_count": sbom_count, "model_id": normalized_model_id}
+        )
+    # --- End Model Existence Check ---
+
+    
+    # --- Main Generation Logic --- 
     try:
         # Try different import paths for AIBOMGenerator
         generator = None
@@ -391,32 +461,31 @@ async def generate_form(
                     logger.error("Could not import AIBOMGenerator from any known location")
                     raise ImportError("Could not import AIBOMGenerator from any known location")
 
-        # Generate AIBOM
+        # Generate AIBOM (pass SANITIZED ID)
         aibom = generator.generate_aibom(
-            model_id=model_id,
+            model_id=sanitized_model_id, # Use sanitized ID
             include_inference=include_inference,
             use_best_practices=use_best_practices
         )
         enhancement_report = generator.get_enhancement_report()
 
-        # Save AIBOM to file
+        # Save AIBOM to file, use industry term ai_sbom in file name
         # Corrected: Removed unnecessary backslashes around '/' and '_'
         # Save AIBOM to file using normalized ID
-        filename = f"{normalized_model_id.replace('/', '_')}_aibom.json"
+        filename = f"{normalized_model_id.replace('/', '_')}_ai_sbom.json"
         filepath = os.path.join(OUTPUT_DIR, filename)
 
         with open(filepath, "w") as f:
             json.dump(aibom, f, indent=2)
 
         # --- Log Generation Event ---
-        log_sbom_generation(model_id)
+        log_sbom_generation(sanitized_model_id) # Use sanitized ID
         sbom_count = get_sbom_count() # Refresh count after logging
         # --- End Log ---
 
         download_url = f"/output/{filename}"
 
         # Create download and UI interaction scripts
-        # Corrected: Removed unnecessary backslashes in script string
         download_script = f"""
         <script>
             function downloadJSON() {{
@@ -461,17 +530,16 @@ async def generate_form(
         """
 
         # Get completeness score or create a comprehensive one if not available
+        # Use sanitized_model_id
         completeness_score = None
-        # Corrected: Removed unnecessary backslash in 'get_completeness_score'
         if hasattr(generator, 'get_completeness_score'):
             try:
-                completeness_score = generator.get_completeness_score(model_id)
+                completeness_score = generator.get_completeness_score(sanitized_model_id)
                 logger.info("Successfully retrieved completeness_score from generator")
             except Exception as e:
                 logger.error(f"Completeness score error from generator: {str(e)}")
 
         # If completeness_score is None or doesn't have field_checklist, use comprehensive one
-        # Corrected: Removed unnecessary backslash in doesn't and 'field_checklist'
         if completeness_score is None or not isinstance(completeness_score, dict) or 'field_checklist' not in completeness_score:
             logger.info("Using comprehensive completeness_score with field_checklist")
             completeness_score = create_comprehensive_completeness_score(aibom)
@@ -537,15 +605,19 @@ async def generate_form(
                 "display_names": display_names,
                 "tooltips": tooltips,
                 "weights": weights,
-                "sbom_count": sbom_count # Pass count
+                "sbom_count": sbom_count,
+                "display_names": display_names,
+                "tooltips": tooltips,
+                "weights": weights
             }
         )
+    # --- Main Exception Handling --- 
     except Exception as e:
         logger.error(f"Error generating AI SBOM: {str(e)}")
-        # Ensure count is passed to error template as well
         sbom_count = get_sbom_count() # Refresh count just in case
+        # Pass count, added normalized model ID
         return templates.TemplateResponse(
-            "error.html", {"request": request, "error": str(e), "sbom_count": sbom_count, "model_id": normalized_model_id} # Pass count, added normalized model ID
+            "error.html", {"request": request, "error": str(e), "sbom_count": sbom_count, "model_id": normalized_model_id}
         )
 
 @app.get("/download/{filename}")