mabac/abac/models.py

from django.contrib.auth.models import AbstractUser
from django.db import models
from django.utils import timezone

import pickle
from base64 import b64encode, b64decode
from random import randint
from phe import paillier
from uuid import uuid4


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 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.CharField(max_length=2048)


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)->(bool, int, int):
        x = randint(0,1000)
        y = randint(0,1000)
        try:
            paillier_public_key = paillier.PaillierPublicKey(n=int(user.public_key))
        except:
            return False, 0, 0
        encrypted_sum = paillier_public_key.encrypt(x)

        for rule_attribute in self.ruleattribute_set.all():
            if not rule_attribute.attribute_type.is_private:
                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, 0, 0
                if not self.attribute_satisfies_rule(user_attribute, rule_attribute):
                    return False, 0, 0
            else:
                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, 0, 0
                encrypted_attribute_value = paillier.EncryptedNumber(paillier_public_key, int(user_attribute.value))
                encrypted_sum += (encrypted_attribute_value-int(rule_attribute.value))*y
        return True, x, encrypted_sum.ciphertext()
    
    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)

class TrustedAuthority(models.Model):
    name = models.CharField(max_length=255)
    public_key = models.TextField()


class EncryptedData(models.Model):
    file = models.ForeignKey(File, on_delete=models.CASCADE)
    x_values = models.TextField()  # Store the x values as comma-separated
    token = models.UUIDField(default=uuid4, unique=True, editable=False)