# -*- coding: utf-8 -*-
import six
import shutil
import tempfile
from contextlib import contextmanager

import pytest
from django.conf import settings
from django.core.exceptions import ImproperlyConfigured
from django.db import connection, models
from django.forms.widgets import Textarea, TextInput
from django.test import TestCase
from django.test.utils import override_settings

from .testapp.models import Secret

# Only perform encrypted fields tests if keyczar is present. Resolves
# http://github.com/django-extensions/django-extensions/issues/#issue/17
try:
    from django_extensions.db.fields.encrypted import BaseEncryptedField, EncryptedCharField, EncryptedTextField
    from keyczar import keyczar, keyczart, keyinfo  # NOQA
    keyczar_active = True
except ImportError:
    keyczar_active = False

# Locations of both private and public keys.
KEY_LOCS = {}


@pytest.fixture(scope="class")
def keyczar_keys(request):
    # If KeyCzar is available, set up the environment.
    if keyczar_active:
        # Create an RSA private key.
        keys_dir = tempfile.mkdtemp("django_extensions_tests_keyzcar_rsa_dir")
        keyczart.Create(keys_dir, "test", keyinfo.DECRYPT_AND_ENCRYPT, asymmetric=True)
        keyczart.AddKey(keys_dir, "PRIMARY", size=4096)
        KEY_LOCS['DECRYPT_AND_ENCRYPT'] = keys_dir

        # Create an RSA public key.
        pub_dir = tempfile.mkdtemp("django_extensions_tests_keyzcar_pub_dir")
        keyczart.PubKey(keys_dir, pub_dir)
        KEY_LOCS['ENCRYPT'] = pub_dir

    # cleanup crypto key temp dirs
    def cleanup():
        for name, path in KEY_LOCS.items():
            shutil.rmtree(path)
    request.addfinalizer(cleanup)


@contextmanager
def keys(purpose, mode=None):
    """
    A context manager that sets up the correct KeyCzar environment for a test.

    Arguments:
        purpose: Either keyczar.keyinfo.DECRYPT_AND_ENCRYPT or
                 keyczar.keyinfo.ENCRYPT.
        mode: If truthy, settings.ENCRYPTED_FIELD_MODE will be set to (and then
              reverted from) this value. If falsy, settings.ENCRYPTED_FIELD_MODE
              will not be changed. Optional. Default: None.

    Yields:
        A Keyczar subclass for the stated purpose. This will be keyczar.Crypter
        for DECRYPT_AND_ENCRYPT or keyczar.Encrypter for ENCRYPT. In addition,
        settings.ENCRYPTED_FIELD_KEYS_DIR will be set correctly, and then
        reverted when the manager exits.
    """
    # Store the original settings so we can restore when the manager exits.
    orig_setting_dir = getattr(settings, 'ENCRYPTED_FIELD_KEYS_DIR', None)
    orig_setting_mode = getattr(settings, 'ENCRYPTED_FIELD_MODE', None)
    try:
        if mode:
            settings.ENCRYPTED_FIELD_MODE = mode

        if purpose == keyinfo.DECRYPT_AND_ENCRYPT:
            settings.ENCRYPTED_FIELD_KEYS_DIR = KEY_LOCS['DECRYPT_AND_ENCRYPT']
            yield keyczar.Crypter.Read(settings.ENCRYPTED_FIELD_KEYS_DIR)
        else:
            settings.ENCRYPTED_FIELD_KEYS_DIR = KEY_LOCS['ENCRYPT']
            yield keyczar.Encrypter.Read(settings.ENCRYPTED_FIELD_KEYS_DIR)
    except Exception:
        raise  # Reraise any exceptions.
    finally:
        # Restore settings.
        settings.ENCRYPTED_FIELD_KEYS_DIR = orig_setting_dir
        if mode:
            if orig_setting_mode:
                settings.ENCRYPTED_FIELD_MODE = orig_setting_mode
            else:
                del settings.ENCRYPTED_FIELD_MODE


