Added helper to traverse database

keemobile · Sep 29, 2021 · 9f8dcd0 · 9f8dcd0
1 parent 38dae06
commit 9f8dcd0
Show file tree

Hide file tree

Showing 3 changed files with 44 additions and 14 deletions.
diff --git a/kotpass/src/main/kotlin/io/github/anvell/kotpass/database/KeePassDatabase.kt b/kotpass/src/main/kotlin/io/github/anvell/kotpass/database/KeePassDatabase.kt
@@ -2,10 +2,7 @@ package io.github.anvell.kotpass.database
 
 import io.github.anvell.kotpass.database.header.DatabaseHeader
 import io.github.anvell.kotpass.database.header.DatabaseInnerHeader
-import io.github.anvell.kotpass.models.DatabaseContent
-import io.github.anvell.kotpass.models.Entry
-import io.github.anvell.kotpass.models.Group
-import io.github.anvell.kotpass.models.Meta
+import io.github.anvell.kotpass.models.*
 import java.util.*
 
 sealed class KeePassDatabase {
@@ -77,6 +74,10 @@ sealed class KeePassDatabase {
     }
 }
 
+fun KeePassDatabase.traverse(
+    block: (DatabaseElement) -> Unit
+) = content.group.traverse(block)
+
 fun KeePassDatabase.findGroup(
     predicate: (Group) -> Boolean
 ): Pair<Group?, Group>? {

diff --git a/kotpass/src/main/kotlin/io/github/anvell/kotpass/models/Group.kt b/kotpass/src/main/kotlin/io/github/anvell/kotpass/models/Group.kt
@@ -22,6 +22,25 @@ data class Group(
     val customData: Map<String, CustomDataValue> = mapOf()
 ) : DatabaseElement {
 
+    fun traverse(
+        block: (DatabaseElement) -> Unit
+    ) {
+        val stack = Stack<Group>()
+        stack.push(this)
+
+        while (!stack.empty()) {
+            val current = stack.pop()
+            block(current)
+
+            for (entry in current.entries) {
+                block(entry)
+            }
+            for (group in current.groups) {
+                stack.push(group)
+            }
+        }
+    }
+
     fun findChildGroup(
         predicate: (Group) -> Boolean
     ): Pair<Group, Group>? {

diff --git a/kotpass/src/test/kotlin/io/github/anvell/kotpass/database/KeePassDatabaseSpec.kt b/kotpass/src/test/kotlin/io/github/anvell/kotpass/database/KeePassDatabaseSpec.kt
@@ -2,16 +2,9 @@ package io.github.anvell.kotpass.database
 
 import io.github.anvell.kotpass.constants.BasicFields
 import io.github.anvell.kotpass.cryptography.EncryptedValue
-import io.github.anvell.kotpass.database.modifiers.cleanupHistory
-import io.github.anvell.kotpass.database.modifiers.modifyEntry
-import io.github.anvell.kotpass.database.modifiers.modifyGroup
-import io.github.anvell.kotpass.database.modifiers.moveEntry
-import io.github.anvell.kotpass.database.modifiers.moveGroup
-import io.github.anvell.kotpass.database.modifiers.removeEntry
-import io.github.anvell.kotpass.database.modifiers.removeGroup
-import io.github.anvell.kotpass.database.modifiers.withHistory
-import io.github.anvell.kotpass.database.modifiers.withRecycleBin
+import io.github.anvell.kotpass.database.modifiers.*
 import io.github.anvell.kotpass.io.decodeBase64ToArray
+import io.github.anvell.kotpass.models.DatabaseElement
 import io.github.anvell.kotpass.models.DeletedObject
 import io.github.anvell.kotpass.models.TimeData
 import io.github.anvell.kotpass.resources.DatabaseRes
@@ -24,7 +17,6 @@ import java.io.ByteArrayInputStream
 import java.io.ByteArrayOutputStream
 import java.time.Instant
 import java.time.Period
-import java.util.*
 
 class KeePassDatabaseSpec : DescribeSpec({
 
@@ -79,6 +71,24 @@ class KeePassDatabaseSpec : DescribeSpec({
     }
 
     describe("Database search") {
+        it("Traverse database") {
+            val database = loadDatabase(
+                rawData = DatabaseRes.GroupsAndEntries.DbGroupsAndEntries,
+                passphrase = "1"
+            )
+            val result = mutableSetOf<DatabaseElement>()
+            database.traverse { result += it }
+
+            result.map { it.uuid } shouldContainAll setOf(
+                DatabaseRes.GroupsAndEntries.Group1,
+                DatabaseRes.GroupsAndEntries.Group2,
+                DatabaseRes.GroupsAndEntries.Group3,
+                DatabaseRes.GroupsAndEntries.Entry1,
+                DatabaseRes.GroupsAndEntries.Entry2,
+                DatabaseRes.GroupsAndEntries.Entry3
+            )
+        }
+
         it("Finds entries with specific title") {
             val database = loadDatabase(
                 rawData = DatabaseRes.GroupsAndEntries.DbGroupsAndEntries,