From e521af228178a0c2815d5a9f501c48dcde96b299 Mon Sep 17 00:00:00 2001
From: JSC <jschoisy@carea5.com>
Date: Fri, 11 Jul 2025 15:10:36 +0200
Subject: [PATCH] Add vulnerability counts to project response in get_projects
 endpoint

---
 main.py | 56 +++++++++++++++++++++++++++++++++++++++++++++++++++++++-
 1 file changed, 55 insertions(+), 1 deletion(-)

diff --git a/main.py b/main.py
index 5958272..defa4b9 100644
--- a/main.py
+++ b/main.py
@@ -155,6 +155,7 @@ class ProjectResponse(BaseModel):
     is_active: bool
     created_at: datetime
     updated_at: datetime
+    vulnerability_counts: Optional[dict] = None
 
 
 class FileResponse(BaseModel):
@@ -287,6 +288,7 @@ async def get_projects(
     skip: int = 0,
     limit: int = 100,
     active_only: bool = True,
+    include_vulnerability_counts: bool = False,
     db: Session = Depends(get_db)
 ):
     query = db.query(Project)
@@ -294,7 +296,59 @@ async def get_projects(
         query = query.filter(Project.is_active == True)
     
     projects = query.offset(skip).limit(limit).all()
-    return projects
+    
+    if not include_vulnerability_counts:
+        return projects
+    
+    # Add vulnerability counts for each project
+    result = []
+    for project in projects:
+        # Get all images for this project
+        project_images = db.query(Image).join(FileImageUsage).join(File).filter(
+            File.project_id == project.id,
+            Image.is_active == True,
+            FileImageUsage.is_active == True
+        ).distinct().all()
+        
+        # Count vulnerabilities by severity for all images in this project
+        vulnerability_counts = {
+            'critical': 0,
+            'high': 0,
+            'medium': 0,
+            'low': 0,
+            'unspecified': 0,
+            'total': 0
+        }
+        
+        for image in project_images:
+            vulnerabilities = db.query(Vulnerability).filter(
+                Vulnerability.image_id == image.id,
+                Vulnerability.is_active == True
+            ).all()
+            
+            for vuln in vulnerabilities:
+                severity = vuln.severity.lower()
+                if severity in vulnerability_counts:
+                    vulnerability_counts[severity] += 1
+                else:
+                    vulnerability_counts['unspecified'] += 1
+                vulnerability_counts['total'] += 1
+        
+        project_dict = {
+            "id": project.id,
+            "gitlab_id": project.gitlab_id,
+            "name": project.name,
+            "path": project.path,
+            "web_url": project.web_url,
+            "last_scanned": project.last_scanned,
+            "is_active": project.is_active,
+            "created_at": project.created_at,
+            "updated_at": project.updated_at,
+            "vulnerability_counts": vulnerability_counts
+        }
+        result.append(project_dict)
+    
+    return result
 
 
 @app.get("/projects/{project_id}", response_model=ProjectResponse)