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, unique=True, help_text="What is the name of your organization?")
    email = models.EmailField(unique=True, help_text="What is your email?")
    employee_headcount = models.CharField(max_length=20, help_text="What is your organization's current employee headcount?")
    annual_revenue = models.CharField(max_length=20, help_text="What is your organization's annual revenue range?")
    critical_applications = models.CharField(max_length=20, help_text="How many critical business applications do your employees use daily?")
    compliance_frameworks = models.JSONField(help_text="Which regulatory frameworks is your organization required to comply with?")  # Stores selected compliance frameworks as a list
    industry_sector = models.CharField(max_length=255,help_text="What is your primary industry sector?")
    it_dependency = models.CharField(max_length=255, help_text="On a scale from 1-10, how dependent is your business operations on technology?")
    remote_workforce_percentage = models.CharField(max_length=20, help_text="What percentage of your workforce operates remotely?")
    third_party_vendor_access = models.CharField(max_length=20, help_text="How many third-party vendors have access to your systems?")
    internal_software_development = models.CharField(max_length=20, help_text="What is the extent of your internal software development activities?")
    geographic_scope = models.CharField(max_length=20, null=True, blank=True, help_text="What is your organization's geographic operational scope?")
    customer_base = models.CharField(max_length=255, null=True, blank=True, help_text="How would you characterize your customer base distribution?")
    customer_type = models.CharField(max_length=20, null=True, blank=True, help_text="What is your primary customer type?")
    product_portfolio = models.CharField(max_length=20, null=True, blank=True, help_text="How diversified is your product/service portfolio?")
    supplier_base = models.CharField(max_length=20, null=True, blank=True, help_text="What is your supplier base structure?")
    it_infrastructure = models.JSONField(null=True, blank=True, help_text="What is your primary IT infrastructure model?")  # Stores selected IT infrastructure types as a list
    integration_level = models.CharField(max_length=20, null=True, blank=True, help_text="How integrated are your critical business systems?")
    network_infrastructure = models.CharField(max_length=20, null=True, blank=True, help_text="What best describes your organization's network infrastructure model?")
    change_rate = models.CharField(max_length=20, null=True, blank=True, help_text="How frequently does your organization undergo significant technology or business changes?")
    threat_actors = models.JSONField(null=True, blank=True, help_text="Which types of threat actors are most relevant to your organization (e.g., cybercriminals, insiders, nation-states)?")
    sensitive_data_types = models.JSONField(null=True, blank=True, help_text="Stores applicable status and business impact rating (1-5) for each sensitive data type. Example: {'personal': {'applicable': True, 'impact': 4}, ...}")
    risks = models.ManyToManyField('Risk', related_name='organizations', blank=True)
    expert_analysis = models.BooleanField(null=True, 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)

    STATUS_WAITING = 'waiting'
    STATUS_DONE = 'done'
    STATUS_INCOMPLETE = 'incomplete'
    STATUS_CHOICES = (
        (STATUS_WAITING, 'Waiting'),
        (STATUS_DONE, 'Done'),
        (STATUS_INCOMPLETE, 'Incomplete'),
    )
    status = models.CharField(max_length=16, choices=STATUS_CHOICES, default=STATUS_WAITING)
    risk_explanations = models.JSONField(default=dict, blank=True, help_text="Map of risk_id to explanation for each risk in this document.")

    key_findings = models.TextField(blank=True, null=True, help_text="Key findings")
    recomendations = models.TextField(blank=True, null=True, help_text="Recommendations")
    
    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):
    id = models.AutoField(primary_key=True)

    subcategory = models.CharField(max_length=64, unique=True, null=False, blank=False, help_text='NIST Subcategory')
    function = models.TextField(null=True, blank=True, help_text="Function")
    category = models.TextField(null=True, blank=True,help_text="Category")
    implementation_examples = models.TextField(null=True, blank=True, help_text="Implementation Examples")
    effectiveness_monitoring_examples = models.TextField(null=True, blank=True, help_text="Effectiveness Monitoring Examples")
    documentation_score = models.IntegerField(null=True, blank=True, help_text="Documentation Score")
    implementation_score = models.IntegerField(null=True, blank=True, help_text="Implementation Score")

    def __str__(self):
        return f"{self.id} ({self.subcategory} - {self.function or self.category or ''})".rstrip(" -() ")
    
class DocumentRiskControl(models.Model):
    document = models.ForeignKey(Document, on_delete=models.CASCADE)
    risk = models.ForeignKey(Risk, on_delete=models.CASCADE)
    control = models.ForeignKey(Control, on_delete=models.CASCADE)
    weight = models.IntegerField()
    likelihood = models.IntegerField(null=True, blank=True)

    class Meta:
        unique_together = ('document', 'risk', 'control')


class DemoCode(models.Model):
    code = models.CharField(max_length=10, unique=True)
    created_at = models.DateTimeField(auto_now_add=True)
    used = models.BooleanField(default=False)
    used_at = models.DateTimeField(null=True, blank=True)
    company = models.ForeignKey(Organization, null=True, blank=True, on_delete=models.SET_NULL)

    def __str__(self):
        return (f"{self.code} - {'Used' if self.used else 'Available'}")

class PaymentRecord(models.Model):
    company = models.ForeignKey(Organization, on_delete=models.CASCADE)
    amount = models.DecimalField(max_digits=10, decimal_places=2)
    currency = models.CharField(max_length=10)
    payment_date = models.DateTimeField(auto_now_add=True)
    transaction_id = models.CharField(max_length=255, unique=True)

    def __str__(self):
        return f"{self.company} - {self.amount} {self.currency}"