neptune/core/models.py

191 lines
6.4 KiB
Python

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"),
(60, "Every minute"),
(900, "Every 15 minutes"),
(1800, "Every 30 minutes"),
(3600, "Every hour"),
(14400, "Every 4 hours"),
(86400, "Every day"),
)
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]
@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)
interval = models.IntegerField(choices=INTERVAL_CHOICES, default=0)
window = models.CharField(max_length=255, null=True, blank=True)
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}"