gdit-back/models.py

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()