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 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):
    """
    Represents an information security risk.
    """

    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, "Low (< 1,000 € – minor operational impact)"),
        (2, "Medium (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(max_length=255)
    description = models.TextField(max_length=225, blank=True, null=True)
    asset = models.CharField(max_length=255, blank=True, null=True)
    process = models.CharField(max_length=255, blank=True, null=True)
    category = models.CharField(max_length=255, blank=True, null=True)
    created_at = models.DateTimeField(auto_now_add=True,)
    updated_at = models.DateTimeField(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
    """

    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.
    """
    STATUS_CHOICES = [
        ("planned", "Planned"),
        ("in_progress", "In progress"),
        ("completed", "Completed"),
        ("verified", "Verified"),
        ("rejected", "Rejected"),
    ]

    title = models.CharField(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.
    """
    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
    """
    STATUS_CHOICES = [
        ("open", "Opened"),
        ("in_progress", "In Progress"),
        ("closed", "Closed"),
    ]
    title = models.CharField(max_length=255)
    description = models.TextField(blank=True, null=True)
    date_reported = models.DateField(blank=True, null=True)
    reported_by = models.ForeignKey(settings.AUTH_USER_MODEL, 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]}..."