wireguard_webadmin/gatekeeper/forms.py

import pyotp
from crispy_forms.helper import FormHelper
from crispy_forms.layout import Layout, Submit, HTML, Div, Field
from django import forms
from django.utils.translation import gettext_lazy as _

from gatekeeper.models import GatekeeperUser, GatekeeperGroup, AuthMethod, AuthMethodAllowedDomain, \
    AuthMethodAllowedEmail


class GatekeeperUserForm(forms.ModelForm):
    class Meta:
        model = GatekeeperUser
        fields = ['username', 'email', 'password', 'totp_secret']
        labels = {
            'username': _('Username'),
            'email': _('Email'),
            'password': _('Password'),
            'totp_secret': _('TOTP Secret'),
        }

    def __init__(self, *args, **kwargs):
        cancel_url = kwargs.pop('cancel_url', '#')
        super().__init__(*args, **kwargs)

        self.helper = FormHelper()
        self.helper.layout = Layout(
            Div(
                Div('username', css_class='col-md-6'),
                Div('email', css_class='col-md-6'),
                css_class='row'
            ),
            Div(
                Div(Field('password', type='password'), css_class='col-md-6'),
                Div('totp_secret', css_class='col-md-6'),
                css_class='row'
            ),
            Div(
                Div(
                    Submit('submit', _('Save'), css_class='btn btn-primary'),
                    HTML(f'<a href="{cancel_url}" class="btn btn-secondary">{_("Cancel")}</a>'),
                    css_class='col-12 d-flex justify-content-end gap-2 mt-3'
                ),
                css_class='row'
            )
        )


class GatekeeperGroupForm(forms.ModelForm):
    class Meta:
        model = GatekeeperGroup
        fields = ['name', 'users']
        labels = {
            'name': _('Group Name'),
            'users': _('Members'),
        }

    def __init__(self, *args, **kwargs):
        cancel_url = kwargs.pop('cancel_url', '#')
        super().__init__(*args, **kwargs)

        self.helper = FormHelper()
        self.helper.layout = Layout(
            Div(
                Div('name', css_class='col-md-12'),
                css_class='row'
            ),
            Div(
                Div('users', css_class='col-md-12'),
                css_class='row'
            ),
            Div(
                Div(
                    Submit('submit', _('Save'), css_class='btn btn-primary'),
                    HTML(f'<a href="{cancel_url}" class="btn btn-secondary">{_("Cancel")}</a>'),
                    css_class='col-12 d-flex justify-content-end gap-2 mt-3'
                ),
                css_class='row'
            )
        )


