From 53af2a573900a9a76e3af417317915b73894dddf Mon Sep 17 00:00:00 2001 From: JSC Date: Fri, 11 Jul 2025 18:29:31 +0200 Subject: [PATCH] Add vulnerability counts and image counts to project and file responses --- main.py | 200 +++++++++++++++++++++++++++++++++++++++++++++++++++----- 1 file changed, 183 insertions(+), 17 deletions(-) diff --git a/main.py b/main.py index defa4b9..c18c792 100644 --- a/main.py +++ b/main.py @@ -168,6 +168,8 @@ class FileResponse(BaseModel): is_active: bool created_at: datetime updated_at: datetime + image_count: Optional[int] = None + vulnerability_counts: Optional[dict] = None class ImageResponse(BaseModel): @@ -181,6 +183,7 @@ class ImageResponse(BaseModel): created_at: datetime updated_at: datetime usage_count: Optional[int] = None + vulnerability_counts: Optional[dict] = None class VulnerabilityResponse(BaseModel): @@ -320,19 +323,37 @@ async def get_projects( '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 + # Count vulnerabilities by severity for all images in this project using SQL COUNT queries + if project_images: + image_ids = [image.id for image in project_images] + vulnerability_counts = { + 'critical': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'critical' + ).count(), + 'high': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'high' + ).count(), + 'medium': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'medium' + ).count(), + 'low': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'low' + ).count(), + 'unspecified': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'unspecified' + ).count(), + } + vulnerability_counts['total'] = sum(vulnerability_counts.values()) project_dict = { "id": project.id, @@ -364,13 +385,85 @@ async def get_project_files( project_id: int, skip: int = 0, limit: int = 100, + include_image_counts: bool = False, db: Session = Depends(get_db) ): files = db.query(File).filter( File.project_id == project_id, File.is_active == True ).offset(skip).limit(limit).all() - return files + + if not include_image_counts: + return files + + # Add image and vulnerability counts for each file + result = [] + for file in files: + # Count distinct images in this file + distinct_images = db.query(Image).join(FileImageUsage).filter( + FileImageUsage.file_id == file.id, + FileImageUsage.is_active == True, + Image.is_active == True + ).distinct().all() + + # Count vulnerabilities by severity for all images in this file + vulnerability_counts = { + 'critical': 0, + 'high': 0, + 'medium': 0, + 'low': 0, + 'unspecified': 0, + 'total': 0 + } + + if distinct_images: + # Count vulnerabilities by severity for all images in this file using SQL COUNT queries + image_ids = [image.id for image in distinct_images] + vulnerability_counts = { + 'critical': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'critical' + ).count(), + 'high': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'high' + ).count(), + 'medium': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'medium' + ).count(), + 'low': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'low' + ).count(), + 'unspecified': db.query(Vulnerability).filter( + Vulnerability.image_id.in_(image_ids), + Vulnerability.is_active == True, + Vulnerability.severity == 'unspecified' + ).count(), + } + vulnerability_counts['total'] = sum(vulnerability_counts.values()) + + file_dict = { + "id": file.id, + "project_id": file.project_id, + "file_path": file.file_path, + "branch": file.branch, + "file_type": file.file_type, + "last_scanned": file.last_scanned, + "is_active": file.is_active, + "created_at": file.created_at, + "updated_at": file.updated_at, + "image_count": len(distinct_images), + "vulnerability_counts": vulnerability_counts + } + result.append(file_dict) + + return result @app.get("/projects/{project_id}/images", response_model=List[ImageResponse]) @@ -378,6 +471,7 @@ async def get_project_images( project_id: int, skip: int = 0, limit: int = 100, + include_vulnerability_counts: bool = False, db: Session = Depends(get_db) ): images = db.query(Image).join(FileImageUsage).join(File).filter( @@ -386,7 +480,7 @@ async def get_project_images( FileImageUsage.is_active == True ).distinct().offset(skip).limit(limit).all() - # Add usage count for each image + # Add usage count and optionally vulnerability counts for each image result = [] for image in images: usage_count = db.query(FileImageUsage).filter( @@ -406,6 +500,39 @@ async def get_project_images( "updated_at": image.updated_at, "usage_count": usage_count } + + if include_vulnerability_counts: + # Count vulnerabilities by severity for this image using SQL COUNT queries + vulnerability_counts = { + 'critical': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'critical' + ).count(), + 'high': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'high' + ).count(), + 'medium': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'medium' + ).count(), + 'low': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'low' + ).count(), + 'unspecified': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'unspecified' + ).count(), + } + vulnerability_counts['total'] = sum(vulnerability_counts.values()) + image_dict["vulnerability_counts"] = vulnerability_counts + result.append(image_dict) return result @@ -450,7 +577,11 @@ async def get_images( @app.get("/images/{image_id}", response_model=ImageResponse) -async def get_image(image_id: int, db: Session = Depends(get_db)): +async def get_image( + image_id: int, + include_vulnerability_counts: bool = False, + db: Session = Depends(get_db) +): image = db.query(Image).filter(Image.id == image_id).first() if not image: raise HTTPException(status_code=404, detail="Image not found") @@ -460,7 +591,7 @@ async def get_image(image_id: int, db: Session = Depends(get_db)): FileImageUsage.is_active == True ).count() - return { + result = { "id": image.id, "image_name": image.image_name, "tag": image.tag, @@ -472,6 +603,40 @@ async def get_image(image_id: int, db: Session = Depends(get_db)): "updated_at": image.updated_at, "usage_count": usage_count } + + if include_vulnerability_counts: + # Count vulnerabilities by severity for this image using SQL COUNT queries + vulnerability_counts = { + 'critical': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'critical' + ).count(), + 'high': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'high' + ).count(), + 'medium': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'medium' + ).count(), + 'low': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'low' + ).count(), + 'unspecified': db.query(Vulnerability).filter( + Vulnerability.image_id == image.id, + Vulnerability.is_active == True, + Vulnerability.severity == 'unspecified' + ).count(), + } + vulnerability_counts['total'] = sum(vulnerability_counts.values()) + result["vulnerability_counts"] = vulnerability_counts + + return result @app.get("/images/{image_id}/vulnerabilities", response_model=List[VulnerabilityResponse]) @@ -681,6 +846,7 @@ async def get_scan_job(job_id: int, db: Session = Depends(get_db)): } + @app.get("/scan/status") async def get_scan_status(db: Session = Depends(get_db)): """Check if there are any running or pending scans"""