58 lines
1.4 KiB
Python
58 lines
1.4 KiB
Python
"""API Key 加密存储与掩码展示。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import base64
|
|
import hashlib
|
|
import os
|
|
|
|
from cryptography.fernet import Fernet, InvalidToken
|
|
|
|
from . import config
|
|
|
|
|
|
def _load_key() -> bytes:
|
|
"""优先使用环境变量，其次复用本地密钥文件，都没有则生成新的。"""
|
|
env = os.getenv("AI_IMAGE_SECRET_KEY")
|
|
if env:
|
|
return env.encode("utf-8")
|
|
if config.SECRET_PATH.exists():
|
|
return config.SECRET_PATH.read_bytes().strip()
|
|
key = Fernet.generate_key()
|
|
config.SECRET_PATH.parent.mkdir(parents=True, exist_ok=True)
|
|
config.SECRET_PATH.write_bytes(key)
|
|
return key
|
|
|
|
|
|
def _fernet() -> Fernet:
|
|
raw = _load_key()
|
|
try:
|
|
return Fernet(raw)
|
|
except Exception:
|
|
derived = base64.urlsafe_b64encode(hashlib.sha256(raw).digest())
|
|
return Fernet(derived)
|
|
|
|
|
|
def encrypt(plain: str | None) -> str:
|
|
if not plain:
|
|
return ""
|
|
return _fernet().encrypt(plain.encode("utf-8")).decode("ascii")
|
|
|
|
|
|
def decrypt(token: str | None) -> str:
|
|
if not token:
|
|
return ""
|
|
try:
|
|
return _fernet().decrypt(token.encode("ascii")).decode("utf-8")
|
|
except (InvalidToken, ValueError):
|
|
return ""
|
|
|
|
|
|
def mask(plain: str | None) -> str:
|
|
"""展示用掩码：sk-a****abcd。"""
|
|
if not plain:
|
|
return ""
|
|
if len(plain) <= 8:
|
|
return plain[:2] + "****"
|
|
return f"{plain[:3]}****-{plain[-4:]}"
|