class AuthMethodForm(forms.ModelForm):
    totp_pin = forms.CharField(
        label=_('TOTP Validation PIN'),
        max_length=6,
        required=False,
        help_text=_('Enter a 6-digit PIN generated by your authenticator app to validate the secret.')
    )

    class Meta:
        model = AuthMethod
        fields = [
            'name', 'auth_type', 'totp_secret', 
            'oidc_provider', 'oidc_client_id', 'oidc_client_secret'
        ]
        labels = {
            'name': _('Name'),
            'auth_type': _('Authentication Type'),
            'totp_secret': _('Global TOTP Secret'),
            'oidc_provider': _('OIDC Provider URL'),
            'oidc_client_id': _('OIDC Client ID'),
            'oidc_client_secret': _('OIDC Client Secret'),
        }

    def __init__(self, *args, **kwargs):
        cancel_url = kwargs.pop('cancel_url', '#')
        super().__init__(*args, **kwargs)

        if self.instance and self.instance.pk:
            self.fields['auth_type'].disabled = True

        self.helper = FormHelper()
        self.helper.layout = Layout(
            Div(
                Div('name', css_class='col-md-6'),
                Div('auth_type', css_class='col-md-6'),
                css_class='row'
            ),
            Div(
                Div('totp_secret', css_class='col-md-6'),
                Div('totp_pin', css_class='col-md-6'),
                css_class='row'
            ),
            Div(
                Div('oidc_provider', css_class='col-md-12'),
                css_class='row'
            ),
            Div(
                Div('oidc_client_id', css_class='col-md-6'),
                Div('oidc_client_secret', css_class='col-md-6'),
                css_class='row'
            ),
            Div(
                Div(
                    Submit('submit', _('Save'), css_class='btn btn-primary'),
                    HTML(f'<a href="{cancel_url}" class="btn btn-secondary">{_("Cancel")}</a>'),
                    css_class='col-12 d-flex justify-content-end gap-2 mt-3'
                ),
                css_class='row'
            )
        )

    def clean(self):
        cleaned_data = super().clean()
        auth_type = cleaned_data.get('auth_type')
        totp_secret = cleaned_data.get('totp_secret')
        oidc_provider = cleaned_data.get('oidc_provider')
        oidc_client_id = cleaned_data.get('oidc_client_id')
        oidc_client_secret = cleaned_data.get('oidc_client_secret')

        if auth_type == 'local_password':
            if totp_secret:
                self.add_error('totp_secret', _('TOTP secret must be empty for Local Password authentication.'))
            if cleaned_data.get('totp_pin'):
                self.add_error('totp_pin', _('TOTP validation PIN must be empty for Local Password authentication.'))
            if oidc_provider or oidc_client_id or oidc_client_secret:
                self.add_error(None, _('OIDC fields must be empty for Local Password authentication.'))
            
            existing_local = AuthMethod.objects.filter(auth_type='local_password')
            if self.instance and self.instance.pk:
                existing_local = existing_local.exclude(pk=self.instance.pk)
            if existing_local.exists():
                self.add_error('auth_type', _('Only one Local Password authentication method can be configured.'))
        elif auth_type == 'totp':
            if oidc_provider or oidc_client_id or oidc_client_secret:
                self.add_error(None, _('OIDC fields must be empty for TOTP authentication.'))
            if not totp_secret:
                self.add_error('totp_secret', _('TOTP secret is required for TOTP authentication.'))
            else:
                totp_pin = cleaned_data.get('totp_pin')
                if not totp_pin:
                    self.add_error('totp_pin', _('Please provide a PIN to validate the TOTP secret.'))
                else:
                    try:
                        totp = pyotp.TOTP(totp_secret)
                        if not totp.verify(totp_pin):
                            self.add_error('totp_pin', _('Invalid TOTP PIN.'))
                    except Exception:
                        self.add_error('totp_secret', _('Invalid TOTP secret format. Must be a valid Base32 string.'))
        elif auth_type == 'oidc':
            if totp_secret:
                self.add_error('totp_secret', _('TOTP secret must be empty for OIDC authentication.'))
            if cleaned_data.get('totp_pin'):
                self.add_error('totp_pin', _('TOTP validation PIN must be empty for OIDC authentication.'))

        return cleaned_data

class AuthMethodAllowedDomainForm(forms.ModelForm):
    class Meta:
        model = AuthMethodAllowedDomain
        fields = ['auth_method', 'domain']
        labels = {
            'auth_method': _('Authentication Method'),
            'domain': _('Domain'),
        }

    def __init__(self, *args, **kwargs):
        cancel_url = kwargs.pop('cancel_url', '#')
        super().__init__(*args, **kwargs)

        self.helper = FormHelper()
        self.helper.layout = Layout(
            Div(
                Div('auth_method', css_class='col-md-6'),
                Div('domain', css_class='col-md-6'),
                css_class='row'
            ),
            Div(
                Div(
                    Submit('submit', _('Save'), css_class='btn btn-primary'),
                    HTML(f'<a href="{cancel_url}" class="btn btn-secondary">{_("Cancel")}</a>'),
                    css_class='col-12 d-flex justify-content-end gap-2 mt-3'
                ),
                css_class='row'
            )
        )


class AuthMethodAllowedEmailForm(forms.ModelForm):
    class Meta:
        model = AuthMethodAllowedEmail
        fields = ['auth_method', 'email']
        labels = {
            'auth_method': _('Authentication Method'),
            'email': _('Email'),
        }

    def __init__(self, *args, **kwargs):
        cancel_url = kwargs.pop('cancel_url', '#')
        super().__init__(*args, **kwargs)

        self.helper = FormHelper()
        self.helper.layout = Layout(
            Div(
                Div('auth_method', css_class='col-md-6'),
                Div('email', css_class='col-md-6'),
                css_class='row'
            ),
            Div(
                Div(
                    Submit('submit', _('Save'), css_class='btn btn-primary'),
                    HTML(f'<a href="{cancel_url}" class="btn btn-secondary">{_("Cancel")}</a>'),
                    css_class='col-12 d-flex justify-content-end gap-2 mt-3'
                ),
                css_class='row'
            )
        )