from datetime import datetime from enum import Enum from typing import Optional from sqlalchemy import ( Boolean, Column, DateTime, ForeignKey, Integer, String, Text, UniqueConstraint, create_engine, ) from sqlalchemy.ext.declarative import declarative_base from sqlalchemy.orm import relationship, sessionmaker Base = declarative_base() class VulnerabilitySeverity(Enum): CRITICAL = "critical" HIGH = "high" MEDIUM = "medium" LOW = "low" UNSPECIFIED = "unspecified" class IgnoreType(Enum): IMAGE = "image" FILE = "file" PROJECT = "project" class Project(Base): __tablename__ = "projects" id = Column(Integer, primary_key=True) gitlab_id = Column(Integer, unique=True, nullable=False) name = Column(String(255), nullable=False) path = Column(String(255), nullable=False) web_url = Column(String(500), nullable=False) last_scanned = Column(DateTime, nullable=True) is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) files = relationship("File", back_populates="project", cascade="all, delete-orphan") ignore_rules = relationship( "IgnoreRule", back_populates="project", cascade="all, delete-orphan" ) class File(Base): __tablename__ = "files" id = Column(Integer, primary_key=True) project_id = Column(Integer, ForeignKey("projects.id"), nullable=False) file_path = Column(String(500), nullable=False) branch = Column(String(100), nullable=False) file_type = Column(String(50), nullable=False) # docker-compose, dockerfile, gitlab-ci last_scanned = Column(DateTime, nullable=True) is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) project = relationship("Project", back_populates="files") image_usages = relationship("FileImageUsage", back_populates="file", cascade="all, delete-orphan") __table_args__ = ( UniqueConstraint("project_id", "file_path", "branch", name="unique_file_branch"), ) class Image(Base): __tablename__ = "images" id = Column(Integer, primary_key=True) image_name = Column(String(500), nullable=False) tag = Column(String(100), nullable=True) registry = Column(String(255), nullable=True) full_image_name = Column(String(1000), nullable=False, unique=True) last_seen = Column(DateTime, default=datetime.utcnow) is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) vulnerabilities = relationship( "Vulnerability", back_populates="image", cascade="all, delete-orphan" ) file_usages = relationship( "FileImageUsage", back_populates="image", cascade="all, delete-orphan" ) class FileImageUsage(Base): __tablename__ = "file_image_usages" id = Column(Integer, primary_key=True) file_id = Column(Integer, ForeignKey("files.id"), nullable=False) image_id = Column(Integer, ForeignKey("images.id"), nullable=False) first_seen = Column(DateTime, default=datetime.utcnow) last_seen = Column(DateTime, default=datetime.utcnow) is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) file = relationship("File", back_populates="image_usages") image = relationship("Image", back_populates="file_usages") __table_args__ = ( UniqueConstraint("file_id", "image_id", name="unique_file_image"), ) class Vulnerability(Base): __tablename__ = "vulnerabilities" id = Column(Integer, primary_key=True) 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) severity = Column(String(20), nullable=False) title = Column(String(500), nullable=True) description = Column(Text, nullable=True) cvss_score = Column(String(10), nullable=True) published_date = Column(DateTime, nullable=True) fixed_version = Column(String(100), nullable=True) scan_date = Column(DateTime, default=datetime.utcnow) is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) image = relationship("Image", back_populates="vulnerabilities") scan_job = relationship("ScanJob", back_populates="vulnerabilities") __table_args__ = ( UniqueConstraint( "image_id", "vulnerability_id", name="unique_image_vulnerability" ), ) class IgnoreRule(Base): __tablename__ = "ignore_rules" id = Column(Integer, primary_key=True) project_id = Column(Integer, ForeignKey("projects.id"), nullable=True) ignore_type = Column(String(20), nullable=False) target = Column(String(1000), nullable=False) reason = Column(Text, nullable=True) created_by = Column(String(100), nullable=True) is_active = Column(Boolean, default=True) created_at = Column(DateTime, default=datetime.utcnow) updated_at = Column(DateTime, default=datetime.utcnow, onupdate=datetime.utcnow) project = relationship("Project", back_populates="ignore_rules") class ScanJob(Base): __tablename__ = "scan_jobs" id = Column(Integer, primary_key=True) job_type = Column(String(50), nullable=False) # discovery, vulnerability status = Column(String(20), nullable=False) # pending, running, completed, failed project_id = Column(Integer, ForeignKey("projects.id"), nullable=True) started_at = Column(DateTime, nullable=True) completed_at = Column(DateTime, nullable=True) error_message = Column(Text, nullable=True) results_summary = Column(Text, nullable=True) created_at = Column(DateTime, default=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" engine = create_engine(DATABASE_URL, echo=False) SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine) def create_tables(): Base.metadata.create_all(bind=engine) def get_db(): db = SessionLocal() try: yield db finally: db.close()