KineticCompat class for migrating legacy KinBase iOS private key accounts to Kinetic keypairs

KineticCompat.swift

//
//  KineticCompat.swift
//  pause-for
//
//  Created by Samuel Dowd on 10/26/22.
//  Copyright © 2022 Oroboros. All rights reserved.
//

import Foundation
import Promises
import KinBase
import Kinetic
import Solana
import struct Solana.PublicKey
import struct Solana.Transaction

class KineticCompat {
    static func moveKin(to destinationAddress: String, from kinEnvironment: KinEnvironment, appConfig: AppConfig, updateStatus: @escaping (String) -> ()) -> Promise<Void> {
        // First get all the keys stored in the KinBase storage
        return kinEnvironment.allAccountIds().then { (allPublicKeys) -> [KinBase.PublicKey] in
            let initialCount = allPublicKeys.count
            if initialCount > 0 {
                updateStatus("Removing unused wallets...")
                // Filter duplicated or non-signing wallets from list
                var keysToProcess: [KinBase.PublicKey] = []
                allPublicKeys.forEach { key in
                    if !keysToProcess.contains(key) && kinEnvironment.storage.hasPrivateKey(key) {
                        keysToProcess.append(key)
                    }
                }
                updateStatus("Found these keys stored: \(String(describing: keysToProcess))")
                return keysToProcess
            } else {
                updateStatus("Already clean, all done.")
                return []
            }
        // Then get the balance on each of those accounts
        }.then { keysToProcess -> Promise<[Void]> in
            all(keysToProcess.map { key -> Promise<Void> in
                Task {
                    let balance = try await AccountAPI.getBalance(
                        environment: appConfig.environment.name,
                        index: appConfig.app.index,
                        accountId: key.base58
                    ).balance
                    let balanceDouble = Double(balance)!
                    // If the account isn't empty, load its private key into kinetic and dump its kin into the new kinetic account
                    if balanceDouble > 0 {
                        updateStatus("\(key.base58): balance of \(balance), creating transfer")
                        let accountContext = KinAccountContext.Builder(env: kinEnvironment).useExistingAccount(key).build()
                        accountContext.storage.getAccount(key).then { account in
                            Task {
                                if let account = account {
                                    await moveKinFrom(
                                        kinBaseAccount: account,
                                        withBalance: balance,
                                        to: destinationAddress,
                                        appConfig: appConfig
                                    )
                                    updateStatus("Transferred \(balance) quarks from \(account.publicKey.base58) to \(destinationAddress)")
                                }
                            }
                        }
                    } else {
                        updateStatus("\(key.base58): Empty account, no op")
                    }
                }
                return Promise(())
            })
        }.then { _ in // Once all payments sent, return
            updateStatus("Threads spun up for each stored account")
            return Promise(())
        }
    }

    static func moveKinFrom(kinBaseAccount: KinAccount, withBalance balance: String, to destinationAddress: String, appConfig: AppConfig) async {
        let solanaAccount = LegacyAccount(account: kinBaseAccount)
        try? await makeTransfer(amount: balance, destination: destinationAddress, owner: solanaAccount, appConfig: appConfig)
    }

    struct LegacyAccount: Account {
        var publicKey: PublicKey
        var keyPair: KeyPair

        func sign(serializedMessage: Data) throws -> Data {
            keyPair.sign(serializedMessage)
        }

        init(account: KinAccount) {
            publicKey = PublicKey(string: account.publicKey.base58)!
            keyPair = KeyPair(publicKey: account.publicKey, privateKey: account.privateKey!)
        }
    }

    static func makeTransfer(
        amount: String,
        destination: String,
        owner: Account,
        appConfig: AppConfig
    ) async throws -> Kinetic.Transaction {
        let mint = appConfig.mint
        let amount = try addDecimals(amount: amount, decimals: 5).description

        let latestBlockhashResponse = try await TransactionAPI.getLatestBlockhash(environment: appConfig.environment.name, index: appConfig.app.index)

        let tx = try generateMakeTransferTransaction(
            amount: amount,
            blockhash: latestBlockhashResponse.blockhash,
            destination: destination,
            mintDecimals: mint.decimals,
            mintFeePayer: mint.feePayer,
            mintPublicKey: mint.publicKey,
            owner: owner
        )

        let makeTransferRequest = try MakeTransferRequest(
            commitment: Commitment.confirmed,
            environment: appConfig.environment.name,
            index: appConfig.app.index,
            mint: mint.publicKey,
            lastValidBlockHeight: latestBlockhashResponse.lastValidBlockHeight,
            referenceId: nil,
            referenceType: nil,
            tx: await serializeTransaction(tx)
        )
        return try await TransactionAPI.makeTransfer(makeTransferRequest: makeTransferRequest)
    }

    static func generateMakeTransferTransaction(
        amount: String,
        blockhash: String,
        destination: String,
        mintDecimals: Int,
        mintFeePayer: String,
        mintPublicKey: String,
        owner: Account
    ) throws -> Transaction {
        // Create objects from Response
        let mintKey = PublicKey(string: mintPublicKey)!
        let feePayerKey = PublicKey(string: mintFeePayer)!
        let ownerPublicKey = owner.publicKey

        let ownerTokenAccount = try getAssociatedTokenAddress(ownerPublicKey: ownerPublicKey, mintKey: mintKey)
        let destinationTokenAccount = try getAssociatedTokenAddress(ownerPublicKey: PublicKey(string: destination)!, mintKey: mintKey)

        var instructions: [TransactionInstruction] = []

        instructions.append(
            TokenProgram.transferCheckedInstruction(
                programId: PublicKey.tokenProgramId,
                source: ownerTokenAccount,
                mint: mintKey,
                destination: destinationTokenAccount,
                owner: ownerPublicKey,
                multiSigners: [],
                amount: UInt64(amount)!,
                decimals: UInt8(mintDecimals)
            )
        )

        var transaction = Transaction(
            signatures: [Transaction.Signature(signature: nil, publicKey: ownerPublicKey)],
            feePayer: feePayerKey,
            instructions: instructions,
            recentBlockhash: blockhash
        )

        transaction.partialSign(signers: [owner])

        return transaction
    }

    static func getAssociatedTokenAddress(ownerPublicKey: PublicKey, mintKey: PublicKey) throws -> PublicKey {
        guard case let .success(ownerTokenAccount) = PublicKey.associatedTokenAddress(walletAddress: ownerPublicKey, tokenMintAddress: mintKey)
        else {
            throw NSError()
        }
        return ownerTokenAccount
    }

    static func serializeTransaction(_ transaction: Transaction) async throws -> String {
        var transaction = transaction
        let serializedRes: Data? = await withCheckedContinuation { continuation in
            let serializedRes = transaction.serialize(requiredAllSignatures: false, verifySignatures: false)
            switch serializedRes {
            case .success(let serialized):
                continuation.resume(returning: serialized)
            case .failure(let e):
                continuation.resume(returning: nil)
            }
        }
        guard let serialized = serializedRes else {
            throw NSError()
        }
        return serialized.base64EncodedString()
    }
}