From e0c2cf5f661cc338de2e50627263eea0c51a5226 Mon Sep 17 00:00:00 2001 From: KnugiHK <24708955+KnugiHK@users.noreply.github.com> Date: Tue, 13 Jun 2023 19:44:16 +0800 Subject: [PATCH] Implement iOS avatar #48 --- Whatsapp_Chat_Exporter/__main__.py | 22 ++++++--- Whatsapp_Chat_Exporter/data_model.py | 15 +++++- Whatsapp_Chat_Exporter/extract_iphone.py | 62 +++++++++++++++++++----- Whatsapp_Chat_Exporter/utility.py | 26 ++++++++-- Whatsapp_Chat_Exporter/whatsapp.html | 38 ++++++++++++--- 5 files changed, 135 insertions(+), 28 deletions(-) diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py index dafcfa1..96f9523 100644 --- a/Whatsapp_Chat_Exporter/__main__.py +++ b/Whatsapp_Chat_Exporter/__main__.py @@ -161,6 +161,13 @@ def main(): action='store_true', help="Assume the first message in a chat as sent by me (must be used together with -e)" ) + parser.add_argument( + "--no-avatar", + dest="no_avatar", + default=False, + action='store_true', + help="Do not render avatar in HTML output" + ) args = parser.parse_args() # Check for updates @@ -261,7 +268,7 @@ def main(): if os.path.isfile(msg_db): with sqlite3.connect(msg_db) as db: db.row_factory = sqlite3.Row - messages(db, data) + messages(db, data, args.media) media(db, data, args.media) vcard(db, data) if not args.no_html: @@ -271,7 +278,8 @@ def main(): args.template, args.embedded, args.offline, - args.size + args.size, + args.no_avatar ) else: print( @@ -283,20 +291,20 @@ def main(): if os.path.isdir(args.media): media_path = os.path.join(args.output, args.media) if os.path.isdir(media_path): - print("Media directory already exists in output directory. Skipping...") + print("\nMedia directory already exists in output directory. Skipping...", end="\n") else: if not args.move_media: if os.path.isdir(media_path): - print("WhatsApp directory already exists in output directory. Skipping...") + print("\nWhatsApp directory already exists in output directory. Skipping...", end="\n") else: - print("Copying media directory...") + print("\nCopying media directory...", end="\n") shutil.copytree(args.media, media_path) else: try: shutil.move(args.media, f"{args.output}/") except PermissionError: - print("Cannot remove original WhatsApp directory. " - "Perhaps the directory is opened?") + print("\nCannot remove original WhatsApp directory. " + "Perhaps the directory is opened?", end="\n") else: extract_exported.messages(args.exported, data, args.assume_first_as_me) if not args.no_html: diff --git a/Whatsapp_Chat_Exporter/data_model.py b/Whatsapp_Chat_Exporter/data_model.py index 28f957d..51463f1 100644 --- a/Whatsapp_Chat_Exporter/data_model.py +++ b/Whatsapp_Chat_Exporter/data_model.py @@ -1,13 +1,26 @@ +import os from datetime import datetime from typing import Union +from Whatsapp_Chat_Exporter.utility import Device class ChatStore(): - def __init__(self, name=None): + def __init__(self, type, name=None, media=None): if name is not None and not isinstance(name, str): raise TypeError("Name must be a string or None") self.name = name self.messages = {} + if media is not None: + if type == Device.IOS: + self.my_avatar = os.path.join(media, "Media/Profile/Photo.jpg") + elif type == Device.ANDROID: + self.my_avatar = None # TODO: Add Android support + else: + self.my_avatar = None + else: + self.my_avatar = None + self.their_avatar = None + self.their_avatar_thumb = None def add_message(self, id, message): if not isinstance(message, Message): diff --git a/Whatsapp_Chat_Exporter/extract_iphone.py b/Whatsapp_Chat_Exporter/extract_iphone.py index ac085ca..8a0f739 100644 --- a/Whatsapp_Chat_Exporter/extract_iphone.py +++ b/Whatsapp_Chat_Exporter/extract_iphone.py @@ -1,5 +1,6 @@ #!/usr/bin/python3 +from glob import glob import sqlite3 import json import jinja2 @@ -8,10 +9,10 @@ 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 +from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, rendering, sanitize_except, determine_day, APPLE_TIME, Device -def messages(db, data): +def messages(db, data, media_folder): c = db.cursor() # Get contacts c.execute("""SELECT count() FROM ZWACHATSESSION""") @@ -21,7 +22,17 @@ def messages(db, data): c.execute("""SELECT ZCONTACTJID, ZPARTNERNAME FROM ZWACHATSESSION; """) content = c.fetchone() while content is not None: - data[content["ZCONTACTJID"]] = ChatStore(content["ZPARTNERNAME"]) + data[content["ZCONTACTJID"]] = ChatStore(Device.IOS, content["ZPARTNERNAME"], media_folder) + path = f'{media_folder}/Media/Profile/{content["ZCONTACTJID"].split("@")[0]}' + avatars = glob(f"{path}*") + if 0 < len(avatars) <= 1: + data[content["ZCONTACTJID"]].their_avatar = avatars[0] + else: + for avatar in avatars: + if avatar.endswith(".thumb"): + data[content["ZCONTACTJID"]].their_avatar_thumb = avatar + elif avatar.endswith(".jpg"): + data[content["ZCONTACTJID"]].their_avatar = avatar content = c.fetchone() # Get message history @@ -49,7 +60,17 @@ def messages(db, data): _id = content["_id"] Z_PK = content["Z_PK"] if _id not in data: - data[_id] = ChatStore() + data[_id] = ChatStore(Device.IOS) + path = f'{media_folder}/Media/Profile/{_id.split("@")[0]}' + avatars = glob(f"{path}*") + if 0 < len(avatars) <= 1: + data[_id].their_avatar = avatars[0] + else: + for avatar in avatars: + if avatar.endswith(".thumb"): + data[_id].their_avatar_thumb = avatar + elif avatar.endswith(".jpg"): + data[_id].their_avatar = avatar ts = APPLE_TIME + content["ZMESSAGEDATE"] message = Message( from_me=content["ZISFROMME"], @@ -232,7 +253,8 @@ def create_html( template=None, embedded=False, offline_static=False, - maximum_size=None + maximum_size=None, + no_avatar=False ): if template is None: template_dir = os.path.dirname(__file__) @@ -243,11 +265,12 @@ def create_html( templateLoader = jinja2.FileSystemLoader(searchpath=template_dir) templateEnv = jinja2.Environment(loader=templateLoader) templateEnv.globals.update(determine_day=determine_day) + templateEnv.globals.update(no_avatar=no_avatar) templateEnv.filters['sanitize_except'] = sanitize_except template = templateEnv.get_template(template_file) total_row_number = len(data) - print(f"\nCreating HTML...(0/{total_row_number})", end="\r") + print(f"\nGenerating chats...(0/{total_row_number})", end="\r") if not os.path.isdir(output_folder): os.mkdir(output_folder) @@ -305,7 +328,10 @@ def create_html( render_box, contact, w3css, - f"{safe_file_name}-{current_page + 1}.html" + f"{safe_file_name}-{current_page + 1}.html", + chat.my_avatar, + chat.their_avatar, + chat.their_avatar_thumb ) render_box = [message] current_size = 0 @@ -323,17 +349,31 @@ def create_html( render_box, contact, w3css, - False + False, + chat.my_avatar, + chat.their_avatar, + chat.their_avatar_thumb ) 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) + rendering( + output_file_name, + template, + name, + chat.get_messages(), + contact, + w3css, + False, + chat.my_avatar, + chat.their_avatar, + chat.their_avatar_thumb + ) if current % 10 == 0: - print(f"Creating HTML...({current}/{total_row_number})", end="\r") + print(f"Generating chats...({current}/{total_row_number})", end="\r") - print(f"Creating HTML...({total_row_number}/{total_row_number})", end="\r") + print(f"Generating chats...({total_row_number}/{total_row_number})", end="\r") if __name__ == "__main__": diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py index c2860cd..cee431e 100644 --- a/Whatsapp_Chat_Exporter/utility.py +++ b/Whatsapp_Chat_Exporter/utility.py @@ -53,20 +53,40 @@ def check_update(): return 0 -def rendering(output_file_name, template, name, msgs, contact, w3css, next): +def rendering( + output_file_name, + template, + name, + msgs, + contact, + w3css, + next, + my_avatar, + their_avatar, + their_avatar_thumb + ): + if their_avatar_thumb is None and their_avatar is not None: + their_avatar_thumb = their_avatar with open(output_file_name, "w", encoding="utf-8") as f: f.write( template.render( name=name, msgs=msgs, - my_avatar=None, - their_avatar=f"WhatsApp/Avatars/{contact}.j", + my_avatar=my_avatar, + their_avatar=their_avatar, + their_avatar_thumb=their_avatar_thumb, w3css=w3css, next=next ) ) +class Device(Enum): + IOS = "ios" + ANDROID = "android" + EXPORTED = "exported" + + # Android Specific CRYPT14_OFFSETS = ( {"iv": 67, "db": 191}, diff --git a/Whatsapp_Chat_Exporter/whatsapp.html b/Whatsapp_Chat_Exporter/whatsapp.html index 8862f0e..bef7d8c 100644 --- a/Whatsapp_Chat_Exporter/whatsapp.html +++ b/Whatsapp_Chat_Exporter/whatsapp.html @@ -58,6 +58,12 @@ border-color: rgba(0,0,0,0); } } + .avatar { + border-radius:50%; + overflow:hidden; + max-width: 64px; + max-height: 64px; + } @@ -77,7 +83,11 @@
You
+ {% if not no_avatar and my_avatar is not none %}
+ {% else %} +
+ {% endif %}
{% if msg.reply is not none %}
@@ -92,7 +102,7 @@
{% endif %} {% if msg.meta == true or msg.media == false and msg.data is none %} -
+

{{ msg.data or 'Not supported WhatsApp internal message' }}

{% else %} @@ -110,7 +120,7 @@ {% elif "/" in msg.mime %} -
+

The file cannot be displayed here, however it should be located at here

{% else %} @@ -124,7 +134,13 @@ {% endif %}
-
+ {% if not no_avatar and my_avatar is not none %} +
+ + + +
+ {% endif %}
{% else %}
@@ -138,8 +154,18 @@
{{ msg.time }}
-
+ {% if not no_avatar %} +
+ {% if their_avatar is not none %} + + {% else %} + + {% endif %} +
+ {% else %} +
+ {% endif %}
{% if msg.reply is not none %}
@@ -154,7 +180,7 @@
{% endif %} {% if msg.meta == true or msg.media == false and msg.data is none %} -
+

{{ msg.data or 'Not supported WhatsApp internal message' }}

{% else %} @@ -172,7 +198,7 @@ {% elif "/" in msg.mime %} -
+

The file cannot be displayed here, however it should be located at here

{% else %}