from Crypto.Cipher import PKCS1_OAEP
from Crypto.PublicKey import RSA

from django import forms

from dcim.models import Device
from utilities.forms import (
    BootstrapMixin, BulkImportForm, CSVDataField, FilterChoiceField, SlugField, get_filter_choices,
)

from .models import Secret, SecretRole, UserKey


def validate_rsa_key(key, is_secret=True):
    """
    Validate the format and type of an RSA key.
    """
    try:
        key = RSA.importKey(key)
    except ValueError:
        raise forms.ValidationError("Invalid RSA key. Please ensure that your key is in PEM (base64) format.")
    except Exception as e:
        raise forms.ValidationError("Invalid key detected: {}".format(e))
    if is_secret and not key.has_private():
        raise forms.ValidationError("This looks like a public key. Please provide your private RSA key.")
    elif not is_secret and key.has_private():
        raise forms.ValidationError("This looks like a private key. Please provide your public RSA key.")
    try:
        PKCS1_OAEP.new(key)
    except:
        raise forms.ValidationError("Error validating RSA key. Please ensure that your key supports PKCS#1 OAEP.")


#
# Secret roles
#

class SecretRoleForm(forms.ModelForm, BootstrapMixin):
    slug = SlugField()

    class Meta:
        model = SecretRole
        fields = ['name', 'slug']


#
# Secrets
#

class SecretForm(forms.ModelForm, BootstrapMixin):
    private_key = forms.CharField(required=False, widget=forms.HiddenInput())
    plaintext = forms.CharField(max_length=65535, required=False, label='Plaintext',
                                widget=forms.TextInput(attrs={'class': 'requires-private-key'}))
    plaintext2 = forms.CharField(max_length=65535, required=False, label='Plaintext (verify)')

    class Meta:
        model = Secret
        fields = ['role', 'name', 'plaintext', 'plaintext2']

    def clean(self):
        if self.cleaned_data['plaintext']:
            validate_rsa_key(self.cleaned_data['private_key'])

    def clean_plaintext2(self):
        plaintext = self.cleaned_data['plaintext']
        plaintext2 = self.cleaned_data['plaintext2']
        if plaintext != plaintext2:
            raise forms.ValidationError("The two given plaintext values do not match. Please check your input.")


class SecretFromCSVForm(forms.ModelForm):
    device = forms.ModelChoiceField(queryset=Device.objects.all(), required=False, to_field_name='name',
                                    error_messages={'invalid_choice': 'Device not found.'})
    role = forms.ModelChoiceField(queryset=SecretRole.objects.all(), to_field_name='name',
                                  error_messages={'invalid_choice': 'Invalid secret role.'})
    plaintext = forms.CharField()

    class Meta:
        model = Secret
        fields = ['device', 'role', 'name', 'plaintext']

    def save(self, *args, **kwargs):
        s = super(SecretFromCSVForm, self).save(*args, **kwargs)
        s.plaintext = str(self.cleaned_data['plaintext'])
        return s


class SecretImportForm(BulkImportForm, BootstrapMixin):
    private_key = forms.CharField(widget=forms.HiddenInput())
    csv = CSVDataField(csv_form=SecretFromCSVForm, widget=forms.Textarea(attrs={'class': 'requires-private-key'}))


class SecretBulkEditForm(forms.Form, BootstrapMixin):
    pk = forms.ModelMultipleChoiceField(queryset=Secret.objects.all(), widget=forms.MultipleHiddenInput)
    role = forms.ModelChoiceField(queryset=SecretRole.objects.all())
    name = forms.CharField(max_length=100, required=False)


class SecretFilterForm(forms.Form, BootstrapMixin):
    role = FilterChoiceField(choices=get_filter_choices(SecretRole, id_field='slug', count_field='secrets'))


#
# UserKeys
#

class UserKeyForm(forms.ModelForm, BootstrapMixin):

    class Meta:
        model = UserKey
        fields = ['public_key']
        help_texts = {
            'public_key': "Enter your public RSA key. Keep the private one with you; you'll need it for decryption.",
        }

    def clean_public_key(self):
        key = self.cleaned_data['public_key']

        # Validate the RSA key format.
        validate_rsa_key(key, is_secret=False)

        return key


class ActivateUserKeyForm(forms.Form):
    _selected_action = forms.ModelMultipleChoiceField(queryset=UserKey.objects.all(), label='User Keys')
    secret_key = forms.CharField(label='Your private key', widget=forms.Textarea(attrs={'class': 'vLargeTextField'}))