contexted-v3/src/composables/useEncryption.ts

import { doc, getDoc, setDoc } from 'firebase/firestore'
import { user, db } from '@/composables/useFirebase'
import { decrypt, encrypt, calculateClientKey, generateEncryptionKey } from '@/utils/crypto'
import { preferredNotesSource } from '@/composables/useSettings'
import { activeNotesSource, syncNotesToFirebase, baseNotes } from '@/composables/useNotes'

function getClientKeysFromLocalStorage(): { [uid: string]: string } {
    try {
        return JSON.parse(localStorage.getItem('clientKeys') || '{}')
    } catch (e) {
        return {}
    }
}

export const clientKey = ref<ClientKey>()

export const getClientKey = () => {
    if (!user.value) return
    const clientKeys = getClientKeysFromLocalStorage()
    clientKey.value = clientKeys[user.value?.uid]
}

export const setClientKey = (passphrase: string) => {
    const calculatedClientKey = calculateClientKey(passphrase)
    const verified = verifyClientKey(calculatedClientKey)
    if (!user.value || !verified) return false
    const clientKeys = getClientKeysFromLocalStorage()
    clientKeys[user.value.uid] = calculatedClientKey
    localStorage.setItem('clientKeys', JSON.stringify(clientKeys))
    clientKey.value = calculatedClientKey
    getEncryptionKey()
    return true
}

export const verifyClientKey = (clientKey: ClientKey) => {
    try {
        if (!encryptedEncryptionKey.value) throw new Error('Encryption key is null')
        if (!clientKey) throw new Error('Client key is null')
        decrypt(encryptedEncryptionKey.value, clientKey)
        return true
    } catch (e) {
        console.log(e)
        return false
    }
}

const removeClientKey = () => {
    if (!user.value) return
    const clientKeys = structuredClone(getClientKeysFromLocalStorage())
    delete clientKeys[user.value?.uid]
    localStorage.setItem('clientKeys', JSON.stringify(clientKeys))
}

const encryptedEncryptionKey = ref<EncryptedEncryptionKey | null>()

async function getEncryptedEncryptionKey(): Promise<EncryptedEncryptionKey | void> {
    if (!user.value || !db.value) return
    const data = (await getDoc(doc(db.value, 'encryptionKeys', user.value.uid))).data()
    return data?.key
}

async function setEncryptedEncryptionKey(encryptedEncryptionKey: EncryptedEncryptionKey | null) {
    if (!user.value || !db.value) return
    const docRef = doc(db.value, 'encryptionKeys', user.value.uid)
    await setDoc(docRef, { key: encryptedEncryptionKey })
}

export const encryptionKey = ref<EncryptionKey | null>()

export async function getEncryptionKey() {
    encryptedEncryptionKey.value = (await getEncryptedEncryptionKey()) || undefined
    if (encryptedEncryptionKey.value && clientKey.value) {
        encryptionKey.value = decrypt(encryptedEncryptionKey.value, clientKey.value)
    } else if (!encryptedEncryptionKey.value) {
        encryptionKey.value = null
    } else {
        encryptionKey.value = undefined
    }
}

export const passphraseRequired = computed(() => {
    return Boolean(encryptedEncryptionKey.value && !clientKey.value)
})

const decryptNote = (note: BaseNote, key: EncryptionKey) => {
    try {
        return {
            ...note,
            title: decrypt(note.title, key),
            content: decrypt(note.content, key)
        }
    } catch (error: any) {
        console.error(error)
        return note
    }
}

export const decryptNotes = (notes: BaseNotes, encryptionKey: EncryptionKey) => {
    const decryptedNotes = Object.fromEntries(
        Object.entries(notes).map(([noteId, note]) => [
            noteId,
            { ...decryptNote(note, encryptionKey) }
        ])
    )
    return decryptedNotes
}

const encryptNote = (note: BaseNote, key: EncryptionKey) => {
    return {
        ...note,
        title: encrypt(note.title, key),
        content: encrypt(note.content, key)
    }
}

export const encryptNotes = (notes: BaseNotes, encryptionKey: EncryptionKey) => {
    const encryptedNotes = Object.fromEntries(
        Object.entries(notes).map(([noteId, note]) => [
            noteId,
            { ...encryptNote(note, encryptionKey) }
        ])
    )
    return encryptedNotes
}

export const verifyPassphrase = (passphrase: string) => {
    const calculatedClientKey = calculateClientKey(passphrase)
    return calculatedClientKey === clientKey.value
}

export const disableEncryption = async (passphrase: string) => {
    if (!encryptionKey.value) return "Encryption key doesn't exist."
    if (!verifyPassphrase(passphrase)) return 'Passphrase is incorrect.'
    preferredNotesSource.value = 'firebase'
    if (activeNotesSource.value !== 'firebase') throw Error('Something went wrong.')
    await setEncryptedEncryptionKey(null)
    encryptedEncryptionKey.value = null
    encryptionKey.value = undefined
    removeClientKey()
    await syncNotesToFirebase(baseNotes.value)
    getEncryptionKey()
}

export const enableEncryption = async (passphrase: string) => {
    preferredNotesSource.value = 'firebase'
    if (activeNotesSource.value !== 'firebase') throw Error('Something went wrong.')
    const candidateEncryptionKey = generateEncryptionKey()
    const candidateClientKey = calculateClientKey(passphrase)
    const candidateEncryptedEncryptionKey = encrypt(candidateEncryptionKey, candidateClientKey)
    await setEncryptedEncryptionKey(candidateEncryptedEncryptionKey)
    encryptedEncryptionKey.value = candidateEncryptedEncryptionKey
    encryptionKey.value = candidateEncryptionKey
    setClientKey(passphrase)
    await syncNotesToFirebase(baseNotes.value)
}

export const clearEncryptionKeys = () => {
    encryptionKey.value = undefined
    encryptedEncryptionKey.value = undefined
}