Add vulnerability counts to project response in get_projects endpoint
This commit is contained in:
56
main.py
56
main.py
@@ -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,7 +296,59 @@ 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()
|
||||||
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)
|
@app.get("/projects/{project_id}", response_model=ProjectResponse)
|
||||||
|
|||||||
Reference in New Issue
Block a user