from aspen_crypto.keys import PublicKey
from console_base.utils.cli_input import readline_input
from lcrequests.umbrella import UMBRELLA_USERS
from uuid import UUID

from django.contrib.auth import get_user_model
from django.core.management.base import BaseCommand, CommandError
from encipher.models import EncryptionKey
from encipher.utils import confirm_umbrella_user


class Command(BaseCommand):
    help = "Import a Public Key for Umbrella API Auth"

    def handle(self, *args, **options):
        User = get_user_model()

        try:
            name = ''
            while not name:
                name = readline_input('Name for this Public Key: ')

            user_cid = None
            while not user_cid:
                cid_or_subject = readline_input('User CID or Umbrella Subject: ').lower()
                if umbrella_user := UMBRELLA_USERS.get(cid_or_subject):
                    user_cid = umbrella_user.cid
                    confirm_umbrella_user(umbrella_user)
                else:
                    try:
                        user_cid = UUID(cid_or_subject)
                        if not User.objects.filter(cid=user_cid).first():
                            self.stdout.write(f'No user found with CID: {cid_or_subject}')
                            user_cid = None
                    except ValueError:
                        self.stdout.write(
                            f'{cid_or_subject} is not a valid CID or Umbrella Subject'
                        )

            key_cid = None
            while not key_cid:
                try:
                    key_cid = UUID(readline_input('Key CID: ').strip("'").strip().lower())
                    if EncryptionKey.objects.filter(cid=key_cid).exists():
                        self.stdout.write('Public Key with this CID already exists')
                        key_cid = None
                except ValueError:
                    self.stdout.write('Not a valid Canonical ID')
                    continue

            public_key = None
            while not public_key:
                pks = readline_input('Public Key PEM: ').strip("'").strip()
                try:
                    public_key = PublicKey.load_pem(pks.encode().decode('unicode_escape').encode())
                except (ValueError, TypeError):
                    self.stdout.write('Not a valid Public Key PEM value')
        except KeyboardInterrupt:
            raise CommandError('Exiting Umbrella Auth import')

        try:
            EncryptionKey.objects.create(
                id=key_cid,
                cid=key_cid,
                user=User.objects.filter(cid=user_cid).first(),
                key=public_key.as_pem.decode(),
                name=name,
            )
        except Exception as e:
            raise CommandError(e) from None
