import logging import uuid import stripe from django.conf import settings 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__) PRIORITY_CHOICES = ( (1, "min"), (2, "low"), (3, "default"), (4, "high"), (5, "max"), ) INTERVAL_CHOICES = ( (0, "On demand"), (5, "Every 5 seconds"), (60, "Every minute"), (900, "Every 15 minutes"), (1800, "Every 30 minutes"), (3600, "Every hour"), (14400, "Every 4 hours"), (86400, "Every day"), ) SERVICE_CHOICES = ( ("ntfy", "NTFY"), ("webhook", "Custom webhook"), ("none", "Disabled"), ) POLICY_CHOICES = ( ("default", "Default: Trigger only when there were no results last time"), ( "change", "Change: Default + trigger when there are no results (if there were before)", ), ("always", "Always: Trigger on every run (not recommended for low intervals)"), ) 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, check=True): sets = NotificationSettings.objects.get_or_create(user=self)[0] if check: if sets.service == "ntfy" and sets.topic is None: return None if sets.service == "webhook" and sets.url is None: return None return sets @property def allowed_indices(self): indices = [settings.INDEX_MAIN] if self.has_perm("core.index_meta"): indices.append(settings.INDEX_META) if self.has_perm("core.index_internal"): indices.append(settings.INDEX_INT) if self.has_perm("core.index_restricted"): if self.has_perm("core.restricted_sources"): indices.append(settings.INDEX_RESTRICTED) return indices 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 = ( ("post_irc", "Can post to IRC"), ("post_discord", "Can post to Discord"), ("use_insights", "Can use the Insights page"), ("use_rules", "Can use the Rules page"), ("rules_scheduled", "Can use the scheduled rules"), ("rules_high_frequency", "Can use the high frequency rules"), ("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): id = models.UUIDField( default=uuid.uuid4, primary_key=True, editable=False, unique=True ) user = models.ForeignKey(User, on_delete=models.CASCADE) name = models.CharField(max_length=255) priority = models.IntegerField(choices=PRIORITY_CHOICES, default=1) topic = models.CharField(max_length=2048, null=True, blank=True) url = models.CharField(max_length=1024, null=True, blank=True) interval = models.IntegerField(choices=INTERVAL_CHOICES, default=0) window = models.CharField(max_length=255, null=True, blank=True) amount = models.PositiveIntegerField(default=1, null=True, blank=True) enabled = models.BooleanField(default=True) data = models.TextField() match = models.JSONField(null=True, blank=True) service = models.CharField(choices=SERVICE_CHOICES, max_length=255, default="ntfy") send_empty = models.BooleanField(default=False) policy = models.CharField(choices=POLICY_CHOICES, max_length=255, default="default") 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 def get_notification_settings(self, check=True): """ Get the notification settings for this rule. Notification rule settings take priority. """ user_settings = self.user.get_notification_settings(check=False) user_settings = user_settings.__dict__ if self.priority is not None: user_settings["priority"] = str(self.priority) if self.topic is not None: user_settings["topic"] = self.topic if self.url is not None: user_settings["url"] = self.url if self.service is not None: user_settings["service"] = self.service if self.send_empty is not None: user_settings["send_empty"] = self.send_empty if check: if user_settings["service"] == "ntfy" and user_settings["topic"] is None: return None if user_settings["service"] == "webhook" and user_settings["url"] is None: return None return user_settings class NotificationSettings(models.Model): user = models.OneToOneField(User, on_delete=models.CASCADE) topic = models.CharField(max_length=2048, null=True, blank=True) url = models.CharField(max_length=1024, null=True, blank=True) service = models.CharField(choices=SERVICE_CHOICES, max_length=255, default="ntfy") def __str__(self): return f"Notification settings for {self.user}"