diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py index 0f78b40..e20fff4 100644 --- a/Whatsapp_Chat_Exporter/__main__.py +++ b/Whatsapp_Chat_Exporter/__main__.py @@ -527,7 +527,7 @@ def handle_decrypt_error(error: int) -> None: exit(5) -def process_contacts(args, data: ChatCollection, contact_store=None) -> None: +def process_contacts(args, data: ChatCollection) -> None: """Process contacts from the database.""" contact_db = args.wa if args.wa else "wa.db" if args.android else "ContactsV2.sqlite" @@ -616,14 +616,10 @@ def handle_media_directory(args) -> None: logger.info(f"Media directory has been copied to the output directory{CLEAR_LINE}") -def create_output_files(args, data: ChatCollection, contact_store=None) -> None: +def create_output_files(args, data: ChatCollection) -> None: """Create output files in the specified formats.""" # Create HTML files if requested if not args.no_html: - # Enrich from vcards if available - if contact_store and not contact_store.is_empty(): - contact_store.enrich_from_vcards(data) - android_handler.create_html( data, args.output, @@ -643,15 +639,12 @@ def create_output_files(args, data: ChatCollection, contact_store=None) -> None: # Create JSON files if requested if args.json and not args.import_json: - export_json(args, data, contact_store) + export_json(args, data) -def export_json(args, data: ChatCollection, contact_store=None) -> None: +def export_json(args, data: ChatCollection) -> None: """Export data to JSON format.""" - # Enrich from vcards if available - if contact_store and not contact_store.is_empty(): - contact_store.enrich_from_vcards(data) - + # TODO: remove all non-target chats from data if filtering is applied? # Convert ChatStore objects to JSON if isinstance(data.get(next(iter(data), None)), ChatStore): data = {jik: chat.to_json() for jik, chat in data.items()} @@ -848,13 +841,17 @@ def main(): logger.info(f"Incremental merge completed successfully.{CLEAR_LINE}") else: # Process contacts - process_contacts(args, data, contact_store) + process_contacts(args, data) + + # Enrich contacts from vCards if needed + if args.android and contact_store and not contact_store.is_empty(): + contact_store.enrich_from_vcards(data) # Process messages, media, and calls process_messages(args, data) # Create output files - create_output_files(args, data, contact_store) + create_output_files(args, data) # Handle media directory handle_media_directory(args) diff --git a/Whatsapp_Chat_Exporter/android_handler.py b/Whatsapp_Chat_Exporter/android_handler.py index d8e2de0..274661b 100644 --- a/Whatsapp_Chat_Exporter/android_handler.py +++ b/Whatsapp_Chat_Exporter/android_handler.py @@ -304,12 +304,10 @@ def _process_single_message(data, content, table_message, timezone_offset): return # Get or create the chat - if not data.get_chat(content["key_remote_jid"]): + current_chat = data.get_chat(content["key_remote_jid"]) + if current_chat is None: current_chat = data.add_chat(content["key_remote_jid"], ChatStore( Device.ANDROID, content["chat_subject"])) - else: - current_chat = data.get_chat(content["key_remote_jid"]) - # Determine sender_jid_row_id if "sender_jid_row_id" in content: sender_jid_row_id = content["sender_jid_row_id"] diff --git a/Whatsapp_Chat_Exporter/vcards_contacts.py b/Whatsapp_Chat_Exporter/vcards_contacts.py index 18cd500..315aa80 100644 --- a/Whatsapp_Chat_Exporter/vcards_contacts.py +++ b/Whatsapp_Chat_Exporter/vcards_contacts.py @@ -1,5 +1,7 @@ import vobject from typing import List, TypedDict +from Whatsapp_Chat_Exporter.data_model import ChatStore +from Whatsapp_Chat_Exporter.utility import Device class ExportedContactNumbers(TypedDict): @@ -21,11 +23,14 @@ class ContactsFromVCards: for number, name in self.contact_mapping: # short number must be a bad contact, lets skip it if len(number) <= 5: - continue - - for chat in filter_chats_by_prefix(chats, number).values(): - if not hasattr(chat, 'name') or (hasattr(chat, 'name') and (chat.name is None or chat.name == '')): - setattr(chat, 'name', name) + continue + chats_search = filter_chats_by_prefix(chats, number).values() + if chats_search: + for chat in chats_search: + if not hasattr(chat, 'name') or (hasattr(chat, 'name') and chat.name is None): + setattr(chat, 'name', name) + else: + chats.add_chat(number + "@s.whatsapp.net", ChatStore(Device.ANDROID, name)) def read_vcards_file(vcf_file_path, default_country_code: str):