quaik8/gdit-back
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Enhance vulnerability scanning by linking vulnerabilities to scan jobs and updating scan_all_images method to accept scan_job_id

Browse Source
This commit is contained in:
JSC
2025-07-10 23:07:19 +02:00
parent 2c64c2c34d
commit fe2600f5bf
3 changed files with 27 additions and 16 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
main.py
View File

@@ -122,7 +122,7 @@ def background_vulnerability_scan(job_id: int):
websocket_manager.notify_scan_started_sync("vulnerability", job_id, "Vulnerability scan started") websocket_manager.notify_scan_started_sync("vulnerability", job_id, "Vulnerability scan started")
vuln_scanner = VulnerabilityScanner() vuln_scanner = VulnerabilityScanner()
vuln_scanner.scan_all_images(db) vuln_scanner.scan_all_images(db, job_id)
if job: if job:
job.status = "completed" job.status = "completed"
@@ -185,6 +185,7 @@ class ImageResponse(BaseModel):
class VulnerabilityResponse(BaseModel): class VulnerabilityResponse(BaseModel):
id: int id: int
image_id: int image_id: int
scan_job_id: Optional[int] = None
vulnerability_id: str vulnerability_id: str
severity: str severity: str
title: Optional[str] = None title: Optional[str] = None

4
models.py
View File

