ISO-27001-Risk-Management/risks/models.py

from django.conf import settings
from django.contrib.auth.models import AbstractUser
from django.core.serializers.json import DjangoJSONEncoder
from django.db import models
from django.utils.translation import gettext_lazy as _
from multiselectfield import MultiSelectField
import datetime
import json

class SafeJSONEncoder(DjangoJSONEncoder):
    def default(self, obj):
        if isinstance(obj, datetime.date):
            return obj.isoformat()
        return super().default(obj)

class User(AbstractUser):
    """
    Custom user model to support both local and SSO users.
    """
    is_sso_user = models.BooleanField(default=False)

    @property
    def risks_owned(self):
        """ All risks where the user is the risk owner. """
        return self.owned_risks.all()
    
    @property
    def controls_responsible(self):
        """ All controls where the user is responsible. """
        return self.responsible_controls.all()

class Risk(models.Model):

    class Meta:
        verbose_name = _("Risk")
        verbose_name_plural = _("Risks")

    LIKELIHOOD_CHOICES = [
        (1, _("Very low – occurs less than once every 5 years")),
        (2, _("Low – once every 1–5 years")),
        (3, _("Likely – once per year or more")),
        (4, _("Very likely – multiple times per year/monthly")),
    ]
    IMPACT_CHOICES = [
        (1, _("Very Low (< 1,000 € – minor operational impact)")),
        (2, _("Low (1,000–5,000 € – local impact)")),
        (3, _("High (5,000–15,000 € – team-level impact)")),
        (4, _("Severe (50,000–100,000 € – regional impact)")),
        (5, _("Critical (> 100,000 € – existential threat)")),
    ]
    CIA_CHOICES = [
        ("1", _("Confidentiality")),
        ("2", _("Integrity")),
        ("3", _("Availability")),
    ]

    # Basic information
    title = models.CharField(_("Title"), max_length=255)
    description = models.TextField(_("Description"), max_length=225, blank=True, null=True)
    asset = models.CharField(_("Asset"), max_length=255, blank=True, null=True)
    process = models.CharField(_("Process"), max_length=255, blank=True, null=True)
    category = models.CharField(_("Category"), max_length=255, blank=True, null=True)
    created_at = models.DateTimeField(_("Created at"), auto_now_add=True)
    updated_at = models.DateTimeField(_("Updated at"), auto_now=True)

    # CIA Protection Goals
    cia = MultiSelectField(choices=CIA_CHOICES, max_length=100, blank=True, null=True)

    # Risk evaluation before controls
    likelihood = models.IntegerField(
        choices=LIKELIHOOD_CHOICES,
        default=1
    )
    impact = models.IntegerField(
        choices=IMPACT_CHOICES,
        default=1
    )

    # Calculated fields
    score = models.IntegerField(editable=False)
    level = models.CharField(max_length=50, editable=False)

    # Ownership
    owner = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        on_delete=models.SET_NULL,
        null=True,
        related_name="owned_risks"
    )

    # Reminder / follow-up date
    follow_up = models.DateField(blank=True, null=True)

    def save(self, *args, **kwargs):
        # Calculate risk score
        self.score = self.likelihood * self.impact

        # Determine level based on score
        if self.score <= 4:
            self.level = "Low"
        elif self.score <= 8:
            self.level = "Medium"
        elif self.score <= 12:
            self.level = "High"
        else:
            self.level = "Critical"

        super().save(*args, **kwargs)

    def __str__(self):
        return f"{self.title} (Score: {self.score}, Level: {self.level})"

class ResidualRisk(models.Model):
    """
    Residual Risk after implementing controls
    """

    class Meta:
        verbose_name = _("Residual Risk")
        verbose_name_plural = _("Residual Risks")

    risk = models.OneToOneField(
        Risk,
        on_delete=models.CASCADE,
        related_name="residual_risk")
    
    likelihood = models.IntegerField(
        choices=Risk.LIKELIHOOD_CHOICES,
        default=1
    )

    impact = models.IntegerField(
        choices=Risk.IMPACT_CHOICES,
        default=1
    )

    score = models.IntegerField(editable=False)

    level = models.CharField(max_length=50, editable=False)

    review_required = models.BooleanField(default=False)

    created_at = models.DateTimeField(auto_now_add=True,)
    updated_at = models.DateTimeField(auto_now=True)

    def save(self, *args, **kwargs):
        # Load previous state (if it exists)
        if self.pk:
            old = ResidualRisk.objects.get(pk=self.pk)
            if old.likelihood != self.likelihood or old.impact != self.impact:
                self.review_required = True

        self.score = self.likelihood * self.impact

        # Determine level based on score
        if self.score <= 4:
            self.level = "Low"
        elif self.score <= 8:
            self.level = "Medium"
        elif self.score <= 12:
            self.level = "High"
        else:
            self.level = "Critical"
        
        super().save(*args, **kwargs)

    def __str__(self):
        return f"Residual Risk for {self.risk.title} (Score: {self.score}, Level: {self.level})"

