Fix save and better write implementation

app/src/main/java/com/kunzisoft/keepass/activities/MainCredentialActivity.kt

@@ -55,13 +55,11 @@ import com.kunzisoft.keepass.autofill.AutofillComponent
 import com.kunzisoft.keepass.autofill.AutofillHelper
 import com.kunzisoft.keepass.biometric.AdvancedUnlockFragment
 import com.kunzisoft.keepass.biometric.AdvancedUnlockManager
-import com.kunzisoft.keepass.database.action.DatabaseTaskProvider
 import com.kunzisoft.keepass.database.element.Database
 import com.kunzisoft.keepass.database.element.MainCredential
 import com.kunzisoft.keepass.database.exception.DuplicateUuidDatabaseException
 import com.kunzisoft.keepass.database.exception.FileNotFoundDatabaseException
 import com.kunzisoft.keepass.education.PasswordActivityEducation
-import com.kunzisoft.keepass.hardware.HardwareKeyResponseHelper
 import com.kunzisoft.keepass.model.*
 import com.kunzisoft.keepass.services.DatabaseTaskNotificationService.Companion.ACTION_DATABASE_LOAD_TASK
 import com.kunzisoft.keepass.services.DatabaseTaskNotificationService.Companion.CIPHER_DATABASE_KEY
@@ -107,17 +105,11 @@ class MainCredentialActivity : DatabaseModeActivity(), AdvancedUnlockFragment.Bu
     private var mReadOnly: Boolean = false
     private var mForceReadOnly: Boolean = false
 
-    private var mHardwareKeyResponseHelper = HardwareKeyResponseHelper(this)
-
     private var mAutofillActivityResultLauncher: ActivityResultLauncher<Intent>? =
         if (Build.VERSION.SDK_INT >= Build.VERSION_CODES.O)
             AutofillHelper.buildActivityResultLauncher(this)
         else null
 
-    override fun initializeDatabaseTaskProvider(): DatabaseTaskProvider {
-        return DatabaseTaskProvider(this, mHardwareKeyResponseHelper)
-    }
-
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
 

app/src/main/java/com/kunzisoft/keepass/activities/legacy/DatabaseActivity.kt

@@ -8,6 +8,7 @@ import com.kunzisoft.keepass.database.action.DatabaseTaskProvider
 import com.kunzisoft.keepass.database.element.Database
 import com.kunzisoft.keepass.model.CipherEncryptDatabase
 import com.kunzisoft.keepass.database.element.MainCredential
+import com.kunzisoft.keepass.hardware.HardwareKeyResponseHelper
 import com.kunzisoft.keepass.tasks.ActionRunnable
 import com.kunzisoft.keepass.viewmodels.DatabaseViewModel
 
@@ -17,10 +18,12 @@ abstract class DatabaseActivity: StylishActivity(), DatabaseRetrieval {
     protected var mDatabaseTaskProvider: DatabaseTaskProvider? = null
     protected var mDatabase: Database? = null
 
+    private var mHardwareKeyResponseHelper = HardwareKeyResponseHelper(this)
+
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)
 
-        mDatabaseTaskProvider = initializeDatabaseTaskProvider()
+        mDatabaseTaskProvider = DatabaseTaskProvider(this, mHardwareKeyResponseHelper)
 
         mDatabaseTaskProvider?.onDatabaseRetrieved = { database ->
             val databaseWasReloaded = database?.wasReloaded == true
@@ -36,10 +39,6 @@ abstract class DatabaseActivity: StylishActivity(), DatabaseRetrieval {
         }
     }
 