@contextmanager
def secret_model():
    """
    A context manager that yields a Secret model defined at runtime.

    All EncryptedField init logic occurs at model class definition time, not at
    object instantiation time. This means that in order to test different keys
    and modes, we must generate a new class definition at runtime, after
    establishing the correct KeyCzar settings. This context manager handles
    that process.

    See https://dynamic-models.readthedocs.io/en/latest/ and
    https://docs.djangoproject.com/en/dev/topics/db/models/
        #differences-between-proxy-inheritance-and-unmanaged-models
    """

    # Create a new class that shadows tests.models.Secret.
    attrs = {
        'name': EncryptedCharField("Name", max_length=Secret._meta.get_field('name').max_length),
        'text': EncryptedTextField("Text"),
        '__module__': 'tests.testapp.models',
        'Meta': type('Meta', (object, ), {
            'managed': False,
            'db_table': Secret._meta.db_table
        })
    }
    yield type('Secret', (models.Model, ), attrs)


@pytest.mark.skipif(keyczar_active is False, reason="Encrypted fields needs that keyczar is installed")
@pytest.mark.usefixtures("admin_user", "keyczar_keys")
class EncryptedFieldsTestCase(TestCase):
    def test_char_field_create(self):
        """
        Uses a private key to encrypt data on model creation.
        Verifies the data is encrypted in the database and can be decrypted.
        """
        with keys(keyinfo.DECRYPT_AND_ENCRYPT) as crypt:
            with secret_model() as model:
                test_val = "Test Secret"
                secret = model.objects.create(name=test_val)

                cursor = connection.cursor()
                query = "SELECT name FROM %s WHERE id = %d" % (model._meta.db_table, secret.id)
                cursor.execute(query)
                db_val, = cursor.fetchone()
                decrypted_val = crypt.Decrypt(db_val[len(EncryptedCharField.prefix):])
                self.assertEqual(test_val, decrypted_val)

    def test_char_field_read(self):
        """
        Uses a private key to encrypt data on model creation.
        Verifies the data is decrypted when reading the value back from the
        model.
        """
        with keys(keyinfo.DECRYPT_AND_ENCRYPT):
            with secret_model() as model:
                test_val = "Test Secret"
                secret = model.objects.create(name=test_val)
                retrieved_secret = model.objects.get(id=secret.id)
                self.assertEqual(test_val, retrieved_secret.name)

    def test_text_field_create(self):
        """
        Uses a private key to encrypt data on model creation.
        Verifies the data is encrypted in the database and can be decrypted.
        """
        with keys(keyinfo.DECRYPT_AND_ENCRYPT) as crypt:
            with secret_model() as model:
                test_val = "Test Secret"
                secret = model.objects.create(text=test_val)
                cursor = connection.cursor()
                query = "SELECT text FROM %s WHERE id = %d" % (model._meta.db_table, secret.id)
                cursor.execute(query)
                db_val, = cursor.fetchone()
                decrypted_val = crypt.Decrypt(db_val[len(EncryptedCharField.prefix):])
                self.assertEqual(test_val, decrypted_val)

    def test_text_field_read(self):
        """
        Uses a private key to encrypt data on model creation.
        Verifies the data is decrypted when reading the value back from the
        model.
        """
        with keys(keyinfo.DECRYPT_AND_ENCRYPT):
            with secret_model() as model:
                test_val = "Test Secret"
                secret = model.objects.create(text=test_val)
                retrieved_secret = model.objects.get(id=secret.id)
                self.assertEqual(test_val, retrieved_secret.text)

    def test_cannot_decrypt(self):
        """
        Uses a public key to encrypt data on model creation.
        Verifies that the data cannot be decrypted using the same key.
        """
        with keys(keyinfo.ENCRYPT, mode=keyinfo.ENCRYPT.name):
            with secret_model() as model:
                test_val = "Test Secret"
                secret = model.objects.create(name=test_val)
                retrieved_secret = model.objects.get(id=secret.id)
                self.assertNotEqual(test_val, retrieved_secret.name)
                self.assertTrue(retrieved_secret.name.startswith(EncryptedCharField.prefix))

    def test_unacceptable_purpose(self):
        """
        Tries to create an encrypted field with a mode mismatch.
        A purpose of "DECRYPT_AND_ENCRYPT" cannot be used with a public key,
        since public keys cannot be used for decryption. This should raise an
        exception.
        """
        with self.assertRaises(keyczar.errors.KeyczarError):
            with keys(keyinfo.ENCRYPT):
                with secret_model():
                    # A KeyCzar exception should get raised during class
                    # definition time, so any code in here would never get run.
                    pass

    def test_decryption_forbidden(self):
        """
        Uses a private key to encrypt data, but decryption is not allowed.
        ENCRYPTED_FIELD_MODE is explicitly set to ENCRYPT, meaning data should
        not be decrypted, even though the key would allow for it.
        """
        with keys(keyinfo.DECRYPT_AND_ENCRYPT, mode=keyinfo.ENCRYPT.name):
            with secret_model() as model:
                test_val = "Test Secret"
                secret = model.objects.create(name=test_val)
                retrieved_secret = model.objects.get(id=secret.id)
                self.assertNotEqual(test_val, retrieved_secret.name)
                self.assertTrue(retrieved_secret.name.startswith(EncryptedCharField.prefix))

    def test_encrypt_public_decrypt_private(self):
        """
        Uses a public key to encrypt, and a private key to decrypt data.
        """
        test_val = "Test Secret"

        # First, encrypt data with public key and save to db.
        with keys(keyinfo.ENCRYPT, mode=keyinfo.ENCRYPT.name):
            with secret_model() as model:
                secret = model.objects.create(name=test_val)
                enc_retrieved_secret = model.objects.get(id=secret.id)
                self.assertNotEqual(test_val, enc_retrieved_secret.name)
                self.assertTrue(enc_retrieved_secret.name.startswith(EncryptedCharField.prefix))

        # Next, retrieve data from db, and decrypt with private key.
        with keys(keyinfo.DECRYPT_AND_ENCRYPT):
            with secret_model() as model:
                retrieved_secret = model.objects.get(id=secret.id)
                self.assertEqual(test_val, retrieved_secret.name)


