feat(#95): 安全模块 - RBAC角色权限+JWT认证+AES-256-GCM数据加密+安全审计哈希链+安全中间件(21个单元测试)
This commit is contained in:
@@ -0,0 +1 @@
|
||||
tests/security/__pycache__/
|
||||
@@ -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.
@@ -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)
|
||||
@@ -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
|
||||
@@ -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")
|
||||
@@ -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})
|
||||
@@ -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()
|
||||
Reference in New Issue
Block a user