Termix/src/backend/utils/user-crypto.ts

import crypto from "crypto";
import { getDb } from "../database/db/index.js";
import { settings } from "../database/db/schema.js";
import { eq } from "drizzle-orm";
import { databaseLogger } from "./logger.js";

interface KEKSalt {
  salt: string;
  iterations: number;
  algorithm: string;
  createdAt: string;
}

interface EncryptedDEK {
  data: string;
  iv: string;
  tag: string;
  algorithm: string;
  createdAt: string;
}

interface UserSession {
  dataKey: Buffer;
  lastActivity: number;
  expiresAt: number;
}

class UserCrypto {
  private static instance: UserCrypto;
  private userSessions: Map<string, UserSession> = new Map();
  private sessionExpiredCallback?: (userId: string) => void;

  private static readonly PBKDF2_ITERATIONS = 100000;
  private static readonly KEK_LENGTH = 32;
  private static readonly DEK_LENGTH = 32;
  private static readonly SESSION_DURATION = 24 * 60 * 60 * 1000;
  private static readonly MAX_INACTIVITY = 6 * 60 * 60 * 1000;

  private constructor() {
    setInterval(
      () => {
        this.cleanupExpiredSessions();
      },
      5 * 60 * 1000,
    );
  }

  static getInstance(): UserCrypto {
    if (!this.instance) {
      this.instance = new UserCrypto();
    }
    return this.instance;
  }

  setSessionExpiredCallback(callback: (userId: string) => void): void {
    this.sessionExpiredCallback = callback;
  }

  async setupUserEncryption(userId: string, password: string): Promise<void> {
    const kekSalt = await this.generateKEKSalt();
    await this.storeKEKSalt(userId, kekSalt);

    const KEK = this.deriveKEK(password, kekSalt);
    const DEK = crypto.randomBytes(UserCrypto.DEK_LENGTH);
    const encryptedDEK = this.encryptDEK(DEK, KEK);
    await this.storeEncryptedDEK(userId, encryptedDEK);

    KEK.fill(0);
    DEK.fill(0);
  }

  async setupOIDCUserEncryption(userId: string): Promise<void> {
    const existingEncryptedDEK = await this.getEncryptedDEK(userId);

    let DEK: Buffer;

    if (existingEncryptedDEK) {
      const systemKey = this.deriveOIDCSystemKey(userId);
      DEK = this.decryptDEK(existingEncryptedDEK, systemKey);
      systemKey.fill(0);
    } else {
      DEK = crypto.randomBytes(UserCrypto.DEK_LENGTH);
      const systemKey = this.deriveOIDCSystemKey(userId);

      try {
        const encryptedDEK = this.encryptDEK(DEK, systemKey);
        await this.storeEncryptedDEK(userId, encryptedDEK);

        const storedEncryptedDEK = await this.getEncryptedDEK(userId);
        if (
          storedEncryptedDEK &&
          storedEncryptedDEK.data !== encryptedDEK.data
        ) {
          DEK.fill(0);
          DEK = this.decryptDEK(storedEncryptedDEK, systemKey);
        } else if (!storedEncryptedDEK) {
          throw new Error("Failed to store and retrieve user encryption key.");
        }
      } finally {
        systemKey.fill(0);
      }
    }

    const now = Date.now();
    this.userSessions.set(userId, {
      dataKey: Buffer.from(DEK),
      lastActivity: now,
      expiresAt: now + UserCrypto.SESSION_DURATION,
    });

    DEK.fill(0);
  }

