from aspen_crypto.keys import Ed25519PrivateKey, Ed25519PublicKey, RSAPrivateKey, RSAPublicKey
from django.conf import settings
from django.test import TestCase
from django.contrib.auth import get_user_model
from django.core.exceptions import ValidationError
from ..choices import SUBJECT
from .. import models, validators
from . import data

User = get_user_model()


class ValidateEncryptionKeyTest(TestCase):
    def test_valid_rsa_pem(self):
        validators.validate_public_key(data.PEM_PUBLIC_RSA.decode())

    def test_valid_ed25519_pem(self):
        validators.validate_public_key(data.PEM_PUBLIC_ED25519.decode())

    def test_valid_rsa_openssh(self):
        validators.validate_public_key(data.OPENSSH_RSA.decode())

    def test_valid_ed25519_openssh(self):
        validators.validate_public_key(data.OPENSSH_ED25519.decode())

    def test_invalid_rsa_pem(self):
        with self.assertRaises(ValidationError):
            validators.validate_public_key(data.PEM_PUBLIC_RSA_INVALID.decode())

    def test_invalid_ed25519_pem(self):
        with self.assertRaises(ValidationError):
            validators.validate_public_key(data.PEM_PUBLIC_ED25519_INVALID.decode())

    def test_invalid_rsa_openssh(self):
        with self.assertRaises(ValidationError):
            validators.validate_public_key(data.OPENSSH_RSA_INVALID.decode())

    def test_invalid_ed25519_openssh(self):
        with self.assertRaises(ValidationError):
            validators.validate_public_key(data.OPENSSH_ED25519_INVALID.decode())


class EncryptionKeyTest(TestCase):
    def setUp(self):
        self.user = User.objects.create_user(username="foo")
        self.user.refresh_from_db()

    def test_create_rsa_keys(self):
        key = models.EncryptionKey.objects.create_rsa_keys(
            name='My-RSA',
            user=self.user,
            subject=SUBJECT.Beacon,
            password_prefix=settings.PRIVATE_KEY_PASSWORD_PREFIX,
        )
        key.refresh_from_db()
        self.assertIsInstance(key, models.EncryptionKey)
        private_key = key.get_private_key()
        self.assertIsInstance(private_key, RSAPrivateKey)

    def test_create_ssh_keys(self):
        key = models.EncryptionKey.objects.create_ssh_keys(
            name='My-SSH',
            user=self.user,
            subject=SUBJECT.DrawBridge,
            password_prefix=settings.PRIVATE_KEY_PASSWORD_PREFIX,
        )
        key.refresh_from_db()
        self.assertIsInstance(key, models.EncryptionKey)
        private_key = key.get_private_key()
        self.assertIsInstance(private_key, Ed25519PrivateKey)

    def test_extract_comment(self):
        pub = models.EncryptionKey(user=self.user, key=data.OPENSSH_ED25519.decode(), name="")
        pub.save()
        self.assertEqual(pub.name, "crgwbr@foo")

    def test_get_key_ed25519(self):
        pub = models.EncryptionKey(user=self.user, key=data.OPENSSH_ED25519.decode())
        pub.save()
        self.assertIsInstance(pub.get_key(), Ed25519PublicKey)

    def test_get_loaded_key_rsa(self):
        pub = models.EncryptionKey(user=self.user, key=data.OPENSSH_RSA.decode(), name="")
        pub.save()
        self.assertIsInstance(pub.get_key(), RSAPublicKey)

    def test_get_loaded_key_invalid(self):
        pub = models.EncryptionKey(user=self.user, key=data.OPENSSH_ED25519_INVALID.decode())
        pub.save()
        with self.assertRaises(ValueError):
            pub.get_key()
