hms/apps/miniprogram/src/utils/secure-storage-aes.ts

/**
 * AES-256-GCM 加密存储 — 替代 XOR 加密
 * 依赖 @noble/ciphers + crypto-polyfill（wx.getRandomValuesSync）
 */
import Taro from '@tarojs/taro';
import { gcm } from '@noble/ciphers/aes.js';

const STORAGE_PREFIX = '_es_';
const AES_MARKER = 'aes:';
const NONCE_LENGTH = 12; // GCM 标准 nonce 长度

declare const wx: {
  getRandomValuesSync?: (params: { length: number }) => ArrayBuffer;
} | undefined;

// 小程序环境无 TextEncoder/TextDecoder，用纯 JS 替代
function utf8Encode(str: string): Uint8Array {
  const bytes: number[] = [];
  for (let i = 0; i < str.length; i++) {
    let code = str.charCodeAt(i);
    if (code < 0x80) {
      bytes.push(code);
    } else if (code < 0x800) {
      bytes.push(0xc0 | (code >> 6), 0x80 | (code & 0x3f));
    } else if (code >= 0xd800 && code <= 0xdbff) {
      // surrogate pair
      const hi = code;
      const lo = str.charCodeAt(++i);
      code = ((hi - 0xd800) << 10) + (lo - 0xdc00) + 0x10000;
      bytes.push(0xf0 | (code >> 18), 0x80 | ((code >> 12) & 0x3f), 0x80 | ((code >> 6) & 0x3f), 0x80 | (code & 0x3f));
    } else {
      bytes.push(0xe0 | (code >> 12), 0x80 | ((code >> 6) & 0x3f), 0x80 | (code & 0x3f));
    }
  }
  return new Uint8Array(bytes);
}

function utf8Decode(bytes: Uint8Array): string {
  let str = '';
  let i = 0;
  while (i < bytes.length) {
    const b = bytes[i++];
    if (b < 0x80) {
      str += String.fromCharCode(b);
    } else if (b < 0xe0) {
      str += String.fromCharCode(((b & 0x1f) << 6) | (bytes[i++] & 0x3f));
    } else if (b < 0xf0) {
      str += String.fromCharCode(((b & 0x0f) << 12) | ((bytes[i++] & 0x3f) << 6) | (bytes[i++] & 0x3f));
    } else {
      const cp = ((b & 0x07) << 18) | ((bytes[i++] & 0x3f) << 12) | ((bytes[i++] & 0x3f) << 6) | (bytes[i++] & 0x3f);
      const hi = ((cp - 0x10000) >> 10) + 0xd800;
      const lo = ((cp - 0x10000) & 0x3ff) + 0xdc00;
      str += String.fromCharCode(hi, lo);
    }
  }
  return str;
}

function getEncryptionKey(): Uint8Array | null {
  const hex = process.env.TARO_APP_ENCRYPTION_KEY || '';
  if (hex && /^[0-9a-fA-F]{64}$/.test(hex)) {
    return new Uint8Array(hex.match(/.{2}/g)!.map((b) => parseInt(b, 16)));
  }
  if (process.env.NODE_ENV === 'production') {
    console.error('[secure-storage] ENCRYPTION_KEY not configured in production');
    return null;
  }
  // dev: warn and use plaintext-compatible mode
  console.warn('[secure-storage] ENCRYPTION_KEY empty — using dev plaintext mode');
  return null;
}

function hasEncryptionKey(): boolean {
  return getEncryptionKey() !== null;
}

function generateNonce(): Uint8Array {
  const nonce = new Uint8Array(NONCE_LENGTH);
  if (typeof globalThis.crypto?.getRandomValues === 'function') {
    globalThis.crypto.getRandomValues(nonce);
  } else if (typeof wx !== 'undefined' && typeof wx.getRandomValuesSync === 'function') {
    const buf = wx.getRandomValuesSync({ length: NONCE_LENGTH });
    const src = new Uint8Array(buf as ArrayBuffer);
    nonce.set(src);
  } else {
    for (let i = 0; i < NONCE_LENGTH; i++) {
      nonce[i] = (Math.random() * 256) | 0;
    }
  }
  return nonce;
}

