From a49a911e03255ba143f96003dfd8fe956e1d7632 Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Thu, 15 Jun 2023 16:19:35 +0800
Subject: [PATCH] Replace image rendered in the HTML to thumbnail if possible

---
 Whatsapp_Chat_Exporter/data_model.py |  1 +
 Whatsapp_Chat_Exporter/extract.py    | 21 +++++++++++++++++++--
 Whatsapp_Chat_Exporter/whatsapp.html |  4 ++--
 3 files changed, 22 insertions(+), 4 deletions(-)

diff --git a/Whatsapp_Chat_Exporter/data_model.py b/Whatsapp_Chat_Exporter/data_model.py
index 51463f1..5859dc5 100644
--- a/Whatsapp_Chat_Exporter/data_model.py
+++ b/Whatsapp_Chat_Exporter/data_model.py
@@ -61,6 +61,7 @@ class Message():
         self.reply = None
         self.quoted_data = None
         self.caption = None
+        self.thumb = None # Android specific
     
     def to_json(self):
         return {
diff --git a/Whatsapp_Chat_Exporter/extract.py b/Whatsapp_Chat_Exporter/extract.py
index ed67227..8d1ea58 100644
--- a/Whatsapp_Chat_Exporter/extract.py
+++ b/Whatsapp_Chat_Exporter/extract.py
@@ -11,6 +11,7 @@ import hmac
 from pathlib import Path
 from mimetypes import MimeTypes
 from hashlib import sha256
+from base64 import b64decode
 from Whatsapp_Chat_Exporter.data_model import ChatStore, Message
 from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, Device, rendering, sanitize_except, determine_day, Crypt
 from Whatsapp_Chat_Exporter.utility import brute_force_offset, CRYPT14_OFFSETS
@@ -412,10 +413,14 @@ def media(db, data, media_folder):
                         file_path,
                         message_url,
                         mime_type,
-                        media_key
+                        media_key,
+                        file_hash,
+						thumbnail
                  FROM message_media
                     INNER JOIN messages
                         ON message_media.message_row_id = messages._id
+					LEFT JOIN media_hash_thumbnail
+						ON message_media.file_hash = media_hash_thumbnail.media_hash
                 ORDER BY messages.key_remote_jid ASC"""
         )
     except sqlite3.OperationalError:
@@ -424,7 +429,9 @@ def media(db, data, media_folder):
                     file_path,
                     message_url,
                     mime_type,
-                    media_key
+                    media_key,
+                    file_hash,
+                    thumbnail
                 FROM message_media
                 INNER JOIN message
                     ON message_media.message_row_id = message._id
@@ -432,10 +439,14 @@ def media(db, data, media_folder):
                     ON chat._id = message.chat_row_id
                 INNER JOIN jid
                     ON jid._id = chat.jid_row_id
+                LEFT JOIN media_hash_thumbnail
+						ON message_media.file_hash = media_hash_thumbnail.media_hash
                 ORDER BY jid.raw_string ASC"""
         )
     content = c.fetchone()
     mime = MimeTypes()
+    if not os.path.isdir(f"{media_folder}/thumbnails"):
+        os.mkdir(f"{media_folder}/thumbnails")
     while content is not None:
         file_path = f"{media_folder}/{content['file_path']}"
         message = data[content["key_remote_jid"]].messages[content["message_row_id"]]
@@ -465,6 +476,12 @@ def media(db, data, media_folder):
             message.data = "The media is missing"
             message.mime = "media"
             message.meta = True
+        if content["thumbnail"] is not None:
+            thumb_path = f"{media_folder}/thumbnails/{b64decode(content['file_hash']).hex()}.png"
+            if not os.path.isfile(thumb_path):
+                with open(thumb_path, "wb") as f:
+                    f.write(content["thumbnail"])
+            message.thumb = thumb_path
         i += 1
         if i % 100 == 0:
             print(f"Processing media...({i}/{total_row_number})", end="\r")
diff --git a/Whatsapp_Chat_Exporter/whatsapp.html b/Whatsapp_Chat_Exporter/whatsapp.html
index fc960c6..7ddf58c 100644
--- a/Whatsapp_Chat_Exporter/whatsapp.html
+++ b/Whatsapp_Chat_Exporter/whatsapp.html
@@ -128,7 +128,7 @@
 								{{ msg.data | sanitize_except() }}
 								{% else %}
 									{% if "image/" in msg.mime %}
-									<a href="{{ msg.data }}"><img src="{{ msg.data }}" /></a>
+									<a href="{{ msg.data }}"><img src="{{ msg.thumb if msg.thumb is not none else msg.data }}" /></a>
 									{% elif "audio/" in msg.mime %}
 									<audio controls="controls" autobuffer="autobuffer">
 										<source src="{{ msg.data }}" />
@@ -206,7 +206,7 @@
 								{{ msg.data | sanitize_except() }}
 								{% else %}
 									{% if "image/" in msg.mime %}
-									<a href="{{ msg.data }}"><img src="{{ msg.data }}" /></a>
+									<a href="{{ msg.data }}"><img src="{{ msg.thumb if msg.thumb is not none else msg.data }}" /></a>
 									{% elif "audio/" in msg.mime %}
 									<audio controls preload="auto">
 										<source src="{{ msg.data }}" />