From 7a4019c5baba6caf9c928a1aac61296bb7d97e63 Mon Sep 17 00:00:00 2001
From: Simon McLoughlin <simon.d.mcl@gmail.com>
Date: Mon, 25 Sep 2023 12:03:12 +0100
Subject: [PATCH] add checksum validity check to Mnemonic isValid func

---
 .../KukaiCryptoSwift/Mnemonic/Mnemonic.swift  | 38 ++++++++++++++++++-
 .../KukaiCryptoSwiftTests/MnemonicTests.swift |  5 ++-
 2 files changed, 41 insertions(+), 2 deletions(-)

diff --git a/Sources/KukaiCryptoSwift/Mnemonic/Mnemonic.swift b/Sources/KukaiCryptoSwift/Mnemonic/Mnemonic.swift
index cc9c01e..21362d3 100644
--- a/Sources/KukaiCryptoSwift/Mnemonic/Mnemonic.swift
+++ b/Sources/KukaiCryptoSwift/Mnemonic/Mnemonic.swift
@@ -3,6 +3,7 @@
 
 import Foundation
 import CommonCrypto
+import CryptoKit
 
 public enum MnemonicError: Swift.Error {
 	case seedDerivationFailed
@@ -122,6 +123,41 @@ public struct Mnemonic: Equatable, Codable {
 		phrase = String(repeating: "0", count: phrase.count)
 	}
 	
+	/**
+	 Derive the checksum portion of an array of bits
+	 */
+	public static func deriveChecksumBits(_ bytes: [UInt8]) -> String {
+		let ENT = bytes.count * 8;
+		let CS = ENT / 32
+		
+		let hash = SHA256.hash(data: bytes)
+		let hashbits = String(hash.flatMap { ("00000000" + String($0, radix:2)).suffix(8) })
+		return String(hashbits.prefix(CS))
+	}
+	
+	/**
+	 Verify the chechsum of the supplied words to esnure its a valid phrase
+	 */
+	public static func isValidChecksum(phrase: [String], wordlist: WordList = WordList.english) -> Bool {
+		let wordL = wordlist.words
+		var bits = ""
+		for word in phrase {
+			guard let i = wordL.firstIndex(of: word) else { return false }
+			bits += ("00000000000" + String(i, radix: 2)).suffix(11)
+		}
+		
+		let dividerIndex = bits.count / 33 * 32
+		let entropyBits = String(bits.prefix(dividerIndex))
+		let checksumBits = String(bits.suffix(bits.count - dividerIndex))
+		
+		let regex = try! NSRegularExpression(pattern: "[01]{1,8}", options: .caseInsensitive)
+		let entropyBytes = regex.matches(in: entropyBits, options: [], range: NSRange(location: 0, length: entropyBits.count)).map {
+			UInt8(strtoul(String(entropyBits[Range($0.range, in: entropyBits)!]), nil, 2))
+		}
+		
+		return checksumBits == deriveChecksumBits(entropyBytes)
+	}
+	
 	/**
 	 Check a mnemonic is of the correct length, and is made up of valid BIP39 words
 	 */
@@ -140,6 +176,6 @@ public struct Mnemonic: Equatable, Codable {
 			}
 		}
 		
-		return true
+		return Mnemonic.isValidChecksum(phrase: words, wordlist: vocabulary)
 	}
 }
diff --git a/Tests/KukaiCryptoSwiftTests/MnemonicTests.swift b/Tests/KukaiCryptoSwiftTests/MnemonicTests.swift
index 9e18efb..d6c7ae8 100644
--- a/Tests/KukaiCryptoSwiftTests/MnemonicTests.swift
+++ b/Tests/KukaiCryptoSwiftTests/MnemonicTests.swift
@@ -68,12 +68,15 @@ final class MnemonicTests: XCTestCase {
 		XCTAssert(mnemonic3.isValid() == false)
 		
 		let mnemonic4 = try Mnemonic(seedPhrase: "remember smile trip tumble era cube worry fuel bracket eight kitten inform remember smile trip tumble era cube worry fuel bracket eight kitten inform")
-		XCTAssert(mnemonic4.isValid() == true)
+		XCTAssert(mnemonic4.isValid() == false)
 		
 		let mnemonic5 = try Mnemonic(seedPhrase: "remember smile trip tumble era cube worry fuel bracket eight kitten inform remember smile trip tumble era cube worry fuel bracket eight kitten infomr")
 		XCTAssert(mnemonic5.isValid() == false)
 		
 		let mnemonic6 = try Mnemonic(seedPhrase: "tell me more about your awesome but totally invalid mnemonic word1 word2")
 		XCTAssert(mnemonic6.isValid() == false)
+		
+		let mnemonic7 = try Mnemonic(seedPhrase: "remember smile trip tumble era cube worry fuel bracket eight kitten remember")
+		XCTAssert(mnemonic7.isValid() == false)
 	}
 }