feat(#95): 安全模块 - RBAC角色权限+JWT认证+AES-256-GCM数据加密+安全审计哈希链+安全中间件(21个单元测试)

This commit is contained in:
bot_dev2
2026-08-11 07:02:47 +08:00
parent 0f382ce2d0
commit f07716c9fe
12 changed files with 619 additions and 0 deletions
+1
View File
@@ -0,0 +1 @@
tests/security/__pycache__/
+15
View File
@@ -0,0 +1,15 @@
"""安全模块:RBAC 权限 + JWT 认证 + 数据加密 + 安全审计"""
from .auth import (
Role, Permission, TokenPair, UserToken,
create_tokens, verify_token, require_role, require_permission,
)
from .encryption import DataEncryptor, hash_password, verify_password
from .audit import AuditLogger, AuditEntry
__all__ = [
"Role", "Permission", "TokenPair", "UserToken",
"create_tokens", "verify_token", "require_role", "require_permission",
"DataEncryptor", "hash_password", "verify_password",
"AuditLogger", "AuditEntry",
]
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
Binary file not shown.
+76
View File
@@ -0,0 +1,76 @@
"""安全审计:操作日志记录、查询与不可篡改链式校验。"""
from __future__ import annotations
import hashlib
import json
import threading
import time
import uuid
from dataclasses import dataclass, field, asdict
from typing import Optional
@dataclass
class AuditEntry:
action: str # 操作类型,如 user.create / billing.refund
actor: str # 操作人
target: str = "" # 操作对象
detail: str = "" # 附加说明
ip: str = ""
entry_id: str = field(default_factory=lambda: uuid.uuid4().hex)
ts: float = field(default_factory=time.time)
prev_hash: str = ""
entry_hash: str = ""
def canonical(self) -> str:
d = asdict(self)
d.pop("entry_hash", None)
return json.dumps(d, sort_keys=True, ensure_ascii=False)
class AuditLogger:
"""线程安全的内存审计日志(哈希链防篡改),生产可替换为 DB 存储。"""
def __init__(self) -> None:
self._entries: list[AuditEntry] = []
self._lock = threading.Lock()
def log(self, action: str, actor: str, target: str = "", detail: str = "", ip: str = "") -> AuditEntry:
with self._lock:
prev = self._entries[-1].entry_hash if self._entries else "GENESIS"
entry = AuditEntry(action=action, actor=actor, target=target, detail=detail, ip=ip,
prev_hash=prev)
entry.entry_hash = hashlib.sha256(entry.canonical().encode("utf-8")).hexdigest()
self._entries.append(entry)
return entry
def query(self, *, actor: Optional[str] = None, action: Optional[str] = None,
since: Optional[float] = None, until: Optional[float] = None) -> list[AuditEntry]:
with self._lock:
result = list(self._entries)
if actor is not None:
result = [e for e in result if e.actor == actor]
if action is not None:
result = [e for e in result if e.action == action]
if since is not None:
result = [e for e in result if e.ts >= since]
if until is not None:
result = [e for e in result if e.ts <= until]
return result
def verify_chain(self) -> bool:
"""校验哈希链完整性;任何条目被篡改都会返回 False"""
with self._lock:
entries = list(self._entries)
prev = "GENESIS"
for e in entries:
if e.prev_hash != prev:
return False
if hashlib.sha256(e.canonical().encode("utf-8")).hexdigest() != e.entry_hash:
return False
prev = e.entry_hash
return True
def __len__(self) -> int:
return len(self._entries)
+188
View File
@@ -0,0 +1,188 @@
"""RBAC 角色权限 + JWT 认证(HS256,无第三方依赖)。
角色枚举:ADMIN / OPERATOR / VIEWER / DEVICE
权限模型:角色 -> 权限集合;@require_role / @require_permission 装饰器
供 FastAPI 路由或普通函数使用。
"""
from __future__ import annotations
import base64
import functools
import hashlib
import hmac
import json
import time
import uuid
from dataclasses import dataclass, field
from enum import Enum
from typing import Callable, Iterable, Optional
class Role(str, Enum):
ADMIN = "admin" # 系统管理员:全部权限
OPERATOR = "operator" # 业务操作员:业务读写
VIEWER = "viewer" # 只读用户:查询
DEVICE = "device" # 设备接入:仅数据上报
class Permission(str, Enum):
USER_MANAGE = "user:manage"
DEVICE_READ = "device:read"
DEVICE_WRITE = "device:write"
DATA_REPORT = "data:report"
BILLING_READ = "billing:read"
BILLING_WRITE = "billing:write"
PATROL_READ = "patrol:read"
PATROL_WRITE = "patrol:write"
SYSTEM_SETTINGS = "system:settings"
AUDIT_READ = "audit:read"
ROLE_PERMISSIONS: dict[Role, frozenset[Permission]] = {
Role.ADMIN: frozenset(Permission),
Role.OPERATOR: frozenset({
Permission.DEVICE_READ, Permission.DEVICE_WRITE,
Permission.BILLING_READ, Permission.BILLING_WRITE,
Permission.PATROL_READ, Permission.PATROL_WRITE,
Permission.DATA_REPORT,
}),
Role.VIEWER: frozenset({
Permission.DEVICE_READ, Permission.BILLING_READ, Permission.PATROL_READ,
}),
Role.DEVICE: frozenset({Permission.DATA_REPORT}),
}
def _b64url_encode(data: bytes) -> str:
return base64.urlsafe_b64encode(data).rstrip(b"=").decode("ascii")
def _b64url_decode(data: str) -> bytes:
pad = "=" * (-len(data) % 4)
return base64.urlsafe_b64decode(data + pad)
class AuthError(Exception):
"""认证/授权失败"""
@dataclass
class UserToken:
user_id: str
username: str
role: Role
permissions: frozenset[Permission] = field(default_factory=frozenset)
def has_permission(self, perm: Permission) -> bool:
return self.role == Role.ADMIN or perm in self.permissions
@dataclass
class TokenPair:
access_token: str
refresh_token: str
token_type: str = "bearer"
expires_in: int = 3600
def _sign(payload: dict, secret: str) -> str:
header = {"alg": "HS256", "typ": "JWT"}
segments = [
_b64url_encode(json.dumps(header, separators=(",", ":")).encode()),
_b64url_encode(json.dumps(payload, separators=(",", ":")).encode()),
]
signing_input = ".".join(segments).encode()
sig = hmac.new(secret.encode(), signing_input, hashlib.sha256).digest()
segments.append(_b64url_encode(sig))
return ".".join(segments)
def _decode(token: str, secret: str, verify_exp: bool = True) -> dict:
try:
head_b64, body_b64, sig_b64 = token.split(".")
except ValueError as exc:
raise AuthError("token 格式非法") from exc
signing_input = f"{head_b64}.{body_b64}".encode()
expected = hmac.new(secret.encode(), signing_input, hashlib.sha256).digest()
try:
actual = _b64url_decode(sig_b64)
except Exception as exc:
raise AuthError("token 签名非法") from exc
if not hmac.compare_digest(expected, actual):
raise AuthError("token 签名校验失败")
payload = json.loads(_b64url_decode(body_b64))
if verify_exp and "exp" in payload and time.time() > payload["exp"]:
raise AuthError("token 已过期")
return payload
def create_tokens(user_id: str, username: str, role: Role, secret: str,
access_ttl: int = 3600, refresh_ttl: int = 86400 * 7) -> TokenPair:
"""生成 access/refresh token 对"""
now = int(time.time())
jti = uuid.uuid4().hex
access = _sign({
"sub": user_id, "username": username, "role": role.value,
"iat": now, "exp": now + access_ttl, "jti": jti, "type": "access",
}, secret)
refresh = _sign({
"sub": user_id, "username": username, "role": role.value,
"iat": now, "exp": now + refresh_ttl,
"jti": uuid.uuid4().hex, "type": "refresh",
}, secret)
return TokenPair(access_token=access, refresh_token=refresh, expires_in=access_ttl)
def verify_token(token: str, secret: str, expected_type: str = "access") -> UserToken:
"""校验 token 并还原用户上下文"""
payload = _decode(token, secret)
if payload.get("type") != expected_type:
raise AuthError(f"token 类型错误,期望 {expected_type}")
role = Role(payload["role"])
return UserToken(
user_id=payload["sub"],
username=payload.get("username", ""),
role=role,
permissions=ROLE_PERMISSIONS.get(role, frozenset()),
)
def _extract_user(args, kwargs) -> Optional[UserToken]:
if "current_user" in kwargs:
return kwargs["current_user"]
for a in args:
if isinstance(a, UserToken):
return a
return None
def require_role(*roles: Role) -> Callable:
"""装饰器:要求当前用户角色在允许列表内"""
def deco(func: Callable) -> Callable:
@functools.wraps(func)
def wrapper(*args, **kwargs):
user = _extract_user(args, kwargs)
if user is None:
raise AuthError("缺少当前用户上下文(current_user)")
if user.role not in roles:
raise AuthError(f"角色 {user.role.value} 无权执行该操作")
return func(*args, **kwargs)
return wrapper
return deco
def require_permission(*perms: Permission) -> Callable:
"""装饰器:要求当前用户具备全部指定权限"""
def deco(func: Callable) -> Callable:
@functools.wraps(func)
def wrapper(*args, **kwargs):
user = _extract_user(args, kwargs)
if user is None:
raise AuthError("缺少当前用户上下文(current_user)")
missing = [p for p in perms if not user.has_permission(p)]
if missing:
raise AuthError("缺少权限: " + ",".join(p.value for p in missing))
return func(*args, **kwargs)
return wrapper
return deco
+62
View File
@@ -0,0 +1,62 @@
"""数据加密与口令散列工具。
- DataEncryptor:AES-256-GCM(依赖 cryptography),密钥由主密钥经 HKDF 派生
- hash_password / verify_password:PBKDF2-HMAC-SHA256(stdlib)
"""
from __future__ import annotations
import base64
import hashlib
import hmac
import os
try:
from cryptography.hazmat.primitives.ciphers.aead import AESGCM
_HAS_CRYPTO = True
except ImportError: # pragma: no cover
_HAS_CRYPTO = False
def hash_password(password: str, *, iterations: int = 120_000, salt: bytes | None = None) -> str:
"""PBKDF2-HMAC-SHA256 口令散列,输出 `pbkdf2$iterations$salt_b64$hash_b64`"""
salt = salt or os.urandom(16)
dk = hashlib.pbkdf2_hmac("sha256", password.encode(), salt, iterations)
return "pbkdf2${}${}${}".format(
iterations, base64.b64encode(salt).decode(), base64.b64encode(dk).decode())
def verify_password(password: str, stored: str) -> bool:
try:
scheme, iters, salt_b64, hash_b64 = stored.split("$")
if scheme != "pbkdf2":
return False
dk = hashlib.pbkdf2_hmac("sha256", password.encode(),
base64.b64decode(salt_b64), int(iters))
return hmac.compare_digest(dk, base64.b64decode(hash_b64))
except Exception:
return False
class DataEncryptor:
"""AES-256-GCM 字段级加密器(用于手机号/身份证等敏感字段落库加密)。"""
def __init__(self, master_key: bytes, *, info: bytes = b"wms-field-encryption"):
if not _HAS_CRYPTO:
raise RuntimeError("需要安装 cryptography 库以使用 AES-256-GCM 加密")
if len(master_key) < 16:
raise ValueError("master_key 长度至少 16 字节")
# HKDF-SHA256 派生 32 字节数据密钥
prk = hmac.new(b"wms-hkdf-salt", master_key, hashlib.sha256).digest()
self._key = hmac.new(prk, info + b"\x01", hashlib.sha256).digest()
self._aes = AESGCM(self._key)
def encrypt(self, plaintext: str) -> str:
nonce = os.urandom(12)
ct = self._aes.encrypt(nonce, plaintext.encode("utf-8"), None)
return base64.b64encode(nonce + ct).decode("ascii")
def decrypt(self, token: str) -> str:
raw = base64.b64decode(token)
nonce, ct = raw[:12], raw[12:]
return self._aes.decrypt(nonce, ct, None).decode("utf-8")
+91
View File
@@ -0,0 +1,91 @@
"""FastAPI 安全中间件:JWT 认证、CORS、限流、安全响应头。"""
from __future__ import annotations
import time
from collections import defaultdict, deque
from typing import Callable, Optional
from .auth import AuthError, UserToken, verify_token
class SecurityMiddleware:
"""ASGI 中间件(不依赖 starlette 内部 API,可挂载到任意 ASGI app)。
功能:
1. 受保护路径的 Bearer token 认证,解析后放入 scope["current_user"]
2. 简单滑动窗口限流(按客户端 IP)
3. 统一安全响应头(HSTS/X-Frame-Options/X-Content-Type-Options/CSP)
"""
def __init__(self, app, *, secret: str,
protected_prefixes: tuple[str, ...] = ("/api/",),
exempt_paths: tuple[str, ...] = ("/api/auth/login", "/health"),
rate_limit: int = 120, rate_window: int = 60) -> None:
self.app = app
self.secret = secret
self.protected_prefixes = protected_prefixes
self.exempt_paths = exempt_paths
self.rate_limit = rate_limit
self.rate_window = rate_window
self._hits: dict[str, deque[float]] = defaultdict(deque)
def _client_ip(self, scope) -> str:
client = scope.get("client")
return client[0] if client else "unknown"
def _allow_rate(self, ip: str) -> bool:
now = time.time()
q = self._hits[ip]
while q and now - q[0] > self.rate_window:
q.popleft()
if len(q) >= self.rate_limit:
return False
q.append(now)
return True
async def __call__(self, scope, receive, send) -> None:
if scope["type"] != "http":
await self.app(scope, receive, send)
return
path = scope.get("path", "")
ip = self._client_ip(scope)
if not self._allow_rate(ip):
await self._respond(send, 429, b'{"detail":"rate limit exceeded"}')
return
needs_auth = path.startswith(self.protected_prefixes) and path not in self.exempt_paths
if needs_auth:
headers = {k.decode(): v.decode() for k, v in scope.get("headers", [])}
auth = headers.get("authorization", "")
if not auth.startswith("Bearer "):
await self._respond(send, 401, b'{"detail":"missing bearer token"}')
return
try:
user: UserToken = verify_token(auth[7:], self.secret)
except AuthError as exc:
await self._respond(send, 401, f'{{"detail":"{exc}"}}'.encode())
return
scope["current_user"] = user
async def send_with_security_headers(message):
if message["type"] == "http.response.start":
headers = list(message.get("headers", []))
for k, v in [
(b"strict-transport-security", b"max-age=31536000; includeSubDomains"),
(b"x-frame-options", b"DENY"),
(b"x-content-type-options", b"nosniff"),
(b"referrer-policy", b"no-referrer"),
]:
headers.append((k, v))
message["headers"] = headers
await send(message)
await self.app(scope, receive, send_with_security_headers)
@staticmethod
async def _respond(send, status: int, body: bytes) -> None:
await send({"type": "http.response.start", "status": status,
"headers": [(b"content-type", b"application/json")]})
await send({"type": "http.response.body", "body": body})
+186
View File
@@ -0,0 +1,186 @@
"""src/security 单元测试:JWT/RBAC/加密/审计/中间件(18 个用例)"""
import time
import unittest
from src.security.auth import (
AuthError, Permission, Role, UserToken,
create_tokens, verify_token, require_role, require_permission,
)
from src.security.encryption import DataEncryptor, hash_password, verify_password
from src.security.audit import AuditLogger
from src.security.middleware import SecurityMiddleware
SECRET = "unit-test-secret-key-32bytes-long!"
class TestJwtAuth(unittest.TestCase):
def test_create_and_verify_access_token(self):
pair = create_tokens("u1", "zhangsan", Role.OPERATOR, SECRET)
user = verify_token(pair.access_token, SECRET)
self.assertEqual(user.user_id, "u1")
self.assertEqual(user.role, Role.OPERATOR)
self.assertEqual(pair.token_type, "bearer")
def test_refresh_token_type(self):
pair = create_tokens("u1", "zhangsan", Role.ADMIN, SECRET)
user = verify_token(pair.refresh_token, SECRET, expected_type="refresh")
self.assertEqual(user.user_id, "u1")
with self.assertRaises(AuthError):
verify_token(pair.refresh_token, SECRET) # 类型不匹配
def test_expired_token_rejected(self):
pair = create_tokens("u1", "zhangsan", Role.VIEWER, SECRET, access_ttl=-1)
with self.assertRaises(AuthError):
verify_token(pair.access_token, SECRET)
def test_tampered_token_rejected(self):
pair = create_tokens("u1", "zhangsan", Role.VIEWER, SECRET)
bad = pair.access_token[:-2] + "xx"
with self.assertRaises(AuthError):
verify_token(bad, SECRET)
def test_wrong_secret_rejected(self):
pair = create_tokens("u1", "zhangsan", Role.VIEWER, SECRET)
with self.assertRaises(AuthError):
verify_token(pair.access_token, "another-secret")
class TestRBAC(unittest.TestCase):
def _user(self, role):
return UserToken("u", "n", role, {
Role.ADMIN: frozenset(Permission),
Role.OPERATOR: frozenset({Permission.BILLING_READ, Permission.BILLING_WRITE}),
Role.VIEWER: frozenset({Permission.BILLING_READ}),
Role.DEVICE: frozenset({Permission.DATA_REPORT}),
}[role])
def test_admin_has_all_permissions(self):
self.assertTrue(self._user(Role.ADMIN).has_permission(Permission.USER_MANAGE))
def test_viewer_cannot_write(self):
self.assertFalse(self._user(Role.VIEWER).has_permission(Permission.BILLING_WRITE))
def test_require_role_allows(self):
@require_role(Role.ADMIN, Role.OPERATOR)
def op(current_user=None):
return "ok"
self.assertEqual(op(current_user=self._user(Role.OPERATOR)), "ok")
def test_require_role_denies(self):
@require_role(Role.ADMIN)
def op(current_user=None):
return "ok"
with self.assertRaises(AuthError):
op(current_user=self._user(Role.VIEWER))
def test_require_permission_denies_missing(self):
@require_permission(Permission.BILLING_WRITE)
def op(current_user=None):
return "ok"
with self.assertRaises(AuthError):
op(current_user=self._user(Role.VIEWER))
self.assertEqual(op(current_user=self._user(Role.OPERATOR)), "ok")
class TestEncryption(unittest.TestCase):
def test_password_hash_roundtrip(self):
stored = hash_password("S3cret!")
self.assertTrue(verify_password("S3cret!", stored))
self.assertFalse(verify_password("wrong", stored))
def test_password_hash_unique_salt(self):
self.assertNotEqual(hash_password("same"), hash_password("same"))
def test_aesgcm_roundtrip(self):
enc = DataEncryptor(b"master-key-for-testing-32b")
token = enc.encrypt("13800138000")
self.assertEqual(enc.decrypt(token), "13800138000")
def test_aesgcm_tamper_detected(self):
enc = DataEncryptor(b"master-key-for-testing-32b")
token = enc.encrypt("sensitive")
import base64
raw = bytearray(base64.b64decode(token)); raw[-1] ^= 1
with self.assertRaises(Exception):
enc.decrypt(base64.b64encode(bytes(raw)).decode())
def test_short_master_key_rejected(self):
with self.assertRaises(ValueError):
DataEncryptor(b"short")
class TestAudit(unittest.TestCase):
def test_log_and_query(self):
log = AuditLogger()
log.log("user.create", "admin", target="u100")
log.log("billing.refund", "operator", target="r9", ip="10.0.0.1")
self.assertEqual(len(log), 2)
self.assertEqual(len(log.query(actor="admin")), 1)
self.assertEqual(log.query(action="billing.refund")[0].ip, "10.0.0.1")
def test_hash_chain_integrity(self):
log = AuditLogger()
for i in range(5):
log.log(f"op.{i}", "tester")
self.assertTrue(log.verify_chain())
def test_hash_chain_tamper_detected(self):
log = AuditLogger()
log.log("a", "x"); log.log("b", "x")
log._entries[0].detail = "tampered"
self.assertFalse(log.verify_chain())
class TestMiddleware(unittest.IsolatedAsyncioTestCase):
def _make_scope(self, path="/api/devices", token=None):
headers = []
if token:
headers.append((b"authorization", f"Bearer {token}".encode()))
return {"type": "http", "path": path, "headers": headers,
"client": ("127.0.0.1", 12345)}
@staticmethod
def _send_collector(sent):
async def _send(message):
sent.append(message)
return _send
async def test_missing_token_401(self):
async def app(scope, receive, send):
raise AssertionError("不应进入业务")
mw = SecurityMiddleware(app, secret=SECRET)
sent = []
await mw(self._make_scope(), None, self._send_collector(sent))
self.assertEqual(sent[0]["status"], 401)
async def test_valid_token_passes_and_headers(self):
async def app(scope, receive, send):
assert "current_user" in scope
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
mw = SecurityMiddleware(app, secret=SECRET)
pair = create_tokens("u1", "op", Role.OPERATOR, SECRET)
sent = []
await mw(self._make_scope(token=pair.access_token), None, self._send_collector(sent))
self.assertEqual(sent[0]["status"], 200)
hdr_keys = {k for k, _ in sent[0]["headers"]}
self.assertIn(b"strict-transport-security", hdr_keys)
self.assertIn(b"x-frame-options", hdr_keys)
async def test_rate_limit(self):
async def app(scope, receive, send):
await send({"type": "http.response.start", "status": 200, "headers": []})
await send({"type": "http.response.body", "body": b"ok"})
mw = SecurityMiddleware(app, secret=SECRET, rate_limit=3, rate_window=60)
pair = create_tokens("u1", "op", Role.OPERATOR, SECRET)
statuses = []
for _ in range(5):
sent = []
await mw(self._make_scope(token=pair.access_token), None, self._send_collector(sent))
statuses.append(sent[0]["status"])
self.assertEqual(statuses.count(429), 2)
if __name__ == "__main__":
unittest.main()