neptune/core/forms.py

from django import forms
from django.contrib.auth.forms import UserCreationForm
from django.core.exceptions import FieldDoesNotExist
from django.forms import ModelForm

from core.db import QueryError
from core.lib.rules import NotificationRuleData, RuleParseError

from .models import NotificationRule, NotificationSettings, User

# flake8: noqa: E501


class RestrictedFormMixin:
    """
    This mixin is used to restrict the queryset of a form to the current user.
    The request object is passed from the view.
    Fieldargs is used to pass additional arguments to the queryset filter.
    """

    fieldargs = {}

    def __init__(self, *args, **kwargs):
        # self.fieldargs = {}
        self.request = kwargs.pop("request")
        super().__init__(*args, **kwargs)
        for field in self.fields:
            # Check it's not something like a CharField which has no queryset
            if not hasattr(self.fields[field], "queryset"):
                continue

            model = self.fields[field].queryset.model
            # Check if the model has a user field
            try:
                model._meta.get_field("user")
                # Add the user to the queryset filters
                self.fields[field].queryset = model.objects.filter(
                    user=self.request.user, **self.fieldargs.get(field, {})
                )
            except FieldDoesNotExist:
                pass


class NewUserForm(UserCreationForm):
    email = forms.EmailField(required=True)

    class Meta:
        model = User
        fields = (
            "username",
            "email",
            "first_name",
            "last_name",
            "password1",
            "password2",
        )

    def save(self, commit=True):
        user = super(NewUserForm, self).save(commit=False)
        user.email = self.cleaned_data["email"]
        if commit:
            user.save()
        return user


class CustomUserCreationForm(UserCreationForm):
    class Meta:
        model = User
        fields = "__all__"


class NotificationSettingsForm(RestrictedFormMixin, ModelForm):
    class Meta:
        model = NotificationSettings
        fields = (
            "ntfy_topic",
            "ntfy_url",
        )
        help_texts = {
            "ntfy_topic": "The topic to send notifications to.",
            "ntfy_url": "Custom NTFY server. Leave blank to use the default server.",
        }


class NotificationRuleForm(RestrictedFormMixin, ModelForm):
    class Meta:
        model = NotificationRule
        fields = (
            "name",
            "data",
            "interval",
            "window",
            "priority",
            "topic",
            "enabled",
        )
        help_texts = {
            "name": "The name of the rule.",
            "priority": "The notification priority of the rule.",
            "topic": "The topic to send notifications to. Leave blank for default.",
            "enabled": "Whether the rule is enabled.",
            "data": "The notification rule definition.",
            "interval": "How often to run the search. On demand only evaluates messages as they are received.",
            "window": "Time window to search: 1d, 1h, 1m, 1s, etc.",
        }

    def clean(self):
        cleaned_data = super(NotificationRuleForm, self).clean()
        try:
            parsed_data = NotificationRuleData(self.request.user, cleaned_data)
        except RuleParseError as e:
            self.add_error(e.field, f"Parsing error: {e}")
            return
        except QueryError as e:
            self.add_error("data", f"Query error: {e}")
            return

        # Write back the validated data
        # We need this to populate the index and source variable if
        # they are not set
        to_store = str(parsed_data)
        cleaned_data["data"] = to_store

        return cleaned_data