class BaseEncryptedFieldTestCase(TestCase):

    @classmethod
    def setUpClass(cls):
        cls.tmpdir = tempfile.mkdtemp()
        keyczart.Create(cls.tmpdir, "test", keyinfo.DECRYPT_AND_ENCRYPT, asymmetric=True)
        keyczart.AddKey(cls.tmpdir, "PRIMARY", size=4096)

    @classmethod
    def tearDownClass(cls):
        shutil.rmtree(cls.tmpdir)


@pytest.mark.skipif(keyczar_active is False, reason="Encrypted fields needs that keyczar is installed")
class BaseEncryptedFieldExceptions(BaseEncryptedFieldTestCase):
    """Tests for BaseEncryptedField exceptions."""

    def test_should_raise_ImproperlyConfigured_if_invalid_ENCRYPTED_FIELD_MODE_is_set(self):
        with override_settings(ENCRYPTED_FIELD_KEYS_DIR=self.tmpdir, ENCRYPTED_FIELD_MODE='INVALID'):
            with six.assertRaisesRegex(self, ImproperlyConfigured, 'ENCRYPTED_FIELD_MODE must be either DECRYPT_AND_ENCRYPT or ENCRYPT, not INVALID.'):
                BaseEncryptedField()


@pytest.mark.skipif(keyczar_active is False, reason="Encrypted fields needs that keyczar is installed")
class EncryptedTextFieldTests(BaseEncryptedFieldTestCase):
    """Tests for EncryptedTextField."""

    def test_should_return_formfield_with_Textarea_widget(self):
        with override_settings(ENCRYPTED_FIELD_KEYS_DIR=self.tmpdir):
            formfield = EncryptedTextField(max_length=50).formfield()

        self.assertTrue(isinstance(formfield.widget, Textarea))


@pytest.mark.skipif(keyczar_active is False, reason="Encrypted fields needs that keyczar is installed")
class EncryptedCharFieldTests(BaseEncryptedFieldTestCase):
    """Tests for EncryptedCharField."""

    def test_should_return_formfield_with_TextInput_widget(self):
        with override_settings(ENCRYPTED_FIELD_KEYS_DIR=self.tmpdir):
            formfield = EncryptedCharField(max_length=50).formfield()

        self.assertTrue(isinstance(formfield.widget, TextInput))
        self.assertEqual(formfield.max_length, 700)
