Pārlūkot izejas kodu

Add write method

maxep 5 gadi atpakaļ
vecāks
revīzija
29eceebf26

+ 17 - 34
Sources/KDB/Database.swift

@@ -34,16 +34,7 @@ public class Database {
 
         let data: Bytes = try input.read()
         let key = try header.masterKey(from: compositeKey)
-
-        let cipher: Cipher
-
-        if header.cipher.contains(.aes) {
-            cipher = try AESCipher(key: key, iv: header.initialVector)
-        } else if header.cipher.contains(.twofish) {
-            cipher = try Twofish(key: key, iv: header.initialVector)
-        } else {
-            throw KDBError.unsupportedCipher
-        }
+        let cipher = try header.cipher(key: key)
 
         let content = try cipher.decrypt(data: data)
 
@@ -69,36 +60,32 @@ public class Database {
 
     public func write(to output: Output, compositeKey: CompositeKey) throws {
 
-        func groups(in group: Group) -> UInt32 {
-            let count = UInt32(group.childs.count)
-            return group.childs.reduce(count) { $0 + groups(in: $1) }
+        func groups(in group: Group) -> [Group] {
+            return group.childs.reduce(group.childs) { $0 + groups(in: $1) }
         }
 
-        func entries(in group: Group) -> UInt32 {
-            let count = UInt32(group.entries.count)
-            return group.childs.reduce(count) { $0 + entries(in: $1) }
+        func entries(in group: Group) -> [Entry] {
+            return group.childs.reduce(group.entries) { $0 + entries(in: $1) }
         }
 
         let key = try header.masterKey(from: compositeKey)
-        let cipher: Cipher
-
-        if header.cipher.contains(.aes) {
-            cipher = try AESCipher(key: key, iv: header.initialVector)
-        } else if header.cipher.contains(.twofish) {
-            cipher = try Twofish(key: key, iv: header.initialVector)
-        } else {
-            throw KDBError.unsupportedCipher
-        }
+        let cipher = try header.cipher(key: key)
 
-        let content = try Data(from: root)
+        let groups = groups(in: root)
+        let entries = entries(in: root)
+
+        let content = Output()
+        try content.write(groups)
+        try content.write(entries)
+        guard var data = content.bytes else { throw KDBError.noData }
 
         var header = header
-        header.contentHash = SHA256.hash(content)
-        header.groups = groups(in: root)
-        header.entries = entries(in: root)
+        header.contentHash = SHA256.hash(data)
+        header.groups = UInt32(groups.count)
+        header.entries = UInt32(entries.count)
         try output.write(header)
 
-        let data = try cipher.encrypt(data: content)
+        data = try cipher.encrypt(data: data)
         try output.write(data)
     }
 
@@ -142,7 +129,3 @@ private func Root(from input: Input, groups: Int, entries: Int) throws -> Group
 
     return root
 }
-
-private func Data(from root: Group) throws -> Bytes {
-    fatalError()
-}

+ 1 - 0
Sources/KDB/Error.swift

@@ -25,4 +25,5 @@ public enum KDBError: Error {
     case emptyCompositeKey
     case invalidPassword
     case invalidKey
+    case noData
 }

+ 12 - 0
Sources/KDB/Header.swift

@@ -70,6 +70,18 @@ extension Header: Streamable {
 
 extension Header {
 
+    func cipher(key: Bytes) throws -> Cipher {
+        if cipher.contains(.aes) {
+            return try AESCipher(key: key, iv: initialVector)
+        }
+
+        if cipher.contains(.twofish) {
+            return try Twofish(key: key, iv: initialVector)
+        }
+
+        throw KDBError.unsupportedCipher
+    }
+
     func masterKey(from compositeKey: CompositeKey) throws -> Bytes {
         let key = try compositeKey.serialize()
 

+ 1 - 1
Sources/KDBX/Database4.swift

@@ -56,10 +56,10 @@ class Database4: Database {
 
         document = try XML.Document(xml: stream.remaining.data, options: options)
     }
+
     func write(to output: Output, compositeKey: CompositeKey) throws {
         fatalError()
     }
-
 }
 
 private func Unhash(header: Bytes, data: Bytes, key: Bytes) throws -> Bytes {

+ 2 - 4
Sources/KDBX/File.swift

@@ -20,9 +20,7 @@ import Foundation
 import Binary
 
 public let FileSignature: UInt32 = 0x9AA2D903
-
 public let BetaFileFormat: UInt32 = 0xB54BFB66
-
 public let FileFormat: UInt32 = 0xB54BFB67
 
 public struct Version {
@@ -75,10 +73,10 @@ public class File {
 
     public func write(to output: Output, compositeKey: CompositeKey) throws {
         try output.write(FileSignature)
-        try output.write(BetaFileFormat)
+        try output.write(FileFormat)
         try output.write(version)
+        try database.write(to: output, compositeKey: compositeKey)
     }
-
 }
 
 extension Version: Streamable {

+ 3 - 0
Sources/KeePass/Database.swift

@@ -17,9 +17,12 @@
 // along with KeePass. If not, see <https://www.gnu.org/licenses/>.
 
 import Foundation
+import Binary
 
 public protocol Database {
     associatedtype Root: Group
 
     var root: Root { get }
+
+    func write(to output: Output, compositeKey: CompositeKey) throws
 }

+ 6 - 1
Sources/KeePass/KDB.swift

@@ -22,7 +22,12 @@ import KDB
 
 extension CompositeKey: KDB.CompositeKey { }
 
-extension KDB.Database: Database {}
+extension KDB.Database: Database {
+
+    public func write(to output: Output, compositeKey: CompositeKey) throws {
+        try self.write(to: output, compositeKey: compositeKey as KDB.CompositeKey )
+    }
+}
 
 extension TLV where Type == KDB.Entry.Column {
 

+ 5 - 0
Sources/KeePass/KDBX.swift

@@ -26,7 +26,12 @@ let DateFormatter = ISO8601DateFormatter()
 extension CompositeKey: KDBX.CompositeKey { }
 
 extension KDBX.File: Database {
+
     public var root: Element { database.document.root.KeePassFile.Root }
+
+    public func write(to output: Output, compositeKey: CompositeKey) throws {
+        try self.write(to: output, compositeKey: compositeKey as KDBX.CompositeKey )
+    }
 }
 
 extension Element {

+ 5 - 4
Sources/KeePass/KeePass.swift

@@ -40,18 +40,19 @@ public class KeePass {
             throw KeePassError.invalidFileFormat
         }
 
-        let format: FileFormat = try stream.read()
+        let format: UInt32 = try stream.read()
 
         switch format {
-        case .kdb:
+        case KDB.FileFormat:
             return AnyDatabase( try KDB.Database(from: stream, compositeKey: compositeKey) )
-        case .prekdbx, .kdbx:
+        case KDBX.BetaFileFormat, KDBX.FileFormat:
             return AnyDatabase ( try KDBX.File(from: stream, compositeKey: compositeKey) )
+        default:
+            throw KeePassError.invalidFileFormat
         }
     }
 
     public static func open(contentOf xml: URL) throws -> some Database {
         try KDBX.File(xml: xml)
     }
-
 }

+ 8 - 0
Sources/KeePass/TypeErasure.swift

@@ -17,6 +17,7 @@
 // along with KeePass. If not, see <https://www.gnu.org/licenses/>.
 
 import Foundation
+import Binary
 
 @inline(never)
 func _abstract(file: StaticString = #file, line: UInt = #line) -> Never {
@@ -27,12 +28,16 @@ func _abstract(file: StaticString = #file, line: UInt = #line) -> Never {
 
 class _AnyDatabaseBoxBase: Database {
     var root: AnyGroup { _abstract() }
+    func write(to output: Output, compositeKey: CompositeKey) throws { _abstract() }
 }
 
 final class _AnyDatabaseBox<Base>: _AnyDatabaseBoxBase where Base: Database {
     override var root: AnyGroup { AnyGroup( _base.root ) }
     var _base: Base
     init(_ base: Base) { _base = base }
+    override func write(to output: Output, compositeKey: CompositeKey) throws {
+        try _base.write(to: output, compositeKey: compositeKey)
+    }
 }
 
 class AnyDatabase: Database {
@@ -41,6 +46,9 @@ class AnyDatabase: Database {
     init<T>(_ base: T) where T: Database {
         _box = _AnyDatabaseBox(base)
     }
+    func write(to output: Output, compositeKey: CompositeKey) throws {
+        try _box.write(to: output, compositeKey: compositeKey)
+    }
 }
 
 // MARK: - Group

+ 5 - 0
Tests/KeePassTests/KeePassTests.swift

@@ -1,4 +1,5 @@
 import XCTest
+import Binary
 @testable import KeePass
 
 final class KeePassTests: XCTestCase {
@@ -8,6 +9,10 @@ final class KeePassTests: XCTestCase {
         let key = CompositeKey(password: "1234")
         let db = try KeePass.open(contentOf: file, compositeKey: key)
         print(db)
+
+        let stream = Output()
+        try db.write(to: stream, compositeKey: key)
+        print(stream.lenght)
     }
 
     func testKDBX3() throws {