From 506f8e89f48632c06768ce201a3e6108619b3cb6 Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Fri, 16 Jun 2023 19:19:33 +0800
Subject: [PATCH] Implement import json and output HTML

---
 Whatsapp_Chat_Exporter/__main__.py | 35 +++++++++++++++++++++-----
 Whatsapp_Chat_Exporter/utility.py  | 40 +++++++++++++++++++++++++++---
 2 files changed, 66 insertions(+), 9 deletions(-)

diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py
index 88f4ad1..9d0f4e9 100644
--- a/Whatsapp_Chat_Exporter/__main__.py
+++ b/Whatsapp_Chat_Exporter/__main__.py
@@ -6,7 +6,7 @@ import glob
 from Whatsapp_Chat_Exporter import extract, extract_exported, extract_iphone
 from Whatsapp_Chat_Exporter import extract_iphone_media
 from Whatsapp_Chat_Exporter.data_model import ChatStore
-from Whatsapp_Chat_Exporter.utility import Crypt, check_update
+from Whatsapp_Chat_Exporter.utility import Crypt, check_update, import_from_json
 from argparse import ArgumentParser, SUPPRESS
 import os
 import sqlite3
@@ -168,6 +168,13 @@ def main():
         action='store_true',
         help="Do not render avatar in HTML output"
     )
+    parser.add_argument(
+        "--import",
+        dest="import_json",
+        default=False,
+        action='store_true',
+        help="Import JSON file and convert to HTML output"
+    )
     args = parser.parse_args()
 
     # Check for updates
@@ -175,15 +182,21 @@ def main():
         exit(check_update())
 
     # Sanity checks
-    if args.android and args.iphone and args.exported:
+    if args.android and args.iphone and args.exported and args.import_json:
         print("You must define only one device type.")
         exit(1)
-    if not args.android and not args.iphone and not args.exported:
+    if not args.android and not args.iphone and not args.exported and not args.import_json:
         print("You must define the device type.")
         exit(1)
     if args.no_html and not args.json:
         print("You must either specify a JSON output file or enable HTML output.")
         exit(1)
+    if args.import_json and (args.android or args.iphone or args.exported or args.no_html):
+        print("You can only use --import with -j and without --no-html.")
+        exit(1)
+    elif args.import_json and not os.path.isfile(args.json):
+        print("JSON file not found.")
+        exit(1)
 
     data = {}
 
@@ -264,7 +277,7 @@ def main():
         else:
             contact_db = args.wa
 
-    if not args.exported:
+    if not args.exported and not args.import_json:
         if os.path.isfile(msg_db):
             with sqlite3.connect(msg_db) as db:
                 db.row_factory = sqlite3.Row
@@ -307,7 +320,7 @@ def main():
                     except PermissionError:
                         print("\nCannot remove original WhatsApp directory. "
                             "Perhaps the directory is opened?", end="\n")
-    else:
+    elif args.exported:
         extract_exported.messages(args.exported, data, args.assume_first_as_me)
         if not args.no_html:
             extract.create_html(
@@ -320,8 +333,18 @@ def main():
             )
         for file in glob.glob(r'*.*'):
             shutil.copy(file, args.output)
+    elif args.import_json:
+        import_from_json(args.json, data)
+        extract.create_html(
+            data,
+            args.output,
+            args.template,
+            args.embedded,
+            args.offline,
+            args.size
+        )
 
-    if args.json:
+    if args.json and not args.import_json:
         if isinstance(data[next(iter(data))], ChatStore):
             data = {jik: chat.to_json() for jik, chat in data.items()}
         with open(args.json, "w") as f:
diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py
index cee431e..ff57014 100644
--- a/Whatsapp_Chat_Exporter/utility.py
+++ b/Whatsapp_Chat_Exporter/utility.py
@@ -1,7 +1,8 @@
+import json
 from bleach import clean as sanitize
 from markupsafe import Markup
 from datetime import datetime
-from enum import Enum
+from enum import IntEnum, StrEnum
 
 
 MAX_SIZE = 4 * 1024 * 1024  # Default 4MB
@@ -81,12 +82,45 @@ def rendering(
         )
 
 
-class Device(Enum):
+class Device(StrEnum):
     IOS = "ios"
     ANDROID = "android"
     EXPORTED = "exported"
 
 
+def import_from_json(json_file, data):
+    from Whatsapp_Chat_Exporter.data_model import ChatStore, Message
+    with open(json_file, "r") as f:
+        temp_data = json.loads(f.read())
+    total_row_number = len(tuple(temp_data.keys()))
+    print(f"Importing chats from JSON...(0/{total_row_number})", end="\r")
+    for index, (jid, chat_data) in enumerate(temp_data.items()):
+        chat = ChatStore(chat_data["type"], chat_data["name"])
+        chat.my_avatar = chat_data["my_avatar"]
+        chat.their_avatar = chat_data["their_avatar"]
+        chat.their_avatar_thumb = chat_data["their_avatar_thumb"]
+        for id, msg in chat_data["messages"].items():
+            message = Message(
+                msg["from_me"],
+                msg["timestamp"],
+                msg["time"],
+                msg["key_id"],
+            )
+            message.media = msg["media"]
+            message.meta = msg["meta"]
+            message.data = msg["data"]
+            message.sender = msg["sender"]
+            message.safe = msg["safe"]
+            message.reply = msg["reply"]
+            message.quoted_data = msg["quoted_data"]
+            message.caption = msg["caption"]
+            message.thumb = msg["thumb"]
+            message.sticker = msg["sticker"]
+            chat.add_message(id, message)
+        data[jid] = chat
+        print(f"Importing chats from JSON...({index + 1}/{total_row_number})", end="\r")
+
+
 # Android Specific
 CRYPT14_OFFSETS = (
     {"iv": 67, "db": 191},
@@ -97,7 +131,7 @@ CRYPT14_OFFSETS = (
 )
 
 
-class Crypt(Enum):
+class Crypt(IntEnum):
     CRYPT15 = 15
     CRYPT14 = 14
     CRYPT12 = 12