Implement import json and output HTML

This commit is contained in:
KnugiHK
2023-06-16 19:19:33 +08:00
parent fff11b26a5
commit 506f8e89f4
2 changed files with 66 additions and 9 deletions
+29 -6
View File
@@ -6,7 +6,7 @@ import glob
from Whatsapp_Chat_Exporter import extract, extract_exported, extract_iphone from Whatsapp_Chat_Exporter import extract, extract_exported, extract_iphone
from Whatsapp_Chat_Exporter import extract_iphone_media from Whatsapp_Chat_Exporter import extract_iphone_media
from Whatsapp_Chat_Exporter.data_model import ChatStore 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 from argparse import ArgumentParser, SUPPRESS
import os import os
import sqlite3 import sqlite3
@@ -168,6 +168,13 @@ def main():
action='store_true', action='store_true',
help="Do not render avatar in HTML output" 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() args = parser.parse_args()
# Check for updates # Check for updates
@@ -175,15 +182,21 @@ def main():
exit(check_update()) exit(check_update())
# Sanity checks # 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.") print("You must define only one device type.")
exit(1) 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.") print("You must define the device type.")
exit(1) exit(1)
if args.no_html and not args.json: if args.no_html and not args.json:
print("You must either specify a JSON output file or enable HTML output.") print("You must either specify a JSON output file or enable HTML output.")
exit(1) 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 = {} data = {}
@@ -264,7 +277,7 @@ def main():
else: else:
contact_db = args.wa contact_db = args.wa
if not args.exported: if not args.exported and not args.import_json:
if os.path.isfile(msg_db): if os.path.isfile(msg_db):
with sqlite3.connect(msg_db) as db: with sqlite3.connect(msg_db) as db:
db.row_factory = sqlite3.Row db.row_factory = sqlite3.Row
@@ -307,7 +320,7 @@ def main():
except PermissionError: except PermissionError:
print("\nCannot remove original WhatsApp directory. " print("\nCannot remove original WhatsApp directory. "
"Perhaps the directory is opened?", end="\n") "Perhaps the directory is opened?", end="\n")
else: elif args.exported:
extract_exported.messages(args.exported, data, args.assume_first_as_me) extract_exported.messages(args.exported, data, args.assume_first_as_me)
if not args.no_html: if not args.no_html:
extract.create_html( extract.create_html(
@@ -320,8 +333,18 @@ def main():
) )
for file in glob.glob(r'*.*'): for file in glob.glob(r'*.*'):
shutil.copy(file, args.output) 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): if isinstance(data[next(iter(data))], ChatStore):
data = {jik: chat.to_json() for jik, chat in data.items()} data = {jik: chat.to_json() for jik, chat in data.items()}
with open(args.json, "w") as f: with open(args.json, "w") as f:
+37 -3
View File
@@ -1,7 +1,8 @@
import json
from bleach import clean as sanitize from bleach import clean as sanitize
from markupsafe import Markup from markupsafe import Markup
from datetime import datetime from datetime import datetime
from enum import Enum from enum import IntEnum, StrEnum
MAX_SIZE = 4 * 1024 * 1024 # Default 4MB MAX_SIZE = 4 * 1024 * 1024 # Default 4MB
@@ -81,12 +82,45 @@ def rendering(
) )
class Device(Enum): class Device(StrEnum):
IOS = "ios" IOS = "ios"
ANDROID = "android" ANDROID = "android"
EXPORTED = "exported" 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 # Android Specific
CRYPT14_OFFSETS = ( CRYPT14_OFFSETS = (
{"iv": 67, "db": 191}, {"iv": 67, "db": 191},
@@ -97,7 +131,7 @@ CRYPT14_OFFSETS = (
) )
class Crypt(Enum): class Crypt(IntEnum):
CRYPT15 = 15 CRYPT15 = 15
CRYPT14 = 14 CRYPT14 = 14
CRYPT12 = 12 CRYPT12 = 12