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

View File

@@ -6,7 +6,7 @@ import glob
from Whatsapp_Chat_Exporter import extract, extract_exported, extract_iphone
from Whatsapp_Chat_Exporter import extract_iphone_media
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
import os
import sqlite3
@@ -168,6 +168,13 @@ def main():
action='store_true',
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()
# Check for updates
@@ -175,15 +182,21 @@ def main():
exit(check_update())
# 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.")
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.")
exit(1)
if args.no_html and not args.json:
print("You must either specify a JSON output file or enable HTML output.")
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 = {}
@@ -264,7 +277,7 @@ def main():
else:
contact_db = args.wa
if not args.exported:
if not args.exported and not args.import_json:
if os.path.isfile(msg_db):
with sqlite3.connect(msg_db) as db:
db.row_factory = sqlite3.Row
@@ -307,7 +320,7 @@ def main():
except PermissionError:
print("\nCannot remove original WhatsApp directory. "
"Perhaps the directory is opened?", end="\n")
else:
elif args.exported:
extract_exported.messages(args.exported, data, args.assume_first_as_me)
if not args.no_html:
extract.create_html(
@@ -320,8 +333,18 @@ def main():
)
for file in glob.glob(r'*.*'):
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):
data = {jik: chat.to_json() for jik, chat in data.items()}
with open(args.json, "w") as f:

View File

@@ -1,7 +1,8 @@
import json
from bleach import clean as sanitize
from markupsafe import Markup
from datetime import datetime
from enum import Enum
from enum import IntEnum, StrEnum
MAX_SIZE = 4 * 1024 * 1024 # Default 4MB
@@ -81,12 +82,45 @@ def rendering(
)
class Device(Enum):
class Device(StrEnum):
IOS = "ios"
ANDROID = "android"
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
CRYPT14_OFFSETS = (
{"iv": 67, "db": 191},
@@ -97,7 +131,7 @@ CRYPT14_OFFSETS = (
)
class Crypt(Enum):
class Crypt(IntEnum):
CRYPT15 = 15
CRYPT14 = 14
CRYPT12 = 12