class Control(models.Model):
    """
    A security control/measure linked to a risk.
    """
    class Meta:
        verbose_name = _("Control")
        verbose_name_plural = _("Controls")

    STATUS_CHOICES = [
        ("planned", _("Planned")),
        ("in_progress", _("In progress")),
        ("completed", _("Completed")),
        ("verified", _("Verified")),
        ("rejected", _("Rejected")),
    ]

    title = models.CharField(_("Title"), max_length=255)
    status = models.CharField(max_length=20, choices=STATUS_CHOICES, default="planned")
    due_date = models.DateField(blank=True, null=True)
    responsible = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        on_delete=models.SET_NULL,
        null=True,
        related_name="responsible_controls"
    )
    description = models.TextField(blank=True, null=True)
    wiki_link = models.URLField(blank=True, null=True)
    created_at = models.DateTimeField(auto_now_add=True,)
    updated_at = models.DateTimeField(auto_now=True)

    # Relation to risk
    risks = models.ManyToManyField("Risk", related_name="controls")

    def __str__(self):
        return f"{self.title} ({self.get_status_display()})"

class AuditLog(models.Model):
    """
    Generic audit log entry for tracking changes.
    """

    class Meta:
        verbose_name = _("Auditlog")
        verbose_name_plural = _("Auditlogs")

    ACTION_CHOICES = [
        ("create", "Created"),
        ("update", "Updated"),
        ("delete", "Deleted"),
    ]

    user = models.ForeignKey(
        settings.AUTH_USER_MODEL,
        null=True, blank=True,
        on_delete=models.SET_NULL,
        related_name="audit_logs"
    )

    action = models.CharField(max_length=10, choices=ACTION_CHOICES)
    model = models.CharField(max_length=100)
    object_id = models.CharField(max_length=50)
    changes = models.JSONField(null=True, blank=True, encoder=SafeJSONEncoder)
    timestamp = models.DateTimeField(auto_now_add=True)

    def __str__(self):
        return f"[{self.timestamp}] {self.user} {self.action} {self.model}({self.object_id})"
    
class Incident(models.Model):
    """
    Incidents and related risks
    """

    class Meta:
        verbose_name = _("Incident")
        verbose_name_plural = _("Incidents")

    STATUS_CHOICES = [
        ("open", _("Opened")),
        ("in_progress", _("In Progress")),
        ("closed", _("Closed")),
    ]
    title = models.CharField(_("Title"), max_length=255)
    description = models.TextField(_("Description"), blank=True, null=True)
    date_reported = models.DateField(_("Date reported"), blank=True, null=True)
    reported_by = models.ForeignKey(
        settings.AUTH_USER_MODEL, verbose_name=_("Reported by"),
        null=True, blank=True, on_delete=models.SET_NULL, related_name="incidents"
    )
    status = models.CharField(max_length=12, choices=STATUS_CHOICES)
    related_risks = models.ManyToManyField("Risk", blank=True, related_name="incidents")
    created_at = models.DateTimeField(auto_now_add=True,)
    updated_at = models.DateTimeField(auto_now=True)

class Notification(models.Model):
    user = models.ForeignKey(settings.AUTH_USER_MODEL, on_delete=models.SET_NULL, null=True, blank=True, related_name="notifications")

    message = models.TextField()
    #related_objects =
    created_at = models.DateTimeField(auto_now_add=True)
    read = models.BooleanField(default=False) # Read in WebApp
    sent = models.BooleanField(default=False) # Sent via Mail (optional)

    def __str__(self):
        user_display = self.user.username if self.user else "System"
        return f"{user_display}: {self.message[:50]}..."