From 01f8994143617b457a9c3843fb54b969b11cd910 Mon Sep 17 00:00:00 2001
From: nift4 <nift4@protonmail.com>
Date: Thu, 1 Aug 2024 11:21:44 +0200
Subject: [PATCH] implement sha256 verification

---
 .../java/org/andbootmgr/app/CreatePartFlow.kt | 14 +++++++++---
 .../java/org/andbootmgr/app/WizardActivity.kt | 22 +++++++++++++++++--
 2 files changed, 31 insertions(+), 5 deletions(-)

diff --git a/app/src/main/java/org/andbootmgr/app/CreatePartFlow.kt b/app/src/main/java/org/andbootmgr/app/CreatePartFlow.kt
index db4a92ab..8bf58186 100644
--- a/app/src/main/java/org/andbootmgr/app/CreatePartFlow.kt
+++ b/app/src/main/java/org/andbootmgr/app/CreatePartFlow.kt
@@ -39,6 +39,7 @@ import org.json.JSONObject
 import org.json.JSONTokener
 import java.io.FileNotFoundException
 import java.net.URL
+import java.nio.charset.Charset
 
 class CreatePartWizardPageFactory(private val vm: WizardActivityState) {
 	fun get(): List<IWizardPage> {
@@ -851,10 +852,17 @@ private fun Download(c: CreatePartDataHolder) {
 											downloadedFile.delete()
 											downloading = false
 										}
+										val desiredHash = if (i == "_install.sh_") c.scriptShaInet!! else null
 
-										val sink = downloadedFile.sink().buffer()
-										sink.writeAll(response.body!!.source())
-										sink.close()
+										val rawSink = downloadedFile.sink()
+										val sink = if (desiredHash != null) HashingSink.sha256(rawSink) else rawSink
+										val buffer = sink.buffer()
+										buffer.writeAll(response.body!!.source())
+										buffer.close()
+										val realHash = if (desiredHash != null)
+												(sink as HashingSink).hash.hex() else null
+										if (desiredHash != null && realHash != desiredHash)
+											throw IllegalStateException("hash $realHash does not match expected hash $desiredHash")
 
 										if (!call.isCanceled())
 											c.chosen[i] = DledFile(null, downloadedFile)
diff --git a/app/src/main/java/org/andbootmgr/app/WizardActivity.kt b/app/src/main/java/org/andbootmgr/app/WizardActivity.kt
index f68c0d47..bfb61131 100644
--- a/app/src/main/java/org/andbootmgr/app/WizardActivity.kt
+++ b/app/src/main/java/org/andbootmgr/app/WizardActivity.kt
@@ -41,6 +41,8 @@ import java.io.OutputStream
 import java.net.URL
 import java.nio.file.Files
 import java.nio.file.StandardCopyOption
+import java.security.DigestInputStream
+import java.security.MessageDigest
 
 class WizardPageFactory(private val vm: WizardActivityState) {
 	fun get(flow: String): List<IWizardPage> {
@@ -177,6 +179,18 @@ class WizardActivity : ComponentActivity() {
 	}
 }
 
+private class ExpectedDigestInputStream(stream: InputStream?,
+                                        digest: MessageDigest?,
+                                        private val expectedDigest: String
+) : DigestInputStream(stream, digest) {
+	@OptIn(ExperimentalStdlibApi::class)
+	fun doAssert() {
+		val hash = digest.digest().toHexString()
+		if (hash != expectedDigest)
+			throw IllegalArgumentException("digest $hash does not match expected hash $expectedDigest")
+	}
+}
+
 class WizardActivityState(val codename: String) {
 	var btnsOverride = false
 	lateinit var navController: NavHostController
@@ -209,13 +223,14 @@ class WizardActivityState(val codename: String) {
 		inputStream.close()
 		outputStream.flush()
 		outputStream.close()
+		if (inputStream is ExpectedDigestInputStream)
+			inputStream.doAssert()
 		return nread
 	}
 
 	fun flashStream(flashType: String): InputStream {
-		// TODO check sha sum
 		return flashes[flashType]?.let {
-			when (it.first.scheme) {
+			val i = when (it.first.scheme) {
 				"content" ->
 					activity.contentResolver.openInputStream(it.first)
 						?: throw IOException("in == null")
@@ -225,6 +240,9 @@ class WizardActivityState(val codename: String) {
 					URL(it.first.toString()).openStream()
 				else -> null
 			}
+			if (it.second != null)
+				ExpectedDigestInputStream(i, MessageDigest.getInstance("SHA-256"), it.second!!)
+			else i
 		} ?: throw IllegalArgumentException()
 	}