143 lines
4.0 KiB
Python
143 lines
4.0 KiB
Python
# -*- coding: utf-8 -*-
|
||
"""master key 加载与生成。
|
||
|
||
加载优先级（高到低）：
|
||
1. 环境变量 ACCOUNT_MANAGER_MASTER_KEY（base64-urlsafe 字符串，44 字符）
|
||
2. 文件 {DATA_ROOT}/{USER_ID}/account-manager/master.key
|
||
3. 都没有 → raise MasterKeyMissingError，错误码 ERR_MASTER_KEY_MISSING
|
||
|
||
不再自动生成——必须显式 `account init-master-key`。
|
||
"""
|
||
from __future__ import annotations
|
||
|
||
import base64
|
||
import os
|
||
|
||
from cryptography.fernet import Fernet
|
||
|
||
from util.runtime_paths import get_skill_data_dir
|
||
|
||
ENV_MASTER_KEY = "ACCOUNT_MANAGER_MASTER_KEY"
|
||
_MASTER_FILENAME = "master.key"
|
||
|
||
|
||
class MasterKeyMissingError(Exception):
|
||
"""Raised when no valid master key can be resolved."""
|
||
|
||
code = "ERR_MASTER_KEY_MISSING"
|
||
|
||
|
||
class MasterKeyExistsError(Exception):
|
||
"""Raised when master.key already exists and overwrite is not allowed."""
|
||
|
||
code = "ERR_MASTER_KEY_EXISTS"
|
||
|
||
|
||
class MasterKeyInvalidError(Exception):
|
||
"""Raised when a supplied master key string cannot be used as a Fernet key."""
|
||
|
||
code = "ERR_MASTER_KEY_INVALID"
|
||
|
||
|
||
def master_key_path() -> str:
|
||
"""Absolute path to master.key alongside account-manager.db."""
|
||
return os.path.join(get_skill_data_dir(), _MASTER_FILENAME)
|
||
|
||
|
||
def _try_fernet_key(key_material: bytes) -> bytes | None:
|
||
"""Return url-safe key bytes accepted by Fernet, or None if invalid."""
|
||
try:
|
||
raw = key_material.strip()
|
||
if not raw:
|
||
return None
|
||
# Accept stored form: ascii urlsafe base64 string as bytes
|
||
Fernet(raw)
|
||
return raw
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _try_decode_manual_b64(s: str) -> bytes | None:
|
||
"""If string decodes to 32 raw bytes, re-encode as Fernet url-safe key."""
|
||
try:
|
||
pad = "=" * ((4 - len(s) % 4) % 4)
|
||
decoded = base64.urlsafe_b64decode(s + pad)
|
||
if len(decoded) != 32:
|
||
return None
|
||
key_bytes = base64.urlsafe_b64encode(decoded)
|
||
Fernet(key_bytes)
|
||
return key_bytes
|
||
except Exception:
|
||
return None
|
||
|
||
|
||
def _normalize_key_string(s: str) -> bytes | None:
|
||
s = (s or "").strip()
|
||
if not s:
|
||
return None
|
||
fk = _try_fernet_key(s.encode("ascii"))
|
||
if fk is not None:
|
||
return fk
|
||
return _try_decode_manual_b64(s)
|
||
|
||
|
||
def load_master_key() -> bytes:
|
||
"""Resolve master key (env → file); raises MasterKeyMissingError if unavailable."""
|
||
env_val = (os.environ.get(ENV_MASTER_KEY) or "").strip()
|
||
if env_val:
|
||
kb = _normalize_key_string(env_val)
|
||
if kb is not None:
|
||
return kb
|
||
|
||
path = master_key_path()
|
||
try:
|
||
with open(path, "rb") as f:
|
||
raw = f.read()
|
||
txt = raw.strip().decode("ascii", errors="strict")
|
||
kb = _normalize_key_string(txt)
|
||
if kb is not None:
|
||
return kb
|
||
except OSError:
|
||
pass
|
||
except (UnicodeDecodeError, ValueError):
|
||
pass
|
||
|
||
raise MasterKeyMissingError("No valid ACCOUNT_MANAGER_MASTER_KEY or master.key")
|
||
|
||
|
||
def normalize_master_key_input(s: str) -> bytes:
|
||
"""Validate user-supplied key text; raises MasterKeyInvalidError if unusable."""
|
||
kb = _normalize_key_string((s or "").strip())
|
||
if kb is None:
|
||
raise MasterKeyInvalidError("invalid master key material")
|
||
return kb
|
||
|
||
|
||
def generate_master_key() -> bytes:
|
||
"""Return a new Fernet key (44-char url-safe base64 as bytes)."""
|
||
return Fernet.generate_key()
|
||
|
||
|
||
def write_master_key_file(key: bytes, *, allow_overwrite: bool = False) -> str:
|
||
"""Write key bytes to master_key_path(); chmod 0o600 best-effort."""
|
||
path = master_key_path()
|
||
os.makedirs(os.path.dirname(path), exist_ok=True)
|
||
if os.path.exists(path) and not allow_overwrite:
|
||
raise MasterKeyExistsError("master.key already exists; use --rotate to overwrite")
|
||
data = key.strip()
|
||
with open(path, "wb") as f:
|
||
f.write(data)
|
||
try:
|
||
os.chmod(path, 0o600)
|
||
except OSError:
|
||
pass
|
||
return path
|
||
|
||
|
||
def is_master_key_available() -> bool:
|
||
try:
|
||
load_master_key()
|
||
return True
|
||
except MasterKeyMissingError:
|
||
return False
|