From 57ff483fbedcb538da4fb0527a6c19c06b13c2ed Mon Sep 17 00:00:00 2001
From: smartfrigde <37928912+smartfrigde@users.noreply.github.com>
Date: Fri, 4 Feb 2022 19:08:11 +0100
Subject: [PATCH] Add checksum checking + backups

---
 README.md                        |  1 +
 examples/example_build_info.json |  3 +-
 src/Updater.nim                  | 54 ++++++++++++++++++++------------
 3 files changed, 37 insertions(+), 21 deletions(-)

diff --git a/README.md b/README.md
index 1e1e2e2..e2e2b76 100644
--- a/README.md
+++ b/README.md
@@ -5,6 +5,7 @@ Experimental updater for ArmCord made in Nim.
 - **Quick and lightweight** 
 - **One binary is whole updater**
 - **Supports custom update servers**
+- **Checks for checksums**
 ## How to use it?
 ### 1. Client side:
 - Download the latest version of the updater from [here](https://github.com/ArmCord/Updater/releases/latest).
diff --git a/examples/example_build_info.json b/examples/example_build_info.json
index 65b4e9f..ba4df8e 100644
--- a/examples/example_build_info.json
+++ b/examples/example_build_info.json
@@ -1,5 +1,6 @@
 {
   "version": "3.0.4",
   "update_endpoint": "https://armcord.xyz/",
-  "note": "https://discord.gg/uaW5vMY3V6"
+  "note": "https://discord.gg/uaW5vMY3V6",
+  "md5": "3d605bd6a1447afc5e4239b44cf393a3"
 }
diff --git a/src/Updater.nim b/src/Updater.nim
index a2152be..8b87bc1 100644
--- a/src/Updater.nim
+++ b/src/Updater.nim
@@ -1,7 +1,9 @@
 import puppy
 import std/json
 import std/strutils
-from os import fileExists
+import std/os
+import std/md5
+var attempts = 1
 proc checkInternetConnection() =
   try:
     discard fetch("https://example.com/")
@@ -15,7 +17,33 @@ proc checkIfValidDirectory() =
   else:
     echo("Please place the updater in correct path (resources folder). Exiting.")
     quit(1)
-
+proc downloadUpdate(endpoint:string) =
+    echo("New version available!")
+    echo("Getting the latest asar.")
+    try:
+      moveFile("app.asar", "app.asar.old")
+      moveFile("build_info.json", "build_info.json.old")
+      writeFile("app.asar", fetch(endpoint & "app.asar"))
+      writeFile("build_info.json", fetch(endpoint & "latest.json"))
+    except Exception as e:
+      echo("Failed to download the latest asar: ", e.msg)
+      quit(1)
+    echo ("Downloaded the latest asar. Checking checksums")
+    if (parseJson(readFile("build_info.json"))["md5"].getStr() == getMD5(readFile("app.asar"))):
+      echo("Checksums match..")
+    else:
+      attempts = attempts + 1
+      if (attempts > 3):
+        moveFile("app.asar.old", "app.asar")
+        moveFile("build_info.json.old", "build_info.json")
+        echo("Checksums didn't match after 2 attempts to download the update. Reverting the update. Exiting.")
+        quit(1)
+      else:
+        echo("Checksums do not match. Update failed! Retrying")
+        downloadUpdate(endpoint)
+    removeFile("app.asar.old")
+    removeFile("build_info.json.old")
+    echo("Update installation has finished. You should be able to use the new version of ArmCord.")
 when isMainModule:
   echo("ArmCord Updater")
   echo("Platform: " & hostOS)
@@ -33,15 +61,7 @@ when isMainModule:
   var a = parseInt(latestVersion.replace(".", "")) # Remove the dots to compare
   var b = parseInt(currentVersion.replace(".", "")) # Remove the dots to compare
   if (a > b):
-    echo("New version available!")
-    echo("Getting the latest asar.")
-    try:
-      writeFile("app.asar", fetch(endpoint & "app.asar"))
-      writeFile("build_info.json", fetch(endpoint & "latest.json"))
-    except Exception as e:
-      echo("Failed to download the latest asar: ", e.msg)
-      quit(1)
-    echo("Update installation has finished. You should be able to use the new version of ArmCord.")
+    downloadUpdate(endpoint)
   else:
     echo("No new version available.")
     if (a < b):
@@ -50,13 +70,7 @@ when isMainModule:
       write(stdout, "[y/n]: ")
       var input = readLine(stdin)
       if (input == "y"):
-        echo("Getting the latest asar.")
-        try:
-          writeFile("app.asar", fetch(endpoint & "app.asar"))
-          writeFile("build_info.json", fetch(endpoint & "latest.json"))
-        except Exception as e:
-          echo("Failed to download the latest asar: ", e.msg)
-          quit(1)
-        echo("Update installation has finished. You should be able to use the new version of ArmCord.")
+        downloadUpdate(endpoint)
       else:
-        echo("Downgrading won't be performed.")
\ No newline at end of file
+        echo("Downgrading won't be performed.")
+         
\ No newline at end of file