  async authenticateUser(userId: string, password: string): Promise<boolean> {
    try {
      const kekSalt = await this.getKEKSalt(userId);
      if (!kekSalt) return false;

      const KEK = this.deriveKEK(password, kekSalt);
      const encryptedDEK = await this.getEncryptedDEK(userId);
      if (!encryptedDEK) {
        KEK.fill(0);
        return false;
      }

      const DEK = this.decryptDEK(encryptedDEK, KEK);
      KEK.fill(0);

      if (!DEK || DEK.length === 0) {
        databaseLogger.error("DEK is empty or invalid after decryption", {
          operation: "user_crypto_auth_debug",
          userId,
          dekLength: DEK ? DEK.length : 0,
        });
        return false;
      }

      const now = Date.now();

      const oldSession = this.userSessions.get(userId);
      if (oldSession) {
        oldSession.dataKey.fill(0);
      }

      this.userSessions.set(userId, {
        dataKey: Buffer.from(DEK),
        lastActivity: now,
        expiresAt: now + UserCrypto.SESSION_DURATION,
      });

      DEK.fill(0);

      return true;
    } catch (error) {
      databaseLogger.warn("User authentication failed", {
        operation: "user_crypto_auth_failed",
        userId,
        error: error instanceof Error ? error.message : "Unknown",
      });
      return false;
    }
  }

  async authenticateOIDCUser(userId: string): Promise<boolean> {
    try {
      const encryptedDEK = await this.getEncryptedDEK(userId);

      if (!encryptedDEK) {
        await this.setupOIDCUserEncryption(userId);
        return true;
      }

      const systemKey = this.deriveOIDCSystemKey(userId);
      const DEK = this.decryptDEK(encryptedDEK, systemKey);
      systemKey.fill(0);

      if (!DEK || DEK.length === 0) {
        await this.setupOIDCUserEncryption(userId);
        return true;
      }

      const now = Date.now();

      const oldSession = this.userSessions.get(userId);
      if (oldSession) {
        oldSession.dataKey.fill(0);
      }

      this.userSessions.set(userId, {
        dataKey: Buffer.from(DEK),
        lastActivity: now,
        expiresAt: now + UserCrypto.SESSION_DURATION,
      });

      DEK.fill(0);

      return true;
    } catch (error) {
      await this.setupOIDCUserEncryption(userId);
      return true;
    }
  }

  getUserDataKey(userId: string): Buffer | null {
    const session = this.userSessions.get(userId);
    if (!session) {
      return null;
    }

    const now = Date.now();

    if (now > session.expiresAt) {
      this.userSessions.delete(userId);
      session.dataKey.fill(0);
      if (this.sessionExpiredCallback) {
        this.sessionExpiredCallback(userId);
      }
      return null;
    }

    if (now - session.lastActivity > UserCrypto.MAX_INACTIVITY) {
      this.userSessions.delete(userId);
      session.dataKey.fill(0);
      if (this.sessionExpiredCallback) {
        this.sessionExpiredCallback(userId);
      }
      return null;
    }

    session.lastActivity = now;
    return session.dataKey;
  }

  logoutUser(userId: string): void {
    const session = this.userSessions.get(userId);
    if (session) {
      session.dataKey.fill(0);
      this.userSessions.delete(userId);
    }
  }

  isUserUnlocked(userId: string): boolean {
    return this.getUserDataKey(userId) !== null;
  }

