Implement iOS avatar #48

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. "
+                        print("\nCannot remove original WhatsApp directory. "
-                            "Perhaps the directory is opened?")
+                            "Perhaps the directory is opened?", end="\n")
     else:
         extract_exported.messages(args.exported, data, args.assume_first_as_me)
         if not args.no_html:

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):

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__":

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,
+                my_avatar=my_avatar,
-                their_avatar=f"WhatsApp/Avatars/{contact}.j",
+                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},

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;
+		}
 		</style>
 	</head>
 	<body>
@@ -77,7 +83,11 @@
 						<div style="padding-left: 10px; text-align: right; color: #3892da;">You</div>
 					</div>
 					<div class="w3-row">
+						{% if not no_avatar and my_avatar is not none %}
 						<div class="w3-col m10 l10">
+						{% else %}
+						<div class="w3-col m12 l12">
+						{% endif %}
 							<div style="text-align: right;">
 							{% if msg.reply is not none %}
 								<div class="reply">
@@ -92,7 +102,7 @@
 								</div>
 							{% endif %}
 							{% if msg.meta == true or msg.media == false and msg.data is none %}
-								<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar">
+								<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar w3-threequarter w3-center">
 								 	<p>{{ msg.data or 'Not supported WhatsApp internal message' }}</p>
 							    </div>
 							{% else %}
@@ -110,7 +120,7 @@
 										<source src="{{ msg.data }}" />
 									</video>
 									{% elif "/" in msg.mime %}
-										<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar">
+										<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar w3-threequarter w3-center">
 											<p>The file cannot be displayed here, however it should be located at <a href="./{{ msg.data }}">here</a></p>
 										</div>
 									{% else %}
@@ -124,7 +134,13 @@
 								{% endif %}
 							</div>
 						</div>
-						<div class="w3-col m2 l2" style="padding-left: 10px"><img src="{{ my_avatar }}" onerror="this.style.display='none'"></div>
+						{% if not no_avatar and my_avatar is not none %}
+						<div class="w3-col m2 l2" style="padding-left: 10px">
+							<a href="{{ my_avatar }}">
+								<img src="{{ my_avatar }}" onerror="this.style.display='none'" class="avatar">
+							</a>
+						</div>
+						{% endif %}
 					</div>
 					{% else %}
 					<div class="w3-row">
@@ -138,8 +154,18 @@
 						<div style="text-align: right; color:#70777c;">{{ msg.time }}</div>
 					</div>
 					<div class="w3-row">
-						<div class="w3-col m2 l2"><img src="{{ their_avatar }}" onerror="this.style.display='none'"></div>
+						{% if not no_avatar %}
+						<div class="w3-col m2 l2">
+								{% if their_avatar is not none %}
+								<a href="{{ their_avatar }}"><img src="{{ their_avatar_thumb or '' }}" onerror="this.style.display='none'" class="avatar"></a>
+								{% else %}
+								<img src="{{ their_avatar_thumb or '' }}" onerror="this.style.display='none'" class="avatar">
+								{% endif %}
+						</div>
 						<div class="w3-col m10 l10">
+						{% else %}
+						<div class="w3-col m12 l12">
+						{% endif %}
 							<div style="text-align: left;">
 							{% if msg.reply is not none %}
 								<div class="reply">
@@ -154,7 +180,7 @@
 								</div>
 							{% endif %}
 							{% if msg.meta == true or msg.media == false and msg.data is none %}
-								<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar">
+								<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar w3-threequarter w3-center">
 									<p>{{ msg.data or 'Not supported WhatsApp internal message' }}</p>
 								</div>
 							{% else %}
@@ -172,7 +198,7 @@
 										<source src="{{ msg.data }}" />
 									</video>
 									{% elif "/" in msg.mime %}
-										<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar">
+										<div style="text-align: center;" class="w3-panel w3-border-blue w3-pale-blue w3-rightbar w3-leftbar w3-threequarter w3-center">
 											<p>The file cannot be displayed here, however it should be located at <a href="./{{ msg.data }}">here</a></p>
 										</div>
 									{% else %}