113 lines
4.5 KiB
Python
113 lines
4.5 KiB
Python
"""Windows 当前用户范围 DPAPI 封装；生产环境绝不降级为明文。"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import ctypes
|
|
from ctypes import wintypes
|
|
import os
|
|
import re
|
|
from typing import Protocol
|
|
|
|
from cmbuyer_client.core.errors import ProtectionError
|
|
|
|
|
|
class SecretProtector(Protocol):
|
|
def protect(self, plaintext: bytes, *, purpose: str) -> bytes: ...
|
|
|
|
def unprotect(self, ciphertext: bytes, *, purpose: str) -> bytes: ...
|
|
|
|
|
|
class _DataBlob(ctypes.Structure):
|
|
_fields_ = (("cbData", wintypes.DWORD), ("pbData", ctypes.POINTER(ctypes.c_ubyte)))
|
|
|
|
|
|
def _blob(data: bytes) -> tuple[_DataBlob, object]:
|
|
buffer = (ctypes.c_ubyte * len(data)).from_buffer_copy(data) if data else (ctypes.c_ubyte * 1)()
|
|
return _DataBlob(len(data), ctypes.cast(buffer, ctypes.POINTER(ctypes.c_ubyte))), buffer
|
|
|
|
|
|
class DpapiProtector:
|
|
"""使用 CryptProtectData/UI_FORBIDDEN；错误只暴露固定 reason code。"""
|
|
|
|
_UI_FORBIDDEN = 0x1
|
|
_ENTROPY_PREFIX = b"cmbuyer-localstate-v1:"
|
|
_PURPOSE_RE = re.compile(
|
|
r"(?:device-token:[a-z0-9][a-z0-9_-]{0,63}:[0-9a-f-]{36}|"
|
|
r"claim-token:[a-z0-9][a-z0-9_-]{0,63}:[0-9a-f-]{36})",
|
|
flags=re.ASCII,
|
|
)
|
|
|
|
def __init__(self) -> None:
|
|
if os.name != "nt":
|
|
raise ProtectionError("dpapi_requires_windows")
|
|
self._crypt32 = ctypes.WinDLL("crypt32", use_last_error=True)
|
|
self._kernel32 = ctypes.WinDLL("kernel32", use_last_error=True)
|
|
self._crypt32.CryptProtectData.argtypes = (
|
|
ctypes.POINTER(_DataBlob),
|
|
wintypes.LPCWSTR,
|
|
ctypes.POINTER(_DataBlob),
|
|
wintypes.LPVOID,
|
|
wintypes.LPVOID,
|
|
wintypes.DWORD,
|
|
ctypes.POINTER(_DataBlob),
|
|
)
|
|
self._crypt32.CryptProtectData.restype = wintypes.BOOL
|
|
self._crypt32.CryptUnprotectData.argtypes = (
|
|
ctypes.POINTER(_DataBlob),
|
|
ctypes.POINTER(wintypes.LPWSTR),
|
|
ctypes.POINTER(_DataBlob),
|
|
wintypes.LPVOID,
|
|
wintypes.LPVOID,
|
|
wintypes.DWORD,
|
|
ctypes.POINTER(_DataBlob),
|
|
)
|
|
self._crypt32.CryptUnprotectData.restype = wintypes.BOOL
|
|
self._kernel32.LocalFree.argtypes = (wintypes.HLOCAL,)
|
|
self._kernel32.LocalFree.restype = wintypes.HLOCAL
|
|
|
|
def protect(self, plaintext: bytes, *, purpose: str) -> bytes:
|
|
if not isinstance(plaintext, bytes) or not plaintext:
|
|
raise ProtectionError("invalid_plaintext")
|
|
entropy = self._entropy(purpose)
|
|
source, source_buffer = _blob(plaintext)
|
|
entropy_blob, entropy_buffer = _blob(entropy)
|
|
output = _DataBlob()
|
|
if not self._crypt32.CryptProtectData(
|
|
ctypes.byref(source), None, ctypes.byref(entropy_blob), None, None, self._UI_FORBIDDEN, ctypes.byref(output)
|
|
):
|
|
raise ProtectionError("dpapi_protect_failed")
|
|
# ctypes 指针不持有底层 Python buffer；局部引用必须活到系统调用返回。
|
|
del source_buffer, entropy_buffer
|
|
return self._take_output(output, "dpapi_protect_failed")
|
|
|
|
def unprotect(self, ciphertext: bytes, *, purpose: str) -> bytes:
|
|
if not isinstance(ciphertext, bytes) or not ciphertext:
|
|
raise ProtectionError("invalid_ciphertext")
|
|
entropy = self._entropy(purpose)
|
|
source, source_buffer = _blob(ciphertext)
|
|
entropy_blob, entropy_buffer = _blob(entropy)
|
|
output = _DataBlob()
|
|
description = wintypes.LPWSTR()
|
|
if not self._crypt32.CryptUnprotectData(
|
|
ctypes.byref(source), ctypes.byref(description), ctypes.byref(entropy_blob), None, None, self._UI_FORBIDDEN, ctypes.byref(output)
|
|
):
|
|
raise ProtectionError("dpapi_unprotect_failed")
|
|
del source_buffer, entropy_buffer
|
|
if description:
|
|
self._kernel32.LocalFree(ctypes.cast(description, wintypes.HLOCAL))
|
|
return self._take_output(output, "dpapi_unprotect_failed")
|
|
|
|
def _take_output(self, output: _DataBlob, reason: str) -> bytes:
|
|
if not output.pbData or output.cbData <= 0:
|
|
raise ProtectionError(reason)
|
|
try:
|
|
return ctypes.string_at(output.pbData, output.cbData)
|
|
finally:
|
|
self._kernel32.LocalFree(ctypes.cast(output.pbData, wintypes.HLOCAL))
|
|
|
|
@classmethod
|
|
def _entropy(cls, purpose: str) -> bytes:
|
|
if not isinstance(purpose, str) or cls._PURPOSE_RE.fullmatch(purpose) is None:
|
|
raise ProtectionError("invalid_protection_purpose")
|
|
return cls._ENTROPY_PREFIX + purpose.encode("ascii")
|