From e78896b6c4c7bcf77a5b161907ea3c82f82590a5 Mon Sep 17 00:00:00 2001
From: Kyle Maas <kylemaasdev@gmail.com>
Date: Mon, 22 Jan 2024 19:14:58 +0000
Subject: [PATCH] Calculate md5sum on uploads while in tmp, deduplicate code

---
 files/helpers.py  | 20 ++++++++++++--------
 files/models.py   |  9 +++------
 files/tasks.py    |  6 ++----
 uploader/views.py |  7 +++++--
 4 files changed, 22 insertions(+), 20 deletions(-)

diff --git a/files/helpers.py b/files/helpers.py
index 5e78c59..536b018 100644
--- a/files/helpers.py
+++ b/files/helpers.py
@@ -270,13 +270,6 @@ def media_file_info(input_file):
         ret["fail"] = True
         return ret
 
-    cmd = ["md5sum", input_file]
-    stdout = run_command(cmd).get("out")
-    if stdout:
-        md5sum = stdout.split()[0]
-    else:
-        md5sum = ""
-
     cmd = [
         settings.FFPROBE_COMMAND,
         "-loglevel",
@@ -460,10 +453,21 @@ def media_file_info(input_file):
     ret["video_info"] = video_info
     ret["audio_info"] = audio_info
     ret["is_video"] = True
-    ret["md5sum"] = md5sum
     return ret
 
 
+def media_file_md5sum(input_file):
+    """
+    Get the md5sum of a file
+    """
+    cmd = ["md5sum", input_file]
+    stdout = run_command(cmd).get("out")
+    if stdout:
+        return stdout.split()[0]
+    else:
+        return ""
+
+
 def calculate_seconds(duration):
     # returns seconds, given a ffmpeg extracted string
     ret = 0
diff --git a/files/models.py b/files/models.py
index d46d094..88e3dd0 100644
--- a/files/models.py
+++ b/files/models.py
@@ -465,7 +465,8 @@ class Media(models.Model):
                     self.media_info = json.dumps(ret)
                 except TypeError:
                     self.media_info = ""
-                self.md5sum = ret.get("md5sum")
+                if not self.md5sum:
+                    self.md5sum = helpers.media_file_md5sum(self.media_file.path)
                 self.size = helpers.show_file_size(ret.get("file_size"))
             else:
                 self.media_type = ""
@@ -1123,11 +1124,7 @@ class Encoding(models.Model):
                 size = int(stdout.strip())
                 self.size = helpers.show_file_size(size)
         if self.chunk_file_path and not self.md5sum:
-            cmd = ["md5sum", self.chunk_file_path]
-            stdout = helpers.run_command(cmd).get("out")
-            if stdout:
-                md5sum = stdout.strip().split()[0]
-                self.md5sum = md5sum
+            self.md5sum = helpers.media_file_md5sum(self.chunk_file_path)
 
         super(Encoding, self).save(*args, **kwargs)
 
diff --git a/files/tasks.py b/files/tasks.py
index 9559fcb..ab6122e 100644
--- a/files/tasks.py
+++ b/files/tasks.py
@@ -29,6 +29,7 @@ from .helpers import (
     get_file_name,
     get_file_type,
     media_file_info,
+    media_file_md5sum,
     produce_ffmpeg_commands,
     produce_friendly_token,
     rm_file,
@@ -99,10 +100,7 @@ def chunkize_media(self, friendly_token, profiles, force=True):
     chunks_dict = {}
     # calculate once md5sums
     for chunk in chunks:
-        cmd = ["md5sum", chunk]
-        stdout = run_command(cmd).get("out")
-        md5sum = stdout.strip().split()[0]
-        chunks_dict[chunk] = md5sum
+        chunks_dict[chunk] = media_file_md5sum(chunk)
 
     for profile in profiles:
         if media.video_height and media.video_height < profile.resolution:
diff --git a/uploader/views.py b/uploader/views.py
index ce977ea..63471e3 100644
--- a/uploader/views.py
+++ b/uploader/views.py
@@ -9,7 +9,7 @@ from django.http import JsonResponse
 from django.views import generic
 
 from cms.permissions import user_allowed_to_upload
-from files.helpers import rm_file
+from files.helpers import media_file_md5sum, rm_file
 from files.models import Media
 
 from .fineuploader import ChunkedFineUploader
@@ -63,9 +63,12 @@ class FineUploaderView(generic.FormView):
             return self.make_response({"success": True})
         # create media!
         media_file = os.path.join(settings.MEDIA_ROOT, self.upload.real_path)
+        # Precalculate md5sum on tmp file
+        # If running on tmpfs, this should be much faster
+        md5 = media_file_md5sum(media_file)
         with open(media_file, "rb") as f:
             myfile = File(f)
-            new = Media.objects.create(media_file=myfile, user=self.request.user)
+            new = Media.objects.create(media_file=myfile, user=self.request.user, md5sum=md5)
         rm_file(media_file)
         shutil.rmtree(os.path.join(settings.MEDIA_ROOT, self.upload.file_path))
         return self.make_response({"success": True, "media_url": new.get_absolute_url()})