.gitignore

@@ -0,0 +1 @@
+tests/security/__pycache__/

src/security/__init__.py

@@ -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",
+]

src/security/audit.py

@@ -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)

src/security/auth.py

@@ -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

src/security/encryption.py

@@ -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")

src/security/middleware.py

@@ -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})

tests/security/test_security.py

@@ -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()