From be806a368987d8a17a796e1152e507cdb4cf4bad Mon Sep 17 00:00:00 2001
From: Tanmoy Sarkar <57363826+tanmoysrt@users.noreply.github.com>
Date: Thu, 28 Nov 2024 09:44:18 +0530
Subject: [PATCH] feat: show hash of downloaded backup files on site restore
 failure (#144)

---
 agent/site.py  | 25 ++++++++++++++++++++++++-
 agent/utils.py | 22 ++++++++++++++++++++++
 2 files changed, 46 insertions(+), 1 deletion(-)

diff --git a/agent/site.py b/agent/site.py
index 1f13d968..25566118 100644
--- a/agent/site.py
+++ b/agent/site.py
@@ -14,7 +14,7 @@
 from agent.base import AgentException, Base
 from agent.database import Database
 from agent.job import job, step
-from agent.utils import b2mb, get_size
+from agent.utils import b2mb, compute_file_hash, get_size
 
 if TYPE_CHECKING:
     from agent.bench import Bench
@@ -129,6 +129,26 @@ def restore(
         finally:
             self.bench.drop_mariadb_user(self.name, mariadb_root_password, self.database)
 
+    @step("Checksum of Downloaded Backup Files")
+    def calculate_checksum_of_backup_files(self, database_file, public_file, private_file):
+        database_file_sha256 = compute_file_hash(database_file, algorithm="sha256", raise_exception=False)
+
+        data = f"""Database File
+> File Name - {os.path.basename(database_file)}
+> SHA256 Checksum - {database_file_sha256}\n"""
+        if public_file:
+            public_file_sha256 = compute_file_hash(public_file, algorithm="sha256", raise_exception=False)
+            data += f"""\nPublic File
+> File Name - {os.path.basename(public_file)}
+> SHA256 Checksum - {public_file_sha256}\n"""
+        if private_file:
+            private_file_sha256 = compute_file_hash(private_file, algorithm="sha256", raise_exception=False)
+            data += f"""\nPrivate File
+> File Name - {os.path.basename(private_file)}
+> SHA256 Checksum - {private_file_sha256}\n"""
+
+        return {"output": data}
+
     @job("Restore Site")
     def restore_job(
         self,
@@ -149,6 +169,9 @@ def restore_job(
                 files["public"],
                 files["private"],
             )
+        except Exception:
+            self.calculate_checksum_of_backup_files(files["database"], files["public"], files["private"])
+            raise
         finally:
             self.bench.delete_downloaded_files(files["directory"])
         self.uninstall_unavailable_apps(apps)
diff --git a/agent/utils.py b/agent/utils.py
index 47539fd9..1ff580d9 100644
--- a/agent/utils.py
+++ b/agent/utils.py
@@ -1,5 +1,6 @@
 from __future__ import annotations
 
+import hashlib
 import os
 from datetime import datetime, timedelta
 from math import ceil
@@ -116,3 +117,24 @@ def end_execution(
     res["status"] = status or "Success"
     res["output"] = output or res["output"]
     return res
+
+
+def compute_file_hash(file_path, algorithm="sha256", raise_exception=True):
+    try:
+        """Compute the hash of a file using the specified algorithm."""
+        hash_func = hashlib.new(algorithm)
+
+        with open(file_path, "rb") as file:
+            # read in 10MB chunks
+            while chunk := file.read(10000000):
+                hash_func.update(chunk)
+
+        return hash_func.hexdigest()
+    except FileNotFoundError:
+        if raise_exception:
+            raise
+        return "File does not exist"
+    except Exception:
+        if raise_exception:
+            raise
+        return "Failed to compute hash"