  async changeUserPassword(
    userId: string,
    oldPassword: string,
    newPassword: string,
  ): Promise<boolean> {
    try {
      // Validate current password
      const isValid = await this.validatePassword(userId, oldPassword);
      if (!isValid) return false;

      const kekSalt = await this.getKEKSalt(userId);
      if (!kekSalt) return false;

      const oldKEK = this.deriveKEK(oldPassword, kekSalt);
      const encryptedDEK = await this.getEncryptedDEK(userId);
      if (!encryptedDEK) return false;

      // Decrypt DEK with old password's KEK
      const DEK = this.decryptDEK(encryptedDEK, oldKEK);

      // The DEK stays the same - only the KEK changes with the new password
      // This preserves all encrypted user data
      
      // Generate new KEK from new password
      const newKekSalt = await this.generateKEKSalt();
      const newKEK = this.deriveKEK(newPassword, newKekSalt);
      
      // Re-encrypt the same DEK with new KEK
      const newEncryptedDEK = this.encryptDEK(DEK, newKEK);

      // Store new KEK salt and new encrypted DEK
      await this.storeKEKSalt(userId, newKekSalt);
      await this.storeEncryptedDEK(userId, newEncryptedDEK);

      // Trigger database save to persist the encryption key changes
      const { saveMemoryDatabaseToFile } = await import("../database/db/index.js");
      await saveMemoryDatabaseToFile();

      // Clean up sensitive data from memory
      oldKEK.fill(0);
      newKEK.fill(0);

      // Create a copy of DEK for the session before zeroing it out
      const dekCopy = Buffer.from(DEK);

      // Keep user session active with the same DEK
      const now = Date.now();
      const oldSession = this.userSessions.get(userId);
      if (oldSession) {
        oldSession.dataKey.fill(0);
      }
      
      this.userSessions.set(userId, {
        dataKey: dekCopy,
        lastActivity: now,
        expiresAt: now + UserCrypto.SESSION_DURATION,
      });

      // Zero out the original DEK from memory
      DEK.fill(0);

      return true;
    } catch (error) {
      databaseLogger.error("Password change failed", error, {
        operation: "password_change_error",
        userId,
        error: error instanceof Error ? error.message : "Unknown error",
      });
      return false;
    }
  }

  async resetUserPasswordWithPreservedDEK(
    userId: string,
    newPassword: string,
  ): Promise<boolean> {
    try {
      // This method preserves the existing DEK by re-encrypting it with the new password
      // Used for password reset when user is logged in

      // Check if user has an active session with DEK
      const existingDEK = this.getUserDataKey(userId);
      if (!existingDEK) {
        return false;
      }

      // Generate new KEK from new password
      const newKekSalt = await this.generateKEKSalt();
      const newKEK = this.deriveKEK(newPassword, newKekSalt);

      // Re-encrypt the existing DEK with new KEK
      const newEncryptedDEK = this.encryptDEK(existingDEK, newKEK);

      // Store new KEK salt and new encrypted DEK
      await this.storeKEKSalt(userId, newKekSalt);
      await this.storeEncryptedDEK(userId, newEncryptedDEK);

      // Trigger database save to persist the encryption key changes
      const { saveMemoryDatabaseToFile } = await import(
        "../database/db/index.js"
      );
      await saveMemoryDatabaseToFile();

      // Clean up sensitive data
      newKEK.fill(0);

      // Update session activity timestamp
      const session = this.userSessions.get(userId);
      if (session) {
        session.lastActivity = Date.now();
      }

      return true;
    } catch (error) {
      databaseLogger.error("Password reset with preserved DEK failed", error, {
        operation: "password_reset_preserve_error",
        userId,
        error: error instanceof Error ? error.message : "Unknown error",
      });
      return false;
    }
  }

  private async validatePassword(
    userId: string,
    password: string,
  ): Promise<boolean> {
    try {
      const kekSalt = await this.getKEKSalt(userId);
      if (!kekSalt) return false;

      const KEK = this.deriveKEK(password, kekSalt);
      const encryptedDEK = await this.getEncryptedDEK(userId);
      if (!encryptedDEK) return false;

      const DEK = this.decryptDEK(encryptedDEK, KEK);

      KEK.fill(0);
      DEK.fill(0);

      return true;
    } catch (error) {
      return false;
    }
  }

  private cleanupExpiredSessions(): void {
    const now = Date.now();
    const expiredUsers: string[] = [];

    for (const [userId, session] of this.userSessions.entries()) {
      if (
        now > session.expiresAt ||
        now - session.lastActivity > UserCrypto.MAX_INACTIVITY
      ) {
        session.dataKey.fill(0);
        expiredUsers.push(userId);
      }
    }

    expiredUsers.forEach((userId) => {
      this.userSessions.delete(userId);
    });
  }

