workbuddy-portal/workbuddy_portal/crypto.py

# -*- coding: utf-8 -*-
# SPDX-License-Identifier: MIT
# Copyright (c) 2026 Wang Chuanli

"""对称加密 —— 给「数据库里的 Cookie」做静态加密。

为什么要自己写而不用 `cryptography` / `pycryptodome`
----------------------------------------------------
本项目刻意保持零第三方依赖（`requirements.txt` 只有 Flask / waitress / openpyxl），
而这里需要的原语只有两个，都能在 RFC 里逐行对照实现：

  * **ChaCha20** 流密码（RFC 8439 §2.3）—— 加密
  * **HMAC-SHA256**（RFC 2104）—— 认证，采用 **encrypt-then-MAC**

明文密钥不是口令而是 32 字节随机数，所以派生不需要慢速 KDF
（PBKDF2/scrypt 是为「低熵口令」设计的），用 HMAC 做一次密钥分离即可。

密文格式
--------
    v1.<b64(salt)>.<b64(nonce)>.<b64(ciphertext)>.<b64(tag)>

  * salt  —— 16 字节随机，用于把主密钥分离成 enc/mac 两把子密钥
  * nonce —— 12 字节随机，每次加密都重新生成（绝不复用）
  * tag   —— HMAC(mac_key, nonce || ciphertext) 的 SHA-256

**不做压缩**：Cookie 是几百到几千字节的高熵串，压缩比接近 1，
反而会引入 CRIME 类侧信道，不值得。

向后兼容
--------
`decrypt()` 遇到不是 `v1.` 开头的值会**原样返回**，这样从旧版本
（Cookie 明文存在 settings 表）升级过来不会立刻炸；下次写入时自然
会被改写为密文（见 `db.set_secret`）。
"""

import base64
import hashlib
import hmac
import secrets
import struct

PREFIX = "v1."

# RFC 8439 §2.3 的常数："expand 32-byte k"
_CONST = b"expand 32-byte k"
_MASK = 0xFFFFFFFF


# ---------------- ChaCha20 ----------------
def _rotl32(v, c):
    return ((v << c) & _MASK) | (v >> (32 - c))


def _quarter_round(s, a, b, c, d):
    """RFC 8439 §2.1。就地修改 s。"""
    s[a] = (s[a] + s[b]) & _MASK
    s[d] = _rotl32(s[d] ^ s[a], 16)
    s[c] = (s[c] + s[d]) & _MASK
    s[b] = _rotl32(s[b] ^ s[c], 12)
    s[a] = (s[a] + s[b]) & _MASK
    s[d] = _rotl32(s[d] ^ s[a], 8)
    s[c] = (s[c] + s[d]) & _MASK
    s[b] = _rotl32(s[b] ^ s[c], 7)


def chacha20_block(key32, counter, nonce12):
    """产出一个 64 字节的块（RFC 8439 §2.3.2）。"""
    st = (list(struct.unpack("<4I", _CONST))
          + list(struct.unpack("<8I", key32))
          + [counter & _MASK]
          + list(struct.unpack("<3I", nonce12)))
    w = list(st)
    for _ in range(10):                       # 10 组 = 20 轮
        _quarter_round(w, 0, 4, 8, 12)
        _quarter_round(w, 1, 5, 9, 13)
        _quarter_round(w, 2, 6, 10, 14)
        _quarter_round(w, 3, 7, 11, 15)
        _quarter_round(w, 0, 5, 10, 15)
        _quarter_round(w, 1, 6, 11, 12)
        _quarter_round(w, 2, 7, 8, 13)
        _quarter_round(w, 3, 4, 9, 14)
    return struct.pack("<16I", *[(w[i] + st[i]) & _MASK for i in range(16)])


def _keystream(key32, nonce12, n):
    """按需生成 n 字节密钥流。counter 从 1 开始（0 号块留给 Poly1305 用，这里不用）。"""
    out = bytearray()
    counter = 1
    while len(out) < n:
        out += chacha20_block(key32, counter, nonce12)
        counter += 1
    return bytes(out[:n])


def _xor(a, b):
    return bytes(x ^ y for x, y in zip(a, b))


# ---------------- 密钥分离 ----------------
def _derive(master, salt, label):
    """HMAC 做一次密钥分离：主密钥是高熵随机数，一次 HMAC 足够。"""
    return hmac.new(master, salt + label, hashlib.sha256).digest()


def _b64(raw):
    return base64.urlsafe_b64encode(raw).decode("ascii").rstrip("=")


def _unb64(text):
    pad = "=" * (-len(text) % 4)
    return base64.urlsafe_b64decode(text + pad)


# ---------------- 对外接口 ----------------
def is_encrypted(value):
    return isinstance(value, str) and value.startswith(PREFIX)


def encrypt(plaintext, master):
    """加密任意字符串；空值原样返回（不产生「有密文的空值」这种歧义状态）。

    master: 32 字节主密钥（bytes）。返回可直接存库的 ASCII 字符串。
    """
    if plaintext is None or plaintext == "":
        return ""
    if not isinstance(master, (bytes, bytearray)) or len(master) != 32:
        raise ValueError("主密钥必须是 32 字节")
    data = plaintext.encode("utf-8") if isinstance(plaintext, str) else bytes(plaintext)
    salt = secrets.token_bytes(16)
    nonce = secrets.token_bytes(12)
    enc_key = _derive(bytes(master), salt, b"enc")
    mac_key = _derive(bytes(master), salt, b"mac")
    ct = _xor(data, _keystream(enc_key, nonce, len(data)))
    tag = hmac.new(mac_key, nonce + ct, hashlib.sha256).digest()
    return PREFIX + ".".join((_b64(salt), _b64(nonce), _b64(ct), _b64(tag)))


class DecryptError(ValueError):
    """密文被篡改、格式损坏或密钥不对。"""


def decrypt(token, master):
    """解密。

    * 非密文（历史明文、空串）原样返回，便于平滑升级
    * 密文校验失败抛 DecryptError —— **绝不**「失败就返回原值」，
      否则一次篡改会被静默当成合法明文用下去
    """
    if not token or not isinstance(token, str):
        return ""
    if not token.startswith(PREFIX):
        return token                            # 兼容旧的明文存储
    if not isinstance(master, (bytes, bytearray)) or len(master) != 32:
        raise DecryptError("主密钥必须是 32 字节")
    parts = token[len(PREFIX):].split(".")
    if len(parts) != 4:
        raise DecryptError("密文格式不正确")
    try:
        salt, nonce, ct, tag = (_unb64(p) for p in parts)
    except (ValueError, TypeError) as e:
        raise DecryptError("密文 base64 解码失败：%s" % e)
    if len(salt) != 16 or len(nonce) != 12 or len(tag) != 32:
        raise DecryptError("密文长度不合法")
    mac_key = _derive(bytes(master), salt, b"mac")
    want = hmac.new(mac_key, nonce + ct, hashlib.sha256).digest()
    # 先比 MAC 再解密：认证失败时不接触密文，避免 padding/解析类侧信道
    if not hmac.compare_digest(want, tag):
        raise DecryptError("完整性校验失败（密文被篡改或主密钥已更换）")
    enc_key = _derive(bytes(master), salt, b"enc")
    return _xor(ct, _keystream(enc_key, nonce, len(ct))).decode("utf-8")


def fingerprint(plaintext):
    """值指纹：用于「是否换过」的判断，不能反推原文。"""
    if not plaintext:
        return ""
    return hashlib.sha256(plaintext.encode("utf-8")).hexdigest()[:16]