From 18a417adddd34499cf2c3aa80d4de1c1cf71f58d Mon Sep 17 00:00:00 2001
From: Jin <22962980+JYC333@users.noreply.github.com>
Date: Fri, 28 Mar 2025 01:53:53 +0100
Subject: [PATCH] =?UTF-8?q?feat:=20=F0=9F=8E=B8=20add=20totp=20encryption?=
MIME-Version: 1.0
Content-Type: text/plain; charset=UTF-8
Content-Transfer-Encoding: 8bit

---
 src/services/encryption/totp_encryption.ts | 86 ++++++++++++++++++++++
 src/services/options_interface.ts          |  3 +
 2 files changed, 89 insertions(+)
 create mode 100644 src/services/encryption/totp_encryption.ts

diff --git a/src/services/encryption/totp_encryption.ts b/src/services/encryption/totp_encryption.ts
new file mode 100644
index 000000000..d3db7ad67
--- /dev/null
+++ b/src/services/encryption/totp_encryption.ts
@@ -0,0 +1,86 @@
+import optionService from "../options.js";
+import myScryptService from "./my_scrypt.js";
+import { randomSecureToken, toBase64 } from "../utils.js";
+import dataEncryptionService from "./data_encryption.js";
+import type { OptionNames } from "../options_interface.js";
+
+const TOTP_OPTIONS: Record<string, OptionNames> = {
+    SALT: "totpEncryptionSalt",
+    ENCRYPTED_SECRET: "totpEncryptedSecret",
+    VERIFICATION_HASH: "totpVerificationHash"
+} as const;
+
+function verifyTotpSecret(secret: string): boolean {
+    const givenSecretHash = toBase64(myScryptService.getVerificationHash(secret));
+    const dbSecretHash = optionService.getOptionOrNull(TOTP_OPTIONS.VERIFICATION_HASH);
+
+    if (!dbSecretHash) {
+        return false;
+    }
+
+    return givenSecretHash === dbSecretHash;
+}
+
+function setTotpSecret(secret: string): void {
+    if (!secret) {
+        throw new Error("TOTP secret cannot be empty");
+    }
+
+    // 生成新的加密盐值
+    const encryptionSalt = randomSecureToken(32);
+    optionService.setOption(TOTP_OPTIONS.SALT, encryptionSalt);
+
+    // 使用 scrypt 生成验证哈希
+    const verificationHash = toBase64(myScryptService.getVerificationHash(secret));
+    optionService.setOption(TOTP_OPTIONS.VERIFICATION_HASH, verificationHash);
+
+    // 使用数据加密密钥加密 TOTP secret
+    const encryptedSecret = dataEncryptionService.encrypt(
+        Buffer.from(encryptionSalt),
+        secret
+    );
+    optionService.setOption(TOTP_OPTIONS.ENCRYPTED_SECRET, encryptedSecret);
+}
+
+function getTotpSecret(): string | null {
+    const encryptionSalt = optionService.getOptionOrNull(TOTP_OPTIONS.SALT);
+    const encryptedSecret = optionService.getOptionOrNull(TOTP_OPTIONS.ENCRYPTED_SECRET);
+
+    if (!encryptionSalt || !encryptedSecret) {
+        return null;
+    }
+
+    try {
+        const decryptedSecret = dataEncryptionService.decrypt(
+            Buffer.from(encryptionSalt),
+            encryptedSecret
+        );
+
+        if (!decryptedSecret) {
+            return null;
+        }
+
+        return decryptedSecret.toString();
+    } catch (e) {
+        console.error("Failed to decrypt TOTP secret:", e);
+        return null;
+    }
+}
+
+function resetTotpSecret(): void {
+    optionService.setOption(TOTP_OPTIONS.SALT, "");
+    optionService.setOption(TOTP_OPTIONS.ENCRYPTED_SECRET, "");
+    optionService.setOption(TOTP_OPTIONS.VERIFICATION_HASH, "");
+}
+
+function isTotpSecretSet(): boolean {
+    return !!optionService.getOptionOrNull(TOTP_OPTIONS.VERIFICATION_HASH);
+}
+
+export default {
+    verifyTotpSecret,
+    setTotpSecret,
+    getTotpSecret,
+    resetTotpSecret,
+    isTotpSecretSet
+};
diff --git a/src/services/options_interface.ts b/src/services/options_interface.ts
index 258ee9195..cdd9184fb 100644
--- a/src/services/options_interface.ts
+++ b/src/services/options_interface.ts
@@ -51,6 +51,9 @@ export interface OptionDefinitions extends KeyboardShortcutsOptions<KeyboardActi
     // Multi-Factor Authentication
     mfaEnabled: boolean;
     mfaMethod: string;
+    totpEncryptionSalt: string;
+    totpEncryptedSecret: string;
+    totpVerificationHash: string;
     encryptedRecoveryCodes: boolean;
     userSubjectIdentifierSaved: boolean;
     recoveryCodeInitialVector: string;