from aspen_crypto.keys import (
    get_key_password,
    generate_kid,
    Ed25519PrivateKey,
    RSAPrivateKey,
)
from sequential_uuids.generators import uuid_time_nextval
from typing import Optional, TYPE_CHECKING

if TYPE_CHECKING:
    from sandbox.latchstring.models import User as UserType
    from .models import EncryptionKey
from console_base.managers import LCBaseQuerySet

from .choices import STATUS, SUBJECT


class EncryptionQuerySet(LCBaseQuerySet):
    def create(self, **kwargs) -> 'EncryptionKey':
        """
        Attempt to add `kid` based on public key if it's not provided.
        """
        if not kwargs.get('kid'):
            kwargs['kid'] = generate_kid(key=kwargs['key'])

        return super().create(**kwargs)

    def valid(self) -> 'EncryptionKey':
        return self.active().filter(status=STATUS.active)

    def create_rsa_keys(
        self,
        name: str,
        user: 'UserType',
        subject: SUBJECT,
        password_prefix: str = '',
    ) -> Optional['EncryptionKey']:
        """
        Create a new RSA Private / Public Key Pair record for the specified user.
        """
        return self._create_key_pair(
            key_cls=RSAPrivateKey,
            name=name,
            user=user,
            subject=subject,
            password_prefix=password_prefix,
        )

    def create_ssh_keys(
        self,
        name: str,
        user: 'UserType',
        subject: SUBJECT,
        password_prefix: str = '',
    ) -> Optional['EncryptionKey']:
        """
        Create a new SSH Private / Public Key Pair record for the specified user.
        """
        return self._create_key_pair(
            key_cls=Ed25519PrivateKey,
            name=name,
            user=user,
            subject=subject,
            password_prefix=password_prefix,
        )

    def _create_key_pair(
        self,
        key_cls,
        name: str,
        user: 'UserType',
        subject: SUBJECT,
        password_prefix: str = '',
    ) -> Optional['EncryptionKey']:
        """
        Helper method; don't call directly.
        Create new SSH or RSA Public / Private key pair for specified user.
        """
        cid = uuid_time_nextval()
        private_key = key_cls.generate()
        password = get_key_password(cid, password_prefix, ssh=key_cls == Ed25519PrivateKey)

        if not password:
            return

        private_bytes = private_key.private_bytes(password)

        return self.create(
            name=name,
            user=user,
            cid=cid,
            key=private_key.public_key.as_pem.decode(),
            subject=subject,
            private_key=private_bytes,
        )
