feat: create eth2 keystores

koraykoska · koraykoska · commit 4038fccf6000 · 2024-04-10T14:31:23.000+02:00
diff --git a/Package.swift b/Package.swift
@@ -22,6 +22,7 @@ let package = Package(
         .package(url: "https://github.com/krzyzanowskim/CryptoSwift.git", from: "1.8.1"),
         .package(url: "https://github.com/Boilertalk/secp256k1.swift.git", from: "0.1.7"),
         .package(name: "Scrypt", url: "https://github.com/greymass/swift-scrypt.git", from: "1.0.0"),
+        .package(url: "https://github.com/MyEtherWallet/bls-eth-swift.git", .upToNextMajor(from: "1.0.0")),
 
         // Test dependencies
         .package(url: "https://github.com/Quick/Quick.git", from: "5.0.1"),
@@ -34,6 +35,7 @@ let package = Package(
                 .product(name: "CryptoSwift", package: "CryptoSwift"),
                 .product(name: "secp256k1", package: "secp256k1.swift"),
                 .product(name: "Scrypt", package: "Scrypt"),
+                .product(name: "bls-eth-swift", package: "bls-eth-swift"),
             ],
             path: "Sources",
             sources: ["Keystore"]),
diff --git a/Sources/Keystore/KeystoreETH2.swift b/Sources/Keystore/KeystoreETH2.swift
@@ -6,15 +6,28 @@ public struct KeystoreETH2: Codable {
 
     // MARK: - Public API
 
-    /// Creates a new keystore for the given private key and password.
+    /// Creates a keystore for the given privateKey with the given password.
     ///
     /// - parameter privateKey: The private key to encrypt.
     /// - parameter password: The password to use for the encryption.
+    /// - parameter kdf: The key derivation function to use.
+    /// - parameter cipher: The cipher to use for encryption.
+    /// - parameter checksum: The password checksum to use to detect bad passwords during decryption.
+    /// - parameter rounds: The number of rounds for the key derivation function to use. Defaults to a secure number.
     ///
-    /// - throws: Error if any step fails.
-//    public init(privateKey: [UInt8], password: String, kdf: Keystore.Crypto.KDFType = .scrypt, cipher: IVBlockModeType = .ctr, rounds: Int? = nil) throws {
-//        self = try KeystoreFactory.keystore(from: privateKey, password: password, kdf: kdf, cipher: cipher, rounds: rounds)
-//    }
+    /// - returns: The KeystoreETH2 object with the encrypted private key.
+    ///
+    /// - throws: Some `KeystoreETH2Factory.Error` if any step fails.
+    public init(
+        from privateKey: [UInt8],
+        password: String,
+        kdf: KeystoreETH2.KDFModule.KDFType,
+        cipher: KeystoreETH2.CipherModule.CipherType,
+        checksum: KeystoreETH2.ChecksumModule.ChecksumType,
+        rounds: Int = 262144
+    ) throws {
+        self = try KeystoreETH2Factory.keystore(from: privateKey, password: password, kdf: kdf, cipher: cipher, checksum: checksum, rounds: rounds)
+    }
 
     /// Extracts the private key from this keystore with the given password.
     ///
diff --git a/Sources/Keystore/KeystoreETH2Factory.swift b/Sources/Keystore/KeystoreETH2Factory.swift
@@ -1,6 +1,7 @@
 import Foundation
 import CryptoSwift
 import Scrypt
+import bls_framework
 
 public struct KeystoreETH2Factory {
 
@@ -11,7 +12,7 @@ public struct KeystoreETH2Factory {
     ///
     /// - returns: The extracted private key.
     ///
-    /// - throws: Some `KeystoreFactory.Error` if any step fails.
+    /// - throws: Some `KeystoreETH2Factory.Error` if any step fails.
     public static func privateKey(from keystore: KeystoreETH2, password: String) throws -> Array<UInt8> {
         var forbiddenCharacterSet = CharacterSet()
         forbiddenCharacterSet.insert(charactersIn: Unicode.Scalar(0x00)...Unicode.Scalar(0x1f))
@@ -63,6 +64,102 @@ public struct KeystoreETH2Factory {
         return try aes.decrypt([UInt8](ciphertextData))
     }
 
+    /// Creates a keystore for the given privateKey with the given password.
+    ///
+    /// - parameter privateKey: The private key to encrypt.
+    /// - parameter password: The password to use for the encryption.
+    /// - parameter kdf: The key derivation function to use.
+    /// - parameter cipher: The cipher to use for encryption.
+    /// - parameter checksum: The password checksum to use to detect bad passwords during decryption.
+    /// - parameter rounds: The number of rounds for the key derivation function to use. Defaults to a secure number.
+    ///
+    /// - returns: The KeystoreETH2 object with the encrypted private key.
+    ///
+    /// - throws: Some `KeystoreETH2Factory.Error` if any step fails.
+    public static func keystore(
+        from privateKey: [UInt8],
+        password: String,
+        kdf: KeystoreETH2.KDFModule.KDFType,
+        cipher: KeystoreETH2.CipherModule.CipherType,
+        checksum: KeystoreETH2.ChecksumModule.ChecksumType,
+        rounds: Int = 262144
+    ) throws -> KeystoreETH2 {
+        guard privateKey.count == 32 else {
+            throw Error.privateKeyMalformed
+        }
+
+        guard let iv = [UInt8].secureRandom(count: 16), let salt = [UInt8].secureRandom(count: 32) else {
+            throw Error.bytesGenerationFailed
+        }
+
+        let password = password.decomposedStringWithCompatibilityMapping
+            .components(separatedBy: .controlCharacters)
+            .filter({ !$0.isEmpty })
+            .joined(separator: "")
+
+        // Derive key
+        let kdfModule: KeystoreETH2.KDFModule
+        switch kdf {
+        case .scrypt:
+            kdfModule = .init(function: kdf, params: .init(salt: salt.toHexString(), dklen: 32, n: rounds, r: 8, p: 1), message: "")
+        case .pbkdf2:
+            kdfModule = .init(function: kdf, params: .init(salt: salt.toHexString(), dklen: 32, prf: "hmac-sha256", c: rounds), message: "")
+        }
+        let encryptionKey = try deriveKey(password: password, kdf: kdfModule)
+        guard encryptionKey.count >= 32 else {
+            throw Error.kdfFailed
+        }
+
+        // Encrypt
+        let usableKey = encryptionKey[0..<16]
+
+        guard cipher == .aes128Ctr else {
+            throw Error.cipherNotAvailable
+        }
+
+        let aes = try AES(
+            key: [UInt8](usableKey),
+            blockMode: CTR(iv: iv),
+            padding: .noPadding
+        )
+        let cipherMessage = try aes.encrypt(privateKey)
+        let cipherModule = KeystoreETH2.CipherModule(function: cipher, params: .init(iv: iv.toHexString()), message: cipherMessage.toHexString())
+
+        // Checksum
+        guard checksum == .sha256 else {
+            throw Error.checksumNotAvailable
+        }
+
+        let checksumGenerated = try generatePasswordChecksum(decryptionKey: encryptionKey, cipher: cipherModule)
+        let checksumModule = KeystoreETH2.ChecksumModule(function: checksum, params: .init(), message: checksumGenerated.toHexString())
+
+        // BLS pubkey
+
+        try BLSInterface.blsInit()
+
+        var serializedPrivateKey = privateKey
+        var secretKey = blsSecretKey.init()
+        if blsSecretKeyDeserialize(&secretKey, &serializedPrivateKey, numericCast(serializedPrivateKey.count)) <= 0 {
+            throw Error.privateKeyMalformed
+        }
+        var publicKey = blsPublicKey.init()
+        blsGetPublicKey(&publicKey, &secretKey)
+        // Ethereum public key is 48 bytes
+        var publicKeyBytes = Data(count: 48).bytes
+        blsPublicKeySerialize(&publicKeyBytes, 48, &publicKey)
+
+        let ethereumPublicKey = Data(publicKeyBytes)
+
+        return KeystoreETH2(
+            crypto: .init(kdf: kdfModule, checksum: checksumModule, cipher: cipherModule),
+            description: "",
+            pubkey: ethereumPublicKey.toHexString(),
+            path: "",
+            uuid: UUID().uuidString,
+            version: 4
+        )
+    }
+
     private static func deriveKey(
         password: String,
         kdf: KeystoreETH2.KDFModule
@@ -104,17 +201,25 @@ public struct KeystoreETH2Factory {
         ).calculate())
     }
 
-    private static func passwordVerification(
+    private static func generatePasswordChecksum(
         decryptionKey: Data,
-        cipher: KeystoreETH2.CipherModule,
-        checksum: KeystoreETH2.ChecksumModule
-    ) throws -> Bool {
+        cipher: KeystoreETH2.CipherModule
+    ) throws -> Data {
         let dkSlice = decryptionKey[16..<32]
         var preImage = dkSlice
         try preImage.append(contentsOf: cipher.message.dataWithHexString())
 
         let calculatedChecksum = preImage.sha256()
 
+        return calculatedChecksum
+    }
+
+    private static func passwordVerification(
+        decryptionKey: Data,
+        cipher: KeystoreETH2.CipherModule,
+        checksum: KeystoreETH2.ChecksumModule
+    ) throws -> Bool {
+        let calculatedChecksum = try generatePasswordChecksum(decryptionKey: decryptionKey, cipher: cipher)
         return try calculatedChecksum == checksum.message.dataWithHexString()
     }
 
@@ -141,6 +246,9 @@ public struct KeystoreETH2Factory {
         /// The given cipher is not available
         case cipherNotAvailable
 
+        /// The given checksum is not available
+        case checksumNotAvailable
+
         /// Generating random bytes failed
         case bytesGenerationFailed
 
diff --git a/Tests/KeystoreTests/KeystoreETH2Tests/ETH2PrivateKeyExtractionTests.swift b/Tests/KeystoreTests/KeystoreETH2Tests/ETH2PrivateKeyExtractionTests.swift
@@ -24,6 +24,7 @@ class ETH2PrivateKeyExtractionTests: QuickSpec {
                     "eth2_scrypt_test1"
                 ]
 
+                // decrypt
                 for testName in stubTests {
                     guard let testData = loadStub(named: testName), let test = try? decoder.decode(KeystoreETH2Stub.self, from: testData) else {
                         it("should never happen") {
@@ -39,7 +40,39 @@ class ETH2PrivateKeyExtractionTests: QuickSpec {
 
                     let actualPrivateKey = try? [UInt8](test.privateKey.dataWithHexString())
                     it("should extract the private key correctly") {
-                        expect(privateKey) == actualPrivateKey
+                        expect(privateKey!) == actualPrivateKey
+                    }
+                }
+
+                // encrypt then decrypt
+                for testName in stubTests {
+                    guard let testData = loadStub(named: testName), let test = try? decoder.decode(KeystoreETH2Stub.self, from: testData) else {
+                        it("should never happen") {
+                            fail("Stub \(testName) couldn't be loaded")
+                        }
+                        return
+                    }
+
+                    let encryptedKeystore = try? KeystoreETH2(
+                        from: test.privateKey.dataWithHexString().bytes,
+                        password: test.password,
+                        kdf: test.keystore.crypto.kdf.function,
+                        cipher: test.keystore.crypto.cipher.function,
+                        checksum: test.keystore.crypto.checksum.function,
+                        rounds: test.keystore.crypto.kdf.params.n ?? test.keystore.crypto.kdf.params.c ?? 262144
+                    )
+                    it("should not be nil") {
+                        expect(encryptedKeystore).toNot(beNil())
+                    }
+
+                    let decryptedPrivateKey = try? encryptedKeystore!.privateKey(password: test.password)
+                    it("should not be nil") {
+                        expect(decryptedPrivateKey).toNot(beNil())
+                    }
+
+                    let actualPrivateKey = try? [UInt8](test.privateKey.dataWithHexString())
+                    it("should extract the private key correctly") {
+                        expect(decryptedPrivateKey!) == actualPrivateKey
                     }
                 }
             }