  private async generateKEKSalt(): Promise<KEKSalt> {
    return {
      salt: crypto.randomBytes(32).toString("hex"),
      iterations: UserCrypto.PBKDF2_ITERATIONS,
      algorithm: "pbkdf2-sha256",
      createdAt: new Date().toISOString(),
    };
  }

  private deriveKEK(password: string, kekSalt: KEKSalt): Buffer {
    return crypto.pbkdf2Sync(
      password,
      Buffer.from(kekSalt.salt, "hex"),
      kekSalt.iterations,
      UserCrypto.KEK_LENGTH,
      "sha256",
    );
  }

  private deriveOIDCSystemKey(userId: string): Buffer {
    const systemSecret =
      process.env.OIDC_SYSTEM_SECRET || "termix-oidc-system-secret-default";
    const salt = Buffer.from(userId, "utf8");
    return crypto.pbkdf2Sync(
      systemSecret,
      salt,
      100000,
      UserCrypto.KEK_LENGTH,
      "sha256",
    );
  }

  private encryptDEK(dek: Buffer, kek: Buffer): EncryptedDEK {
    const iv = crypto.randomBytes(16);
    const cipher = crypto.createCipheriv("aes-256-gcm", kek, iv);

    let encrypted = cipher.update(dek);
    encrypted = Buffer.concat([encrypted, cipher.final()]);
    const tag = cipher.getAuthTag();

    return {
      data: encrypted.toString("hex"),
      iv: iv.toString("hex"),
      tag: tag.toString("hex"),
      algorithm: "aes-256-gcm",
      createdAt: new Date().toISOString(),
    };
  }

  private decryptDEK(encryptedDEK: EncryptedDEK, kek: Buffer): Buffer {
    const decipher = crypto.createDecipheriv(
      "aes-256-gcm",
      kek,
      Buffer.from(encryptedDEK.iv, "hex"),
    );

    decipher.setAuthTag(Buffer.from(encryptedDEK.tag, "hex"));
    let decrypted = decipher.update(Buffer.from(encryptedDEK.data, "hex"));
    decrypted = Buffer.concat([decrypted, decipher.final()]);

    return decrypted;
  }

  private async storeKEKSalt(userId: string, kekSalt: KEKSalt): Promise<void> {
    const key = `user_kek_salt_${userId}`;
    const value = JSON.stringify(kekSalt);

    const existing = await getDb()
      .select()
      .from(settings)
      .where(eq(settings.key, key));

    if (existing.length > 0) {
      await getDb()
        .update(settings)
        .set({ value })
        .where(eq(settings.key, key));
    } else {
      await getDb().insert(settings).values({ key, value });
    }
  }

  private async getKEKSalt(userId: string): Promise<KEKSalt | null> {
    try {
      const key = `user_kek_salt_${userId}`;
      const result = await getDb()
        .select()
        .from(settings)
        .where(eq(settings.key, key));

      if (result.length === 0) {
        return null;
      }

      return JSON.parse(result[0].value);
    } catch (error) {
      return null;
    }
  }

  private async storeEncryptedDEK(
    userId: string,
    encryptedDEK: EncryptedDEK,
  ): Promise<void> {
    const key = `user_encrypted_dek_${userId}`;
    const value = JSON.stringify(encryptedDEK);

    const existing = await getDb()
      .select()
      .from(settings)
      .where(eq(settings.key, key));

    if (existing.length > 0) {
      await getDb()
        .update(settings)
        .set({ value })
        .where(eq(settings.key, key));
    } else {
      await getDb().insert(settings).values({ key, value });
    }
  }

  private async getEncryptedDEK(userId: string): Promise<EncryptedDEK | null> {
    try {
      const key = `user_encrypted_dek_${userId}`;
      const result = await getDb()
        .select()
        .from(settings)
        .where(eq(settings.key, key));

      if (result.length === 0) {
        return null;
      }

      return JSON.parse(result[0].value);
    } catch (error) {
      return null;
    }
  }
}

export { UserCrypto, type KEKSalt, type EncryptedDEK };