Add vulnerability counts to project response in get_projects endpoint

This commit is contained in:
JSC
2025-07-11 15:10:36 +02:00
parent fe2600f5bf
commit e521af2281

54
main.py
View File

@@ -155,6 +155,7 @@ class ProjectResponse(BaseModel):
is_active: bool is_active: bool
created_at: datetime created_at: datetime
updated_at: datetime updated_at: datetime
vulnerability_counts: Optional[dict] = None
class FileResponse(BaseModel): class FileResponse(BaseModel):
@@ -287,6 +288,7 @@ async def get_projects(
skip: int = 0, skip: int = 0,
limit: int = 100, limit: int = 100,
active_only: bool = True, active_only: bool = True,
include_vulnerability_counts: bool = False,
db: Session = Depends(get_db) db: Session = Depends(get_db)
): ):
query = db.query(Project) query = db.query(Project)
@@ -294,8 +296,60 @@ async def get_projects(
query = query.filter(Project.is_active == True) query = query.filter(Project.is_active == True)
projects = query.offset(skip).limit(limit).all() projects = query.offset(skip).limit(limit).all()
if not include_vulnerability_counts:
return projects 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) @app.get("/projects/{project_id}", response_model=ProjectResponse)
async def get_project(project_id: int, db: Session = Depends(get_db)): async def get_project(project_id: int, db: Session = Depends(get_db)):