NeuroSploit/backend/models/vulnerability.py

"""
NeuroSploit v3 - Vulnerability Models
"""
from datetime import datetime
from typing import Optional, List
from sqlalchemy import String, Integer, Float, Boolean, DateTime, Text, JSON, ForeignKey
from sqlalchemy.orm import Mapped, mapped_column, relationship
from backend.db.database import Base
import uuid


class VulnerabilityTest(Base):
    """Individual vulnerability test record"""
    __tablename__ = "vulnerability_tests"

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id: Mapped[str] = mapped_column(String(36), ForeignKey("scans.id", ondelete="CASCADE"))
    endpoint_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("endpoints.id", ondelete="SET NULL"), nullable=True)

    # Test details
    vulnerability_type: Mapped[str] = mapped_column(String(100))  # xss_reflected, sqli_union, etc.
    payload: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Request/Response
    request_data: Mapped[dict] = mapped_column(JSON, default=dict)
    response_data: Mapped[dict] = mapped_column(JSON, default=dict)

    # Result
    is_vulnerable: Mapped[bool] = mapped_column(Boolean, default=False)
    confidence: Mapped[Optional[float]] = mapped_column(Float, nullable=True)  # 0.0 to 1.0
    evidence: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Timestamps
    tested_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)

    def to_dict(self) -> dict:
        """Convert to dictionary"""
        return {
            "id": self.id,
            "scan_id": self.scan_id,
            "endpoint_id": self.endpoint_id,
            "vulnerability_type": self.vulnerability_type,
            "payload": self.payload,
            "request_data": self.request_data,
            "response_data": self.response_data,
            "is_vulnerable": self.is_vulnerable,
            "confidence": self.confidence,
            "evidence": self.evidence,
            "tested_at": self.tested_at.isoformat() if self.tested_at else None
        }


class Vulnerability(Base):
    """Confirmed vulnerability model"""
    __tablename__ = "vulnerabilities"

    id: Mapped[str] = mapped_column(String(36), primary_key=True, default=lambda: str(uuid.uuid4()))
    scan_id: Mapped[str] = mapped_column(String(36), ForeignKey("scans.id", ondelete="CASCADE"))
    test_id: Mapped[Optional[str]] = mapped_column(String(36), ForeignKey("vulnerability_tests.id", ondelete="SET NULL"), nullable=True)

    # Vulnerability details
    title: Mapped[str] = mapped_column(String(500))
    vulnerability_type: Mapped[str] = mapped_column(String(100))
    severity: Mapped[str] = mapped_column(String(20))  # critical, high, medium, low, info

    # Scoring
    cvss_score: Mapped[Optional[float]] = mapped_column(Float, nullable=True)
    cvss_vector: Mapped[Optional[str]] = mapped_column(String(100), nullable=True)
    cwe_id: Mapped[Optional[str]] = mapped_column(String(50), nullable=True)

    # Details
    description: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    affected_endpoint: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Proof of Concept
    poc_request: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    poc_response: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    poc_payload: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    poc_parameter: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)  # Vulnerable parameter
    poc_evidence: Mapped[Optional[str]] = mapped_column(Text, nullable=True)  # Evidence of vulnerability

    # Remediation
    impact: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    remediation: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    references: Mapped[List] = mapped_column(JSON, default=list)

    # AI Analysis
    ai_analysis: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # PoC Code (executable proof-of-concept: HTML, Python, curl, etc.)
    poc_code: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Screenshots (list of base64 data URIs or filesystem paths)
    screenshots: Mapped[List] = mapped_column(JSON, default=list)

    # Source URL and parameter (for finding_id reconstruction)
    url: Mapped[Optional[str]] = mapped_column(Text, nullable=True)
    parameter: Mapped[Optional[str]] = mapped_column(String(500), nullable=True)

    # Validation status (manual review workflow)
    validation_status: Mapped[str] = mapped_column(String(20), default="ai_confirmed")
    # Values: "ai_confirmed" | "ai_rejected" | "validated" | "false_positive" | "pending_review"
    ai_rejection_reason: Mapped[Optional[str]] = mapped_column(Text, nullable=True)

    # Timestamps
    created_at: Mapped[datetime] = mapped_column(DateTime, default=datetime.utcnow)

    # Relationships
    scan: Mapped["Scan"] = relationship("Scan", back_populates="vulnerabilities")

    def to_dict(self) -> dict:
        """Convert to dictionary"""
        return {
            "id": self.id,
            "scan_id": self.scan_id,
            "test_id": self.test_id,
            "title": self.title,
            "vulnerability_type": self.vulnerability_type,
            "severity": self.severity,
            "cvss_score": self.cvss_score,
            "cvss_vector": self.cvss_vector,
            "cwe_id": self.cwe_id,
            "description": self.description,
            "affected_endpoint": self.affected_endpoint,
            "poc_request": self.poc_request,
            "poc_response": self.poc_response,
            "poc_payload": self.poc_payload,
            "poc_parameter": self.poc_parameter,
            "poc_evidence": self.poc_evidence,
            "impact": self.impact,
            "remediation": self.remediation,
            "references": self.references,
            "ai_analysis": self.ai_analysis,
            "poc_code": self.poc_code,
            "screenshots": self.screenshots or [],
            "url": self.url,
            "parameter": self.parameter,
            "validation_status": self.validation_status or "ai_confirmed",
            "ai_rejection_reason": self.ai_rejection_reason,
            "created_at": self.created_at.isoformat() if self.created_at else None
        }