function aesEncrypt(plaintext: string): string | null {
  const key = getEncryptionKey();
  if (!key) return null;
  const nonce = generateNonce();
  const cipher = gcm(key, nonce);
  const data = utf8Encode(plaintext);
  const ciphertext = cipher.encrypt(data);
  const combined = new Uint8Array(nonce.length + ciphertext.length);
  combined.set(nonce, 0);
  combined.set(ciphertext, nonce.length);
  return AES_MARKER + Taro.arrayBufferToBase64(combined.buffer as ArrayBuffer);
}

function aesDecrypt(encoded: string): string | null {
  try {
    const key = getEncryptionKey();
    if (!key) return null;
    const b64 = encoded.slice(AES_MARKER.length);
    const buf = Taro.base64ToArrayBuffer(b64);
    const combined = new Uint8Array(buf);
    if (combined.length < NONCE_LENGTH) return null;

    const nonce = combined.slice(0, NONCE_LENGTH);
    const ciphertext = combined.slice(NONCE_LENGTH);
    const cipher = gcm(key, nonce);
    const plaintext = cipher.decrypt(ciphertext);
    return utf8Decode(plaintext);
  } catch {
    return null;
  }
}

function fromBase64(b64: string): string {
  try {
    const buffer = Taro.base64ToArrayBuffer(b64);
    return new TextDecoder().decode(new Uint8Array(buffer));
  } catch {
    return '';
  }
}

export function secureSet(key: string, value: string): void {
  if (!value) {
    Taro.removeStorageSync(STORAGE_PREFIX + key);
    return;
  }
  const encrypted = aesEncrypt(value);
  if (encrypted) {
    Taro.setStorageSync(STORAGE_PREFIX + key, encrypted);
    return;
  }
  // 密钥不可用时的降级策略
  if (process.env.NODE_ENV === 'production') {
    // 生产环境不允许明文存储敏感数据
    console.error(`[secure-storage] 拒绝明文写入 production key: ${key} — ENCRYPTION_KEY 未配置`);
    return;
  }
  // dev mode: store plaintext with prefix for compatibility
  Taro.setStorageSync(STORAGE_PREFIX + key, value);
}

export function secureGet(key: string): string {
  const prefixedKey = STORAGE_PREFIX + key;
  const raw = Taro.getStorageSync(prefixedKey);
  if (!raw || typeof raw !== 'string') {
    // fallback: 明文键（兼容 MCP 注入，仅 dev）
    const plain = Taro.getStorageSync(key);
    if (plain && typeof plain === 'string') {
      if (process.env.NODE_ENV === 'production') {
        console.warn('[secure-storage] plaintext fallback in production for key:', key);
      }
      return plain;
    }
    return '';
  }

  // AES 格式
  if (raw.startsWith(AES_MARKER)) {
    const decrypted = aesDecrypt(raw);
    if (decrypted !== null) return decrypted;
    // AES 解密失败
    if (process.env.NODE_ENV === 'production') {
      console.warn('[secure-storage] AES decrypt failed in production for key:', key);
      return '';
    }
    // dev: fallthrough to try other formats
  }

  // 非加密前缀 — dev mode plaintext 或 legacy XOR
  if (!raw.startsWith(AES_MARKER) && hasEncryptionKey()) {
    // key is configured but data isn't AES-encrypted — try legacy XOR
    try {
      const decoded = fromBase64(raw);
      if (decoded) return decoded;
    } catch {
      // fallthrough
    }
  }

  // dev mode: stored as plaintext by secureSet when no key
  return raw;
}

export function secureRemove(key: string): void {
  Taro.removeStorageSync(STORAGE_PREFIX + key);
}

const MIGRATION_KEYS = [
  'access_token', 'refresh_token', 'token_expires_at',
  'user_data', 'user_roles', 'tenant_id', 'wechat_openid',
];

export function migrateLegacyStorage(): void {
  try {
    for (const key of MIGRATION_KEYS) {
      const prefixed = STORAGE_PREFIX + key;
      const already = Taro.getStorageSync(prefixed);
      if (already) {
        // re-encrypt with AES if not already
        if (typeof already === 'string' && !already.startsWith(AES_MARKER)) {
          const value = secureGet(key);
          if (value) secureSet(key, value);
        }
        continue;
      }

      const legacy = Taro.getStorageSync(key);
      if (!legacy || typeof legacy !== 'string') continue;

      secureSet(key, legacy);
      Taro.removeStorageSync(key);
    }
  } catch {
    // migration best-effort
  }
}