Add vulnerability counts to project response in get_projects endpoint
This commit is contained in:
54
main.py
54
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,8 +296,60 @@ async def get_projects(
|
||||
query = query.filter(Project.is_active == True)
|
||||
|
||||
projects = query.offset(skip).limit(limit).all()
|
||||
|
||||
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)
|
||||
async def get_project(project_id: int, db: Session = Depends(get_db)):
|
||||
|
||||
Reference in New Issue
Block a user