old-riskletpy/backend/core/models.py

import uuid

from django.contrib.auth.models import User
from django.db import models
from localflavor.br.br_states import STATE_CHOICES
from django.contrib import admin
import yaml


class UuidModel(models.Model):
    uuid = models.UUIDField(unique=True, editable=False, default=uuid.uuid4)

    class Meta:
        abstract = True


class TimeStampedModel(models.Model):
    created = models.DateTimeField(
        'criado em',
        auto_now_add=True,
        auto_now=False
    )
    modified = models.DateTimeField(
        'modificado em',
        auto_now_add=False,
        auto_now=True
    )

    class Meta:
        abstract = True


class CreatedBy(models.Model):
    created_by = models.ForeignKey(
        User,
        verbose_name='criado por',
        on_delete=models.SET_NULL,
        null=True,
        blank=True,
    )

    class Meta:
        abstract = True


class Organization(models.Model):
    name = models.CharField(max_length=255)
    email = models.EmailField()
    employee_headcount = models.CharField(max_length=20)
    annual_revenue = models.CharField(max_length=20)
    critical_applications = models.CharField(max_length=20)
    compliance_frameworks = models.JSONField()  # Stores selected compliance frameworks as a list
    industry_sector = models.CharField(max_length=255)
    it_dependency = models.IntegerField()
    data_sensitivity = models.CharField(max_length=20)
    network_infrastructure = models.CharField(max_length=20)
    remote_workforce_percentage = models.CharField(max_length=20)
    third_party_vendor_access = models.CharField(max_length=20)
    internal_software_development = models.CharField(max_length=20)
    geographic_scope = models.CharField(max_length=20, null=True, blank=True)
    customer_base = models.CharField(max_length=20, null=True, blank=True)
    customer_type = models.CharField(max_length=20, null=True, blank=True)
    product_portfolio = models.CharField(max_length=20, null=True, blank=True)
    supplier_base = models.CharField(max_length=20, null=True, blank=True)
    it_infrastructure = models.JSONField(null=True, blank=True)  # Stores selected IT infrastructure types as a list
    intellectual_property = models.JSONField(null=True, blank=True)  # Stores selected IP protection types as a list
    sensitive_data = models.JSONField(null=True, blank=True)  # Stores selected sensitive data types as a list
    integration_level = models.CharField(max_length=20, null=True, blank=True)

    risks = models.ManyToManyField('Risk', related_name='organizations', blank=True)

    def __str__(self):
        return self.name


class DocumentSegment(models.Model):
    SEGMENT_TYPES = (
        ('title', 'Title'),
        ('subtitle', 'Subtitle'), 
        ('h1', 'Header 1'),
        ('h2', 'Header 2'),
        ('h3', 'Header 3'),
        ('body', 'Body Text'),
        ('quote', 'Quote')
    )

    document = models.ForeignKey('Document', on_delete=models.CASCADE, related_name='segments')
    segment_type = models.CharField(max_length=20, choices=SEGMENT_TYPES)
    content = models.TextField()
    order = models.PositiveIntegerField()
    modified_at = models.DateTimeField(auto_now=True)

    class Meta:
        ordering = ['order']

    def __str__(self):
        return f"{self.get_segment_type_display()} - {self.content[:50]}"


class Document(models.Model):
    id = models.UUIDField(primary_key=True, default=uuid.uuid4, editable=False)
    organization = models.ForeignKey(Organization, on_delete=models.CASCADE, related_name='documents')
    created_at = models.DateTimeField(auto_now_add=True)
    modified_at = models.DateTimeField(auto_now=True)

    def __str__(self):
        return f"Document for {self.organization.name}"

    def add_segment(self, segment_type, content, position=None):
        """
        Add a new segment at the specified position.
        If position is None, append to the end.
        """
        if position is None:
            # Get the highest order and add 1
            last_order = self.segments.aggregate(models.Max('order'))['order__max']
            new_order = 1 if last_order is None else last_order + 1
        else:
            # Move all segments at and after the position up by 1
            self.segments.filter(order__gte=position).update(order=models.F('order') + 1)
            new_order = position

        return self.segments.create(
            segment_type=segment_type,
            content=content,
            order=new_order
        )
    

class DocumentTemplate(models.Model):
    name = models.CharField(max_length=255, unique=True)
    content = models.TextField(help_text="YAML format content")
    created_at = models.DateTimeField(auto_now_add=True)
    updated_at = models.DateTimeField(auto_now=True)

    def to_dict(self):
        return yaml.safe_load(self.content)


class Risk(models.Model):
    risk_id = models.IntegerField(unique=True)
    category = models.CharField(max_length=255)
    risk_name = models.CharField(max_length=255)
    primary_impact = models.TextField()
    secondary_impact = models.TextField()
    tretiary_impact = models.TextField()
    detection_difficulty = models.CharField(max_length=255)
    recovery_complexity = models.CharField(max_length=255)
    businnes_impact_severity = models.CharField(max_length=255)

    def __str__(self):
        return f"{self.risk_id} - {self.risk_name}"


class Control(models.Model):
    risk = models.ForeignKey(Risk, on_delete=models.CASCADE, related_name="controls")
    description = models.TextField()
    safeguard = models.CharField(max_length=255)
    weight = models.IntegerField()
    
    def __str__(self):
        return f"{self.safeguard} ({self.weight})"