mabac/abac/models.py

import pickle
from base64 import b64encode, b64decode
from django.contrib.auth.models import AbstractUser
from django.db import models
from django.utils import timezone


class User(AbstractUser):
    pass
    public_key = models.TextField()
    private_key = models.TextField()

    def save(self, *args, **kwargs):
        if hasattr(self, '_phe_public_key') and hasattr(self, '_phe_private_key'):
            self.phe_public_key = b64encode(pickle.dumps(self._phe_public_key)).decode('utf-8')
            self.phe_private_key = b64encode(pickle.dumps(self._phe_private_key)).decode('utf-8')
        
        super().save(*args, **kwargs)

    @property
    def deserialized_public_key(self):
        return pickle.loads(b64decode(self.phe_public_key))
        
    @property
    def deserialized_private_key(self):
        return pickle.loads(b64decode(self.phe_private_key))
    
    class Meta:
        permissions = [
            ("can_create_users", "Can create new users"),
        ]


class AttributeType(models.Model):
    DATATYPE_CHOICES = [
        ('string', 'String'),
        ('boolean', 'Boolean'),
        ('integer', 'Integer'),
    ]

    is_private = models.BooleanField(default=False)
    datatype = models.CharField(max_length=15)
    significant_digits = models.PositiveIntegerField(null=True, blank=True)
    name = models.CharField(max_length=40)

    def save(self, *args, **kwargs):
        if self.datatype.startswith('float'):
            if self.significant_digits is None:
                raise ValueError('significant_digits must be set for float datatype')
            self.datatype = f'float_{self.significant_digits}'
        elif self.significant_digits is not None:
            raise ValueError('significant_digits must be None for non-float datatype')
        super().save(*args, **kwargs)

    def __str__(self) -> str:
        return self.name

    
class Attribute(models.Model):
    
    attribute_type = models.ForeignKey(AttributeType, on_delete=models.CASCADE)
    value = models.IntegerField()


class UserAttribute(Attribute):
    user = models.ForeignKey(User, on_delete=models.CASCADE)
    last_modified = models.DateTimeField(auto_now=True)

class File(models.Model):
    owner = models.ForeignKey(User, on_delete=models.CASCADE)
    name = models.CharField(max_length=255)
    file = models.FileField(upload_to='uploads/')
    created_at = models.DateTimeField(auto_now_add=True)
    last_modified = models.DateTimeField(auto_now=True)
    
    def __str__(self):
        return self.name

    class Meta:
        ordering = ['name', 'created_at'] 

class Rule(models.Model):
    name = models.CharField(max_length=255)
    file = models.ForeignKey(File, on_delete=models.CASCADE)
    read = models.BooleanField(default=False)
    write = models.BooleanField(default=False)
    delete = models.BooleanField(default=False)

    def is_satisfied_by(self, user):
        for rule_attribute in self.ruleattribute_set.all():
            try:
                user_attribute = UserAttribute.objects.get(
                    user=user,
                    attribute_type=rule_attribute.attribute_type,
                )
            except UserAttribute.DoesNotExist:
                # The user does not have this attribute; the rule is not satisfied.
                return False

            if not self.attribute_satisfies_rule(user_attribute, rule_attribute):
                return False
        
        return True
    
    def attribute_satisfies_rule(self, user_attribute, rule_attribute):
        operator = rule_attribute.operator
        if operator == 'EQ':
            return user_attribute.value == rule_attribute.value
        elif operator == 'NEQ':
            return not (user_attribute.value == rule_attribute.value)
        elif operator == 'GT':
            return user_attribute.value > rule_attribute.value
        elif operator == 'LT':
            return user_attribute.value < rule_attribute.value
        return False

class RuleAttribute(Attribute):
    OPTION_EQUALS = 'EQ'
    OPTION_DOES_NOT_EQUAL = 'NEQ'
    OPTION_GREATER_THAN = 'GT'
    OPTION_LESS_THAN = 'LT'

    OPTION_CHOICES = [
        ('EQ', 'Equals'),
        ('NEQ', "Doesn't Equal"),
        ('GT', 'Greater Than'),
        ('LT', 'Less Than'),
    ]

    rule = models.ForeignKey(Rule, on_delete=models.CASCADE)
    operator = models.CharField(max_length=3, choices=OPTION_CHOICES, default=OPTION_EQUALS)