@@ -119,6 +119,7 @@ class Vulnerability(Base):
id = Column(Integer, primary_key=True) id = Column(Integer, primary_key=True)
image_id = Column(Integer, ForeignKey("images.id"), nullable=False) image_id = Column(Integer, ForeignKey("images.id"), nullable=False)
scan_job_id = Column(Integer, ForeignKey("scan_jobs.id"), nullable=True)
vulnerability_id = Column(String(100), nullable=False) vulnerability_id = Column(String(100), nullable=False)
severity = Column(String(20), nullable=False) severity = Column(String(20), nullable=False)
title = Column(String(500), nullable=True) title = Column(String(500), nullable=True)
@@ -132,6 +133,7 @@ class Vulnerability(Base):
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
image = relationship("Image", back_populates="vulnerabilities") image = relationship("Image", back_populates="vulnerabilities")
scan_job = relationship("ScanJob", back_populates="vulnerabilities")
__table_args__ = ( __table_args__ = (
UniqueConstraint( UniqueConstraint(
@@ -170,6 +172,8 @@ class ScanJob(Base):
created_at = Column(DateTime, default=datetime.utcnow) created_at = Column(DateTime, default=datetime.utcnow)
updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow)
vulnerabilities = relationship("Vulnerability", back_populates="scan_job", cascade="all, delete-orphan")
DATABASE_URL = "sqlite:///./gitlab_docker_tracker.db" DATABASE_URL = "sqlite:///./gitlab_docker_tracker.db"
engine = create_engine(DATABASE_URL, echo=False) engine = create_engine(DATABASE_URL, echo=False)

36
vulnerability_scanner.py
View File

@@ -23,18 +23,18 @@ class VulnerabilityScanner:
return False return False
def scan_image_vulnerabilities( def scan_image_vulnerabilities(
self, db: Session, image: Image self, db: Session, image: Image, scan_job_id: Optional[int] = None
) -> List[Vulnerability]: ) -> List[Vulnerability]:
vulnerabilities = [] vulnerabilities = []
if self.trivy_available: if self.trivy_available:
vulnerabilities.extend(self._scan_with_trivy(db, image)) vulnerabilities.extend(self._scan_with_trivy(db, image, scan_job_id))
else: else:
vulnerabilities.extend(self._scan_with_api(db, image)) vulnerabilities.extend(self._scan_with_api(db, image, scan_job_id))
return vulnerabilities return vulnerabilities
def _scan_with_trivy(self, db: Session, image: Image) -> List[Vulnerability]: def _scan_with_trivy(self, db: Session, image: Image, scan_job_id: Optional[int] = None) -> List[Vulnerability]:
vulnerabilities = [] vulnerabilities = []
try: try:
@@ -54,7 +54,7 @@ class VulnerabilityScanner:
if result.returncode == 0: if result.returncode == 0:
data = json.loads(result.stdout) data = json.loads(result.stdout)
vulnerabilities = self._parse_trivy_results(db, image, data) vulnerabilities = self._parse_trivy_results(db, image, data, scan_job_id)
else: else:
print(f"Trivy scan failed for {image.full_image_name}: {result.stderr}") print(f"Trivy scan failed for {image.full_image_name}: {result.stderr}")
@@ -64,7 +64,7 @@ class VulnerabilityScanner:
return vulnerabilities return vulnerabilities
def _parse_trivy_results( def _parse_trivy_results(
self, db: Session, image: Image, data: Dict self, db: Session, image: Image, data: Dict, scan_job_id: Optional[int] = None
) -> List[Vulnerability]: ) -> List[Vulnerability]:
vulnerabilities = [] vulnerabilities = []
@@ -72,14 +72,14 @@ class VulnerabilityScanner:
for result in results: for result in results:
vulns = result.get("Vulnerabilities", []) vulns = result.get("Vulnerabilities", [])
for vuln in vulns: for vuln in vulns:
vulnerability = self._create_vulnerability_from_trivy(db, image, vuln) vulnerability = self._create_vulnerability_from_trivy(db, image, vuln, scan_job_id)
if vulnerability: if vulnerability:
vulnerabilities.append(vulnerability) vulnerabilities.append(vulnerability)
return vulnerabilities return vulnerabilities
def _create_vulnerability_from_trivy( def _create_vulnerability_from_trivy(
self, db: Session, image: Image, vuln_data: Dict self, db: Session, image: Image, vuln_data: Dict, scan_job_id: Optional[int] = None
) -> Optional[Vulnerability]: ) -> Optional[Vulnerability]:
vulnerability_id = vuln_data.get("VulnerabilityID") vulnerability_id = vuln_data.get("VulnerabilityID")
if not vulnerability_id: if not vulnerability_id:
@@ -97,6 +97,8 @@ class VulnerabilityScanner:
if existing: if existing:
existing.is_active = True existing.is_active = True
existing.scan_date = datetime.utcnow() existing.scan_date = datetime.utcnow()
if scan_job_id:
existing.scan_job_id = scan_job_id
db.commit() db.commit()
return existing return existing
@@ -128,6 +130,7 @@ class VulnerabilityScanner:
vulnerability = Vulnerability( vulnerability = Vulnerability(
image_id=image.id, image_id=image.id,
scan_job_id=scan_job_id,
vulnerability_id=vulnerability_id, vulnerability_id=vulnerability_id,
severity=severity, severity=severity,
title=title, title=title,
@@ -144,7 +147,7 @@ class VulnerabilityScanner:
return vulnerability return vulnerability
def _scan_with_api(self, db: Session, image: Image) -> List[Vulnerability]: def _scan_with_api(self, db: Session, image: Image, scan_job_id: Optional[int] = None) -> List[Vulnerability]:
vulnerabilities = [] vulnerabilities = []
try: try:
@@ -159,7 +162,7 @@ class VulnerabilityScanner:
if response.status_code == 200: if response.status_code == 200:
data = response.json() data = response.json()
vulnerabilities = self._parse_api_results(db, image, data) vulnerabilities = self._parse_api_results(db, image, data, scan_job_id)
else: else:
print( print(
f"API scan failed for {image.full_image_name}: {response.status_code}" f"API scan failed for {image.full_image_name}: {response.status_code}"
@@ -171,21 +174,21 @@ class VulnerabilityScanner:
return vulnerabilities return vulnerabilities
def _parse_api_results( def _parse_api_results(
self, db: Session, image: Image, data: Dict self, db: Session, image: Image, data: Dict, scan_job_id: Optional[int] = None
) -> List[Vulnerability]: ) -> List[Vulnerability]:
vulnerabilities = [] vulnerabilities = []
cves = data.get("vulnerabilities", []) cves = data.get("vulnerabilities", [])
for cve in cves: for cve in cves:
cve_data = cve.get("cve", {}) cve_data = cve.get("cve", {})
vulnerability = self._create_vulnerability_from_api(db, image, cve_data) vulnerability = self._create_vulnerability_from_api(db, image, cve_data, scan_job_id)
if vulnerability: if vulnerability:
vulnerabilities.append(vulnerability) vulnerabilities.append(vulnerability)
return vulnerabilities return vulnerabilities
def _create_vulnerability_from_api( def _create_vulnerability_from_api(
self, db: Session, image: Image, cve_data: Dict self, db: Session, image: Image, cve_data: Dict, scan_job_id: Optional[int] = None
) -> Optional[Vulnerability]: ) -> Optional[Vulnerability]:
vulnerability_id = cve_data.get("id") vulnerability_id = cve_data.get("id")
if not vulnerability_id: if not vulnerability_id:
@@ -203,6 +206,8 @@ class VulnerabilityScanner:
if existing: if existing:
existing.is_active = True existing.is_active = True
existing.scan_date = datetime.utcnow() existing.scan_date = datetime.utcnow()
if scan_job_id:
existing.scan_job_id = scan_job_id
db.commit() db.commit()
return existing return existing
@@ -243,6 +248,7 @@ class VulnerabilityScanner:
vulnerability = Vulnerability( vulnerability = Vulnerability(
image_id=image.id, image_id=image.id,
scan_job_id=scan_job_id,
vulnerability_id=vulnerability_id, vulnerability_id=vulnerability_id,
severity=severity, severity=severity,
title=title, title=title,
@@ -280,13 +286,13 @@ class VulnerabilityScanner:
else: else:
return "unspecified" return "unspecified"
def scan_all_images(self, db: Session) -> None: def scan_all_images(self, db: Session, scan_job_id: Optional[int] = None) -> None:
images = db.query(Image).filter(Image.is_active == True).all() images = db.query(Image).filter(Image.is_active == True).all()
for image in images: for image in images:
try: try:
print(f"Scanning vulnerabilities for {image.full_image_name}") print(f"Scanning vulnerabilities for {image.full_image_name}")
vulnerabilities = self.scan_image_vulnerabilities(db, image) vulnerabilities = self.scan_image_vulnerabilities(db, image, scan_job_id)
print(f"Found {len(vulnerabilities)} vulnerabilities") print(f"Found {len(vulnerabilities)} vulnerabilities")
except Exception as e: except Exception as e:
print(f"Error scanning {image.full_image_name}: {e}") print(f"Error scanning {image.full_image_name}: {e}")