neptune/core/models.py

import logging

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"),
)


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"),
            ("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)
    priority = models.IntegerField(choices=PRIORITY_CHOICES, default=1)
    topic = models.CharField(max_length=255, 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)

    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

    @property
    def matches(self):
        """
        Get the total number of matches for this rule.
        """
        if isinstance(self.match, dict):
            return f"{sum(list(self.match.values()))}/{len(self.match)}"

    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=255, 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}"