-    open fun initializeDatabaseTaskProvider(): DatabaseTaskProvider {
-        return DatabaseTaskProvider(this)
-    }
-
     override fun onDatabaseRetrieved(database: Database?) {
         mDatabase = database
         mDatabaseViewModel.defineDatabase(database)

app/src/main/java/com/kunzisoft/keepass/database/element/Database.kt

@@ -792,28 +792,39 @@ class Database {
                 try {
                     outputStream = UriUtil.getUriOutputStream(contentResolver, saveUri)
                     outputStream?.let { definedOutputStream ->
-                        val databaseOutput =
-                            mDatabaseKDB?.let {
-                                DatabaseOutputKDB(it, definedOutputStream)
+                        mDatabaseKDB?.let { databaseKDB ->
+                            DatabaseOutputKDB(databaseKDB).apply {
+                                writeDatabase(definedOutputStream) {
+                                    if (mainCredential != null) {
+                                        databaseKDB.deriveMasterKey(
+                                            contentResolver,
+                                            mainCredential
+                                        )
+                                    } else {
+                                        // No master key change
+                                    }
+                                }
                             }
-                            ?: mDatabaseKDBX?.let {
-                                DatabaseOutputKDBX(it, definedOutputStream) {
+                        }
+                        ?: mDatabaseKDBX?.let { databaseKDBX ->
+                            DatabaseOutputKDBX(databaseKDBX).apply {
+                                writeDatabase(definedOutputStream) {
                                     if (mainCredential != null) {
                                         // Build new master key from MainCredential
-                                        mDatabaseKDBX?.deriveMasterKey(
+                                        databaseKDBX.deriveMasterKey(
                                             contentResolver,
                                             mainCredential,
                                             challengeResponseRetriever
                                         )
                                     } else {
                                         // Reuse composite key parts
-                                        mDatabaseKDBX?.deriveCompositeKey(
+                                        databaseKDBX.deriveCompositeKey(
                                             challengeResponseRetriever
                                         )
                                     }
                                 }
                             }
-                        databaseOutput?.output()
+                        }
                     }
                 } catch (e: Exception) {
                     throw IOException(e)

app/src/main/java/com/kunzisoft/keepass/database/element/database/DatabaseKDBX.kt

@@ -52,7 +52,6 @@ import com.kunzisoft.keepass.database.file.DatabaseHeaderKDBX.Companion.FILE_VER
 import com.kunzisoft.keepass.database.file.DatabaseHeaderKDBX.Companion.FILE_VERSION_40
 import com.kunzisoft.keepass.database.file.DatabaseHeaderKDBX.Companion.FILE_VERSION_41
 import com.kunzisoft.keepass.hardware.HardwareKey
-import com.kunzisoft.keepass.database.element.MainCredential
 import com.kunzisoft.keepass.utils.UnsignedInt
 import com.kunzisoft.keepass.utils.longTo8Bytes
 import java.io.IOException

app/src/main/java/com/kunzisoft/keepass/database/file/output/DatabaseOutput.kt

@@ -26,7 +26,7 @@ import java.io.OutputStream
 import java.security.NoSuchAlgorithmException
 import java.security.SecureRandom
 
-abstract class DatabaseOutput<Header : DatabaseHeader> protected constructor(protected var mOutputStream: OutputStream) {
+abstract class DatabaseOutput<Header : DatabaseHeader> {
 
     @Throws(DatabaseOutputException::class)
     protected open fun setIVs(header: Header): SecureRandom {
@@ -44,9 +44,7 @@ abstract class DatabaseOutput<Header : DatabaseHeader> protected constructor(pro
     }
 
     @Throws(DatabaseOutputException::class)
-    abstract fun output()
-
-    @Throws(DatabaseOutputException::class)
-    abstract fun outputHeader(outputStream: OutputStream): Header
+    abstract fun writeDatabase(outputStream: OutputStream,
+                               assignMasterKey: () -> Unit)
 
 }

app/src/main/java/com/kunzisoft/keepass/database/file/output/DatabaseOutputKDB.kt

@@ -39,9 +39,8 @@ import java.security.*
 import javax.crypto.Cipher
 import javax.crypto.CipherOutputStream
 
-class DatabaseOutputKDB(private val mDatabaseKDB: DatabaseKDB,
-                        outputStream: OutputStream)
-    : DatabaseOutput<DatabaseHeaderKDB>(outputStream) {
+class DatabaseOutputKDB(private val mDatabaseKDB: DatabaseKDB)
+    : DatabaseOutput<DatabaseHeaderKDB>() {
 
     private var headerHashBlock: ByteArray? = null
 
@@ -60,14 +59,15 @@ class DatabaseOutputKDB(private val mDatabaseKDB: DatabaseKDB,
     }
 
     @Throws(DatabaseOutputException::class)
-    override fun output() {
+    override fun writeDatabase(outputStream: OutputStream,
+                               assignMasterKey: () -> Unit) {
         // Before we output the header, we should sort our list of groups
         // and remove any orphaned nodes that are no longer part of the tree hierarchy
         // also remove the virtual root not present in kdb
         val rootGroup = mDatabaseKDB.rootGroup
         sortNodesForOutput()
 
-        val header = outputHeader(mOutputStream)
+        val header = outputHeader(outputStream, assignMasterKey)
         val finalKey = getFinalKey(header)
 
         val cipher: Cipher = try {
@@ -80,7 +80,7 @@ class DatabaseOutputKDB(private val mDatabaseKDB: DatabaseKDB,
         }
 
         try {
-            val cos = CipherOutputStream(mOutputStream, cipher)
+            val cos = CipherOutputStream(outputStream, cipher)
             val bos = BufferedOutputStream(cos)
             outputPlanGroupAndEntries(bos)
             bos.flush()
@@ -106,7 +106,8 @@ class DatabaseOutputKDB(private val mDatabaseKDB: DatabaseKDB,
     }
 
     @Throws(DatabaseOutputException::class)
-    override fun outputHeader(outputStream: OutputStream): DatabaseHeaderKDB {
+    private fun outputHeader(outputStream: OutputStream,
+                             assignMasterKey: () -> Unit): DatabaseHeaderKDB {
         // Build header
         val header = DatabaseHeaderKDB()
         header.signature1 = DatabaseHeaderKDB.DBSIG_1
@@ -131,6 +132,9 @@ class DatabaseOutputKDB(private val mDatabaseKDB: DatabaseKDB,
 
         setIVs(header)
 
+        mDatabaseKDB.transformSeed = header.transformSeed
+        assignMasterKey()
+
         // Header checksum
         val headerDigest: MessageDigest = HashManager.getHash256()
 

app/src/main/java/com/kunzisoft/keepass/database/file/output/DatabaseOutputKDBX.kt

@@ -56,10 +56,8 @@ import javax.crypto.CipherOutputStream
 import kotlin.experimental.or
 
 
-class DatabaseOutputKDBX(private val mDatabaseKDBX: DatabaseKDBX,
-                         outputStream: OutputStream,
-                         private val mAssignMasterKey: (() -> Unit))
-    : DatabaseOutput<DatabaseHeaderKDBX>(outputStream) {
+class DatabaseOutputKDBX(private val mDatabaseKDBX: DatabaseKDBX)
+    : DatabaseOutput<DatabaseHeaderKDBX>() {
 
     private var randomStream: StreamCipher? = null
     private lateinit var xml: XmlSerializer
@@ -68,21 +66,22 @@ class DatabaseOutputKDBX(private val mDatabaseKDBX: DatabaseKDBX,
     private var headerHmac: ByteArray? = null
 
     @Throws(DatabaseOutputException::class)
-    override fun output() {
+    override fun writeDatabase(outputStream: OutputStream,
+                               assignMasterKey: () -> Unit) {
 
         try {
-            header = outputHeader(mOutputStream)
+            header = outputHeader(outputStream, assignMasterKey)
 
             val osPlain: OutputStream = if (header!!.version.isBefore(FILE_VERSION_40)) {
-                val cos = attachStreamEncryptor(header!!, mOutputStream)
+                val cos = attachStreamEncryptor(header!!, outputStream)
                 cos.write(header!!.streamStartBytes)
 
                 HashedBlockOutputStream(cos)
             } else {
-                mOutputStream.write(hashOfHeader!!)
-                mOutputStream.write(headerHmac!!)
+                outputStream.write(hashOfHeader!!)
+                outputStream.write(headerHmac!!)
 
-                attachStreamEncryptor(header!!, HmacBlockOutputStream(mOutputStream, mDatabaseKDBX.hmacKey!!))
+                attachStreamEncryptor(header!!, HmacBlockOutputStream(outputStream, mDatabaseKDBX.hmacKey!!))
             }
 
             val xmlOutputStream: OutputStream
@@ -323,13 +322,14 @@ class DatabaseOutputKDBX(private val mDatabaseKDBX: DatabaseKDBX,
     }
 
     @Throws(DatabaseOutputException::class)
-    override fun outputHeader(outputStream: OutputStream): DatabaseHeaderKDBX {
+    private fun outputHeader(outputStream: OutputStream,
+                             assignMasterKey: () -> Unit): DatabaseHeaderKDBX {
         try {
             val header = DatabaseHeaderKDBX(mDatabaseKDBX)
             setIVs(header)
 
-            // TODO Check modification
-            mAssignMasterKey.invoke()
+            mDatabaseKDBX.transformSeed = header.transformSeed
+            assignMasterKey.invoke()
 
             val pho = DatabaseHeaderOutputKDBX(mDatabaseKDBX, header, outputStream)
             pho.output()

app/src/main/java/com/kunzisoft/keepass/services/DatabaseTaskNotificationService.kt

@@ -135,7 +135,8 @@ open class DatabaseTaskNotificationService : LockNotificationService(), Progress
 
         fun removeRequestChallengeListener(requestChallengeListener: RequestChallengeListener) {
             mainScope.launch {
-                // TODO mRequestChallengeListenerChannel.cancel()
+                mRequestChallengeListenerChannel.cancel()
+                mRequestChallengeListenerChannel = Channel(0)
             }
         }
     }

app/src/main/java/com/kunzisoft/keepass/settings/SettingsActivity.kt

@@ -35,9 +35,7 @@ import com.kunzisoft.keepass.R
 import com.kunzisoft.keepass.activities.dialogs.SetMainCredentialDialogFragment
 import com.kunzisoft.keepass.activities.helpers.ExternalFileHelper
 import com.kunzisoft.keepass.activities.legacy.DatabaseLockActivity
-import com.kunzisoft.keepass.database.action.DatabaseTaskProvider
 import com.kunzisoft.keepass.database.element.Database
-import com.kunzisoft.keepass.hardware.HardwareKeyResponseHelper
 import com.kunzisoft.keepass.database.element.MainCredential
 import com.kunzisoft.keepass.tasks.ActionRunnable
 import com.kunzisoft.keepass.timeout.TimeoutHelper
@@ -57,12 +55,6 @@ open class SettingsActivity
     private var toolbar: Toolbar? = null
     private var lockView: FloatingActionButton? = null
 
-    private var mHardwareKeyResponseHelper = HardwareKeyResponseHelper(this)
-
-    override fun initializeDatabaseTaskProvider(): DatabaseTaskProvider {
-        return DatabaseTaskProvider(this, mHardwareKeyResponseHelper)
-    }
-
     override fun onCreate(savedInstanceState: Bundle?) {
         super.onCreate(savedInstanceState)