import logging import stripe from django.contrib.auth.models import AbstractUser from django.db import models from yaml import load from yaml.parser import ParserError from yaml.scanner import ScannerError from core.lib.customers import get_or_create, update_customer_fields try: from yaml import CLoader as Loader except ImportError: from yaml import Loader logger = logging.getLogger(__name__) class Plan(models.Model): name = models.CharField(max_length=255, unique=True) description = models.CharField(max_length=1024, null=True, blank=True) cost = models.IntegerField() product_id = models.CharField(max_length=255, unique=True, null=True, blank=True) image = models.CharField(max_length=1024, null=True, blank=True) def __str__(self): return f"{self.name} (£{self.cost})" class User(AbstractUser): # Stripe customer ID stripe_id = models.CharField(max_length=255, null=True, blank=True) last_payment = models.DateTimeField(null=True, blank=True) plans = models.ManyToManyField(Plan, blank=True) email = models.EmailField(unique=True) def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._original = self def save(self, *args, **kwargs): """ Override the save function to create a Stripe customer. """ if not self.stripe_id: # stripe ID not stored self.stripe_id = get_or_create(self.email, self.first_name, self.last_name) to_update = {} if self.email != self._original.email: to_update["email"] = self.email if self.first_name != self._original.first_name: to_update["first_name"] = self.first_name if self.last_name != self._original.last_name: to_update["last_name"] = self.last_name update_customer_fields(self.stripe_id, **to_update) super().save(*args, **kwargs) def delete(self, *args, **kwargs): if self.stripe_id: stripe.Customer.delete(self.stripe_id) logger.info(f"Deleted Stripe customer {self.stripe_id}") super().delete(*args, **kwargs) def has_plan(self, plan): plan_list = [plan.name for plan in self.plans.all()] return plan in plan_list def get_notification_settings(self): return NotificationSettings.objects.get_or_create(user=self)[0] class Session(models.Model): user = models.ForeignKey(User, on_delete=models.CASCADE) request = models.CharField(max_length=255, null=True, blank=True) session = models.CharField(max_length=255, null=True, blank=True) subscription_id = models.CharField(max_length=255, null=True, blank=True) plan = models.ForeignKey(Plan, null=True, blank=True, on_delete=models.CASCADE) class ContentBlock(models.Model): user = models.ForeignKey(User, on_delete=models.PROTECT) position = models.IntegerField() page = models.CharField(max_length=255, null=True, blank=True) title = models.CharField(max_length=255, null=True, blank=True) column1 = models.TextField(null=True, blank=True) column2 = models.TextField(null=True, blank=True) column3 = models.TextField(null=True, blank=True) image1 = models.CharField(max_length=255, null=True, blank=True) image2 = models.CharField(max_length=255, null=True, blank=True) image3 = models.CharField(max_length=255, null=True, blank=True) def __str__(self): return f"[{self.position}] {self.page} {self.title}" def save(self, *args, **kwargs): """ Override the save function to blank fields. """ if self.column1 == "": self.column1 = None if self.column2 == "": self.column2 = None if self.column3 == "": self.column3 = None if self.image1 == "": self.image1 = None if self.image2 == "": self.image2 = None if self.image3 == "": self.image3 = None super().save(*args, **kwargs) class Perms(models.Model): class Meta: permissions = ( ("bypass_hashing", "Can bypass field hashing"), # ("bypass_blacklist", "Can bypass the blacklist"), # ("bypass_encryption", "Can bypass field encryption"), # ("bypass_obfuscation", "Can bypass field obfuscation"), # ("bypass_delay", "Can bypass data delay"), # ("bypass_randomisation", "Can bypass data randomisation"), # ("post_irc", "Can post to IRC"), ("post_discord", "Can post to Discord"), ("query_search", "Can search with query strings"), # ("use_insights", "Can use the Insights page"), ("index_internal", "Can use the internal index"), ("index_meta", "Can use the meta index"), ("index_restricted", "Can use the restricted index"), ("restricted_sources", "Can access restricted sources"), ) class NotificationRule(models.Model): user = models.ForeignKey(User, on_delete=models.CASCADE) name = models.CharField(max_length=255) enabled = models.BooleanField(default=True) data = models.TextField() def __str__(self): return f"{self.user} - {self.name}" def parse(self): try: parsed = load(self.data, Loader=Loader) except (ScannerError, ParserError) as e: raise ValueError(f"Invalid YAML: {e}") return parsed class NotificationSettings(models.Model): user = models.OneToOneField(User, on_delete=models.CASCADE) ntfy_topic = models.CharField(max_length=255, null=True, blank=True) ntfy_url = models.CharField(max_length=255, null=True, blank=True) def __str__(self): return f"Notification settings for {self.user}"