Refactor

Whatsapp_Chat_Exporter/__main__.py

@@ -261,7 +261,7 @@ def main():
         messages = extract_iphone.messages
         media = extract_iphone.media
         vcard = extract_iphone.vcard
-        create_html = extract_iphone.create_html
+        create_html = extract.create_html
         if args.media is None:
             args.media = "AppDomainGroup-group.net.whatsapp.WhatsApp.shared"
         if args.backup is not None:

Whatsapp_Chat_Exporter/extract.py

@@ -14,7 +14,7 @@ from hashlib import sha256
 from base64 import b64decode, b64encode
 from Whatsapp_Chat_Exporter.data_model import ChatStore, Message
 from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, Device, determine_metadata
-from Whatsapp_Chat_Exporter.utility import rendering, sanitize_except, determine_day, Crypt
+from Whatsapp_Chat_Exporter.utility import rendering, sanitize_except, determine_day, Crypt, get_file_name
 from Whatsapp_Chat_Exporter.utility import brute_force_offset, CRYPT14_OFFSETS
 
 try:
@@ -669,29 +669,14 @@ def create_html(
         w3css_path = os.path.join(static_folder, "w3.css")
         if not os.path.isfile(w3css_path):
             with urllib.request.urlopen(w3css) as resp:
-                with open(w3css_path, "wb") as f:
+                with open(w3css_path, "wb") as f: f.write(resp.read())
-                    f.write(resp.read())
         w3css = os.path.join(offline_static, "w3.css")
 
     for current, contact in enumerate(data):
         chat = data[contact]
         if len(chat.messages) == 0:
             continue
-        phone_number = contact.split('@')[0]
+        safe_file_name, name = get_file_name(contact, chat)
-        if "-" in contact:
-            file_name = ""
-        else:
-            file_name = phone_number
-
-        if chat.name is not None:
-            if file_name != "":
-                file_name += "-"
-            file_name += chat.name.replace("/", "-")
-            name = chat.name
-        else:
-            name = phone_number
-
-        safe_file_name = "".join(x for x in file_name if x.isalnum() or x in "- ")
 
         if maximum_size is not None:
             current_size = 0

Whatsapp_Chat_Exporter/extract_iphone.py

@@ -9,7 +9,7 @@ import shutil
 from pathlib import Path
 from mimetypes import MimeTypes
 from Whatsapp_Chat_Exporter.data_model import ChatStore, Message
-from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, rendering, sanitize_except, determine_day, APPLE_TIME, Device
+from Whatsapp_Chat_Exporter.utility import APPLE_TIME, Device
 
 
 def contacts(db, data):
@@ -282,128 +282,3 @@ def vcard(db, data):
         message.media = True
         message.meta = True
         print(f"Processing vCards...({index + 1}/{total_row_number})", end="\r")
-
-
-def create_html(
-        data,
-        output_folder,
-        template=None,
-        embedded=False,
-        offline_static=False,
-        maximum_size=None,
-        no_avatar=False
-    ):
-    if template is None:
-        template_dir = os.path.dirname(__file__)
-        template_file = "whatsapp.html"
-    else:
-        template_dir = os.path.dirname(template)
-        template_file = os.path.basename(template)
-    template_loader = jinja2.FileSystemLoader(searchpath=template_dir)
-    template_env = jinja2.Environment(loader=template_loader, autoescape=True)
-    template_env.globals.update(
-        determine_day=determine_day,
-        no_avatar=no_avatar
-    )
-    template_env.filters['sanitize_except'] = sanitize_except
-    template = template_env.get_template(template_file)
-
-    total_row_number = len(data)
-    print(f"\nGenerating chats...(0/{total_row_number})", end="\r")
-
-    if not os.path.isdir(output_folder):
-        os.mkdir(output_folder)
-
-    w3css = "https://www.w3schools.com/w3css/4/w3.css"
-    if offline_static:
-        import urllib.request
-        static_folder = os.path.join(output_folder, offline_static)
-        if not os.path.isdir(static_folder):
-            os.mkdir(static_folder)
-        w3css_path = os.path.join(static_folder, "w3.css")
-        if not os.path.isfile(w3css_path):
-            with urllib.request.urlopen(w3css) as resp:
-                with open(w3css_path, "wb") as f: f.write(resp.read())
-        w3css = os.path.join(offline_static, "w3.css")
-
-    for current, contact in enumerate(data):
-        chat = data[contact]
-        if len(chat.messages) == 0:
-            continue
-        phone_number = contact.split('@')[0]
-        if "-" in contact:
-            file_name = ""
-        else:
-            file_name = phone_number
-
-        if chat.name is not None:
-            if file_name != "":
-                file_name += "-"
-            file_name += chat.name.replace("/", "-")
-            name = chat.name
-        else:
-            name = phone_number
-
-        safe_file_name = "".join(x for x in file_name if x.isalnum() or x in "- ")
-
-        if maximum_size is not None:
-            current_size = 0
-            current_page = 1
-            render_box = []
-            if maximum_size == 0:
-                maximum_size = MAX_SIZE
-            last_msg = chat.get_last_message().key_id
-            for message in chat.get_messages():
-                if message.data is not None and not message.meta and not message.media:
-                    current_size += len(message.data) + ROW_SIZE
-                else:
-                    current_size += ROW_SIZE + 100 # Assume media and meta HTML are 100 bytes
-                if current_size > maximum_size:
-                    output_file_name = f"{output_folder}/{safe_file_name}-{current_page}.html"
-                    rendering(
-                        output_file_name,
-                        template,
-                        name,
-                        render_box,
-                        contact,
-                        w3css,
-                        f"{safe_file_name}-{current_page + 1}.html",
-                        chat
-                    )
-                    render_box = [message]
-                    current_size = 0
-                    current_page += 1
-                else:
-                    if message.key_id == last_msg:
-                        if current_page == 1:
-                            output_file_name = f"{output_folder}/{safe_file_name}.html"
-                        else:
-                            output_file_name = f"{output_folder}/{safe_file_name}-{current_page}.html"
-                        rendering(
-                            output_file_name,
-                            template,
-                            name,
-                            render_box,
-                            contact,
-                            w3css,
-                            False,
-                            chat
-                        )
-                    else:
-                        render_box.append(message)
-        else:
-            output_file_name = f"{output_folder}/{safe_file_name}.html"
-            rendering(
-                output_file_name,
-                template,
-                name,
-                chat.get_messages(),
-                contact,
-                w3css,
-                False,
-                chat
-            )
-        if current % 10 == 0:
-            print(f"Generating chats...({current}/{total_row_number})", end="\r")
-
-    print(f"Generating chats...({total_row_number}/{total_row_number})", end="\r")

Whatsapp_Chat_Exporter/utility.py

@@ -3,6 +3,7 @@ from bleach import clean as sanitize
 from markupsafe import Markup
 from datetime import datetime
 from enum import IntEnum
+from Whatsapp_Chat_Exporter.data_model import ChatStore
 try:
     from enum import StrEnum
 except ImportError:
@@ -130,6 +131,26 @@ def import_from_json(json_file, data):
         print(f"Importing chats from JSON...({index + 1}/{total_row_number})", end="\r")
 
 
+def get_file_name(contact: str, chat: ChatStore):
+    if "@" not in contact:
+        raise ValueError("Unexpected contact format: " + contact)
+    phone_number = contact.split('@')[0]
+    if "-" in contact:
+        file_name = ""
+    else:
+        file_name = phone_number
+
+    if chat.name is not None:
+        if file_name != "":
+            file_name += "-"
+        file_name += chat.name.replace("/", "-")
+        name = chat.name
+    else:
+        name = phone_number
+
+    return "".join(x for x in file_name if x.isalnum() or x in "- "), name
+
+
 # Android Specific
 CRYPT14_OFFSETS = (
     {"iv": 67, "db": 191},