openide/platform/credential-store-impl/test/keePass/BaseKeePassFileManagerTest.kt

// Copyright 2000-2024 JetBrains s.r.o. and contributors. Use of this source code is governed by the Apache 2.0 license.
package com.intellij.credentialStore.keePass

import com.intellij.credentialStore.CredentialAttributes
import com.intellij.credentialStore.Credentials
import com.intellij.credentialStore.EncryptionSpec
import com.intellij.credentialStore.createSecureRandom
import com.intellij.openapi.util.ThreadLocalCachedValue
import com.intellij.testFramework.assertions.Assertions.assertThat
import com.intellij.testFramework.rules.InMemoryFsRule
import org.junit.Rule
import java.awt.Component
import java.nio.file.Path
import java.security.SecureRandom

internal val SECURE_RANDOM_CACHE = object : ThreadLocalCachedValue<SecureRandom>() {
  public override fun create() = createSecureRandom()
}

internal open class BaseKeePassFileManagerTest {
  companion object {
    val testCredentialAttributes = CredentialAttributes("foo", "u")
  }

  @JvmField
  @Rule
  val fsRule = InMemoryFsRule()

  protected fun createTestStoreWithCustomMasterKey(baseDir: Path = fsRule.fs.getPath("/")): KeePassCredentialStore {
    val store = createStore(baseDir)
    store.set(testCredentialAttributes, Credentials("u", "p"))
    store.setMasterKey("foo", SECURE_RANDOM_CACHE.value)
    return store
  }

  protected fun createStore() = createStore(fsRule.fs.getPath("/"))

  protected fun checkStoreAfterSuccessfulImport(store: KeePassCredentialStore) {
    store.reload()

    assertThat(store.dbFile).exists()
    assertThat(store.mainKeyFile).exists()

    assertThat(store.get(testCredentialAttributes)!!.getPasswordAsString()).isEqualTo("p")
  }
}

internal class TestKeePassFileManager(
  file: Path,
  masterKeyFile: Path,
  private val masterPasswordRequestAnswer: String? = null,
  private val oldMasterPasswordRequestAnswer: String? = null,
  masterKeyEncryptionSpec: EncryptionSpec = defaultEncryptionSpec
) : KeePassFileManager(file, masterKeyFile, masterKeyEncryptionSpec, lazyOf(SECURE_RANDOM_CACHE.value)) {
  constructor(store: KeePassCredentialStore,
              masterPasswordRequestAnswer: String? = null,
              oldMasterPasswordRequestAnswer: String? = null,
              masterKeyEncryptionSpec: EncryptionSpec = defaultEncryptionSpec) : this(store.dbFile, store.mainKeyFile, masterPasswordRequestAnswer,
                                                                                      oldMasterPasswordRequestAnswer, masterKeyEncryptionSpec)

  var isUnsatisfiedMasterPasswordRequest = false

  override fun requestMainPassword(title: String, topNote: String?, contextComponent: Component?, ok: (value: ByteArray) -> String?): Boolean {
    if (masterPasswordRequestAnswer == null) {
      isUnsatisfiedMasterPasswordRequest = true
      return false
    }
    else {
      assertThat(ok(masterPasswordRequestAnswer.toByteArray())).isNull()
      return true
    }
  }

  override fun requestCurrentAndNewKeys(contextComponent: Component?): Boolean {
    doSetNewMainPassword(oldMasterPasswordRequestAnswer!!.toCharArray(), masterPasswordRequestAnswer!!.toCharArray())
    return true
  }
}

internal fun KeePassCredentialStore.setMasterKey(value: String, secureRandom: SecureRandom) {
  setMainPassword(MainKey(value.toByteArray(), isAutoGenerated = false, encryptionSpec = defaultEncryptionSpec), secureRandom)
}