77 lines
2.5 KiB
Python
77 lines
2.5 KiB
Python
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
from dataclasses import dataclass
|
|
from functools import lru_cache
|
|
|
|
from fastapi import HTTPException
|
|
from cryptography.hazmat.primitives import hashes, serialization
|
|
from cryptography.hazmat.primitives.asymmetric import padding, rsa
|
|
|
|
|
|
@dataclass(frozen=True)
|
|
class LoginPublicKey:
|
|
key_id: str
|
|
public_key_pem: str
|
|
algorithm: str = "RSA-OAEP-SHA256"
|
|
|
|
|
|
def _normalize_env_pem(value: str) -> bytes:
|
|
return value.replace("\\n", "\n").encode("utf-8")
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def _load_private_key():
|
|
private_key_path = os.getenv("AUTH_LOGIN_PRIVATE_KEY_PATH")
|
|
private_key_pem = os.getenv("AUTH_LOGIN_PRIVATE_KEY")
|
|
|
|
if private_key_path:
|
|
with open(private_key_path, "rb") as key_file:
|
|
key_data = key_file.read()
|
|
return serialization.load_pem_private_key(key_data, password=None)
|
|
|
|
if private_key_pem:
|
|
return serialization.load_pem_private_key(_normalize_env_pem(private_key_pem), password=None)
|
|
|
|
return rsa.generate_private_key(public_exponent=65537, key_size=2048)
|
|
|
|
|
|
@lru_cache(maxsize=1)
|
|
def get_login_public_key() -> LoginPublicKey:
|
|
private_key = _load_private_key()
|
|
public_key = private_key.public_key()
|
|
public_der = public_key.public_bytes(
|
|
encoding=serialization.Encoding.DER,
|
|
format=serialization.PublicFormat.SubjectPublicKeyInfo,
|
|
)
|
|
public_pem = public_key.public_bytes(
|
|
encoding=serialization.Encoding.PEM,
|
|
format=serialization.PublicFormat.SubjectPublicKeyInfo,
|
|
).decode("utf-8")
|
|
key_id = hashlib.sha256(public_der).hexdigest()[:16]
|
|
return LoginPublicKey(key_id=key_id, public_key_pem=public_pem)
|
|
|
|
|
|
def decrypt_login_password(password_ciphertext: str, login_key_id: str) -> str:
|
|
key_info = get_login_public_key()
|
|
if login_key_id != key_info.key_id:
|
|
raise HTTPException(status_code=400, detail="登录密钥已过期，请刷新后重试")
|
|
|
|
try:
|
|
ciphertext = base64.b64decode(password_ciphertext.encode("utf-8"), validate=True)
|
|
plaintext = _load_private_key().decrypt(
|
|
ciphertext,
|
|
padding.OAEP(
|
|
mgf=padding.MGF1(algorithm=hashes.SHA256()),
|
|
algorithm=hashes.SHA256(),
|
|
label=None,
|
|
),
|
|
)
|
|
return plaintext.decode("utf-8")
|
|
except HTTPException:
|
|
raise
|
|
except Exception as exc:
|
|
raise HTTPException(status_code=400, detail="登录密码密文无效，请刷新后重试") from exc
|