From 2af0c711f7b3852046f4d436cd55a8a85dd89315 Mon Sep 17 00:00:00 2001
From: 6od9i <6od911@gmail.com>
Date: Thu, 29 Jun 2023 11:01:07 +0400
Subject: [PATCH] Addresses generator (#1)

- Generation amount of addresses added
- Generation wallet with index > 9 fixed
- Test generation addresses added
---
 .../KeystoreManager/BIP32Keystore.swift       | 53 +++++++++++++++++-
 .../localTests/BIP32KeystoreTests.swift       | 54 +++++++++++++++++++
 2 files changed, 105 insertions(+), 2 deletions(-)
 create mode 100644 Tests/web3swiftTests/localTests/BIP32KeystoreTests.swift

diff --git a/Sources/Web3Core/KeystoreManager/BIP32Keystore.swift b/Sources/Web3Core/KeystoreManager/BIP32Keystore.swift
index 4277cca61..003ba7fd5 100755
--- a/Sources/Web3Core/KeystoreManager/BIP32Keystore.swift
+++ b/Sources/Web3Core/KeystoreManager/BIP32Keystore.swift
@@ -177,7 +177,7 @@ public class BIP32Keystore: AbstractKeystore {
         if path.hasPrefix(prefixPath) {
             let upperIndex = (path.range(of: prefixPath)?.upperBound)!
             if upperIndex < path.endIndex {
-                pathAppendix = String(path[path.index(after: upperIndex)])
+                pathAppendix = String(path[path.index(after: upperIndex)..<path.endIndex])
             } else {
                 throw AbstractKeystoreError.encryptionError("out of bounds")
             }
@@ -215,7 +215,56 @@ public class BIP32Keystore: AbstractKeystore {
         guard let serializedRootNode = rootNode.serialize(serializePublic: false) else {throw AbstractKeystoreError.keyDerivationError}
         try encryptDataToStorage(password, data: serializedRootNode, aesMode: self.keystoreParams!.crypto.cipher)
     }
-
+    
+    /// Fast generation addresses for current account
+    /// used for shows wich address user wiil get when changed number of his wallet
+    /// - Parameters:
+    ///   - password: password of seed storage
+    ///   - preffixPath: preffix of Derivation Path without account number
+    ///   - number: number of wallets adresses needed to generate from 0 to number-1
+    /// - Returns: Array of addresses generated from 0 to number bound, or empty array in case of error
+    public func getAddressForAccount(password: String, preffixPath: String, number: Int) -> [EthereumAddress] {
+        guard let decryptedRootNode = try? getPrefixNodeData(password) else {
+            return []
+        }
+        guard let rootNode = HDNode(decryptedRootNode) else {
+            return []
+        }
+        let prefixPath = self.rootPrefix
+        var pathAppendix: String?
+        
+        return [Int](0..<number).compactMap({ number in
+            pathAppendix = nil
+            let path = preffixPath + "/\(number)"
+            if path.hasPrefix(prefixPath) {
+                let upperIndex = (path.range(of: prefixPath)?.upperBound)!
+                if upperIndex < path.endIndex {
+                    pathAppendix = String(path[path.index(after: upperIndex)..<path.endIndex])
+                } else {
+                    return nil
+                }
+                
+                guard pathAppendix != nil else {
+                    return nil
+                }
+                if pathAppendix!.hasPrefix("/") {
+                    pathAppendix = pathAppendix?.trimmingCharacters(in: CharacterSet.init(charactersIn: "/"))
+                }
+            } else {
+                if path.hasPrefix("/") {
+                    pathAppendix = path.trimmingCharacters(in: CharacterSet.init(charactersIn: "/"))
+                }
+            }
+            guard pathAppendix != nil,
+                  rootNode.depth == prefixPath.components(separatedBy: "/").count - 1,
+                  let newNode = rootNode.derive(path: pathAppendix!, derivePrivateKey: true),
+                  let newAddress = Utilities.publicToAddress(newNode.publicKey) else {
+                return nil
+            }
+            return newAddress
+        })
+    }
+    
     fileprivate func encryptDataToStorage(_ password: String, data: Data, dkLen: Int = 32, N: Int = 4096, R: Int = 6, P: Int = 1, aesMode: String = "aes-128-cbc") throws {
         guard data.count == 82 else {
             throw AbstractKeystoreError.encryptionError("Invalid expected data length")
diff --git a/Tests/web3swiftTests/localTests/BIP32KeystoreTests.swift b/Tests/web3swiftTests/localTests/BIP32KeystoreTests.swift
new file mode 100644
index 000000000..0ae108db4
--- /dev/null
+++ b/Tests/web3swiftTests/localTests/BIP32KeystoreTests.swift
@@ -0,0 +1,54 @@
+//
+//  BIP32KeystoreTests.swift
+//  localTests
+//
+//  Created by 6od9i on 29.06.2023.
+//
+
+import Foundation
+import XCTest
+import Web3Core
+
+@testable import web3swift
+
+class BIP32KeystoreTests: XCTestCase {
+    func testAddressGeneration() throws {
+        /// Seed randomly generated for this test
+        let mnemonic = "resource beyond merit enemy foot piece reveal eagle nothing luggage goose spot"
+        let password = "test_password"
+        
+        let addressesCount = 101
+        
+        guard let keystore = try BIP32Keystore(
+            mnemonics: mnemonic,
+            password: password,
+            mnemonicsPassword: "",
+            language: .english,
+            prefixPath: HDNode.defaultPathMetamaskPrefix) else {
+            XCTFail("Keystore has not generated")
+            throw NSError(domain: "0", code: 0)
+        }
+        
+        let addresses = keystore.getAddressForAccount(password: password,
+                                                      preffixPath: HDNode.defaultPathMetamaskPrefix,
+                                                      number: addressesCount)
+        XCTAssertEqual(addresses.count, addressesCount)
+        XCTAssertNotEqual(addresses[11], addresses[1])
+        
+        guard let sameKeystore = try BIP32Keystore(
+            mnemonics: mnemonic,
+            password: password,
+            mnemonicsPassword: "",
+            language: .english,
+            prefixPath: HDNode.defaultPathMetamaskPrefix) else {
+            XCTFail("Keystore has not generated")
+            throw NSError(domain: "0", code: 0)
+        }
+        
+        let walletNumber = addressesCount - 1
+        try sameKeystore.createNewCustomChildAccount(password: password,
+                                                     path: HDNode.defaultPathMetamaskPrefix + "/\(walletNumber)")
+        let address = sameKeystore.addresses?.last?.address
+        XCTAssertEqual(addresses.last?.address, address)
+    }
+}