From 0f167cf29fd51b9e95146c4e93dd0f046e82de44 Mon Sep 17 00:00:00 2001
From: Alejandro Sirgo Rica <asirgo@soleta.eu>
Date: Wed, 2 Oct 2024 14:50:14 +0200
Subject: src: consolidate compute_md5 functions

Add compute_md5 function in src/utils/fs.py

Remove identical md5 functions from src/live/ogOperations.py and
src/utils/tiptorrent.py

Move error checks from ogOperations.py into compute_md5 function in
src/utils/fs.py
---
 src/live/ogOperations.py | 18 +-----------------
 src/utils/fs.py          | 17 +++++++++++++++++
 src/utils/tiptorrent.py  | 14 ++------------
 3 files changed, 20 insertions(+), 29 deletions(-)

(limited to 'src')

diff --git a/src/live/ogOperations.py b/src/live/ogOperations.py
index 83242d2..f1478c3 100644
--- a/src/live/ogOperations.py
+++ b/src/live/ogOperations.py
@@ -200,22 +200,6 @@ class OgLiveOperations:
 
         return data
 
-    def _compute_md5(self, path, bs=2**20):
-        m = hashlib.md5()
-        with open(path, 'rb') as f:
-            while True:
-                buf = f.read(bs)
-                if not buf:
-                    break
-                m.update(buf)
-        return m.hexdigest()
-
-    def _md5_file(self, path):
-        if not os.path.exists(path):
-            raise OgError(f'Failed to calculate checksum, image file {path} does not exist')
-
-        return self._compute_md5(path)
-
     def _write_md5_file(self, path, checksum):
         if not os.path.exists(path):
             raise OgError(f'Failed to calculate checksum, image file {path} does not exist')
@@ -725,7 +709,7 @@ class OgLiveOperations:
         logging.info(f'Writing checksum file {name}.img.full.sum...')
         logging.info('*DO NOT REBOOT OR POWEROFF* the client during this time')
 
-        checksum = self._md5_file(f'/opt/opengnsys/images/{name}.img')
+        checksum = compute_md5(f'/opt/opengnsys/images/{name}.img')
         if checksum == -1:
             raise OgError(f'Cannot access {name}.full.sum file')
 
diff --git a/src/utils/fs.py b/src/utils/fs.py
index d6cf36a..8446f08 100644
--- a/src/utils/fs.py
+++ b/src/utils/fs.py
@@ -8,6 +8,7 @@
 
 import logging
 import os
+import hashlib
 import subprocess
 import shlex
 from src.log import OgError
@@ -340,3 +341,19 @@ def _extend_ntfsresize(partdev):
     proc = subprocess.run(cmd, input=b'y')
     if proc.returncode != 0:
         raise OgError(f'Error growing ntfs filesystem at {partdev}')
+
+def compute_md5(path, bs=2**20):
+    if not os.path.exists(path):
+        raise OgError(f"Failed to calculate checksum, image file {path} does not exist")
+
+    m = hashlib.md5()
+    try:
+        with open(path, 'rb') as f:
+            while True:
+                buf = f.read(bs)
+                if not buf:
+                    break
+                m.update(buf)
+    except Exception as e:
+        raise OgError(f'Failed to calculate checksum for {path}: {e}') from e
+    return m.hexdigest()
diff --git a/src/utils/tiptorrent.py b/src/utils/tiptorrent.py
index 095f17f..7392725 100644
--- a/src/utils/tiptorrent.py
+++ b/src/utils/tiptorrent.py
@@ -6,7 +6,6 @@
 # Free Software Foundation; either version 3 of the License, or
 # (at your option) any later version.
 
-import hashlib
 import logging
 import os
 import shlex
@@ -15,16 +14,7 @@ import subprocess
 import urllib.request
 from src.log import OgError
 from src.utils.cache import *
-
-def _compute_md5(path, bs=2**20):
-    m = hashlib.md5()
-    with open(path, 'rb') as f:
-        while True:
-            buf = f.read(bs)
-            if not buf:
-                break
-            m.update(buf)
-    return m.hexdigest()
+from src.utils.fs import compute_md5
 
 
 def tip_fetch_csum(tip_addr, image_name):
@@ -64,7 +54,7 @@ def tip_write_csum(image_name):
     logging.info('*DO NOT REBOOT OR POWEROFF* the client during this time')
 
     filename = image_path + ".full.sum"
-    csum = _compute_md5(image_path)
+    csum = compute_md5(image_path)
     try:
         with open(filename, 'w') as f:
             f.write(csum)
-- 
cgit v1.2.3-18-g5258