from django.contrib.auth.models import AbstractUser
from django.db import models
import base64
from cryptography.fernet import Fernet


class CustomUser(AbstractUser):
    """사용자 모델 - 기존 필드 + SSH Private Key 관리 필드"""

    # 기존 필드 유지
    grade = models.CharField(max_length=50, blank=True, null=True)
    nhnc_id = models.CharField(max_length=100, blank=True, null=True)
    nhnc_api_tenant_id = models.CharField(max_length=100, blank=True, null=True)

    """사용자 모델 - SSH Private Key 관리 필드"""
    encrypted_private_key = models.BinaryField(blank=True, null=True)  # 암호화된 SSH 키

    def encrypt_private_key(self, private_key: str) -> bytes:
        """SSH Private Key 암호화"""
        cipher = Fernet(self.get_encryption_key())
        return cipher.encrypt(private_key.encode())

    def decrypt_private_key(self) -> str:
        """SSH Private Key 복호화"""
        if self.encrypted_private_key:
            cipher = Fernet(self.get_encryption_key())
            return cipher.decrypt(self.encrypted_private_key).decode()
        return ""

    def save_private_key(self, private_key: str):
        """SSH Private Key 저장 (암호화 후)"""
        self.encrypted_private_key = self.encrypt_private_key(private_key)
        self.save()

    def get_encryption_key(self) -> bytes:
        """Fernet 키 생성 (username 기반)"""
        username_encoded = self.username.encode()  # 사용자 이름을 바이트로 인코딩
        base64_key = base64.urlsafe_b64encode(username_encoded.ljust(32)[:32])
        return base64_key

    def __str__(self):
        return self.username