From 0d626519eca68f8b4124a1d6d5746c786b1e7e95 Mon Sep 17 00:00:00 2001
From: KnugiHK <24708955+KnugiHK@users.noreply.github.com>
Date: Sun, 3 Dec 2023 13:40:21 +0800
Subject: [PATCH] Add support for decrypting contact db from crypt15 backup

---
 Whatsapp_Chat_Exporter/__main__.py | 31 ++++++++++++++++++++++++------
 Whatsapp_Chat_Exporter/extract.py  | 16 +++++++--------
 Whatsapp_Chat_Exporter/utility.py  |  5 +++++
 3 files changed, 38 insertions(+), 14 deletions(-)

diff --git a/Whatsapp_Chat_Exporter/__main__.py b/Whatsapp_Chat_Exporter/__main__.py
index 13985a8..aeac9d5 100644
--- a/Whatsapp_Chat_Exporter/__main__.py
+++ b/Whatsapp_Chat_Exporter/__main__.py
@@ -9,7 +9,7 @@ import glob
 from Whatsapp_Chat_Exporter import extract_exported, extract_iphone
 from Whatsapp_Chat_Exporter import extract, extract_iphone_media
 from Whatsapp_Chat_Exporter.data_model import ChatStore
-from Whatsapp_Chat_Exporter.utility import Crypt, check_update, import_from_json
+from Whatsapp_Chat_Exporter.utility import Crypt, DbType, check_update, import_from_json
 from argparse import ArgumentParser, SUPPRESS
 from sys import exit
 try:
@@ -191,6 +191,13 @@ def main():
         action='store_true',
         help="Preserve the modification timestamp of the extracted files (iOS only)"
     )
+    parser.add_argument(
+        "--wab",
+        "--wa-backup",
+        dest="wab",
+        default=None,
+        help="Path to contact database in crypt15 format"
+    )
     args = parser.parse_args()
 
     # Check for updates
@@ -229,6 +236,10 @@ def main():
             msg_db = "msgstore.db"
         else:
             msg_db = args.db
+        if args.wa is None:
+            contact_db = "wa.db"
+        else:
+            contact_db = args.wa
         if args.key is not None:
             if args.backup is None:
                 print("You must specify the backup file with -b")
@@ -245,7 +256,19 @@ def main():
             elif all(char in string.hexdigits for char in args.key):
                 key = bytes.fromhex(args.key)
             db = open(args.backup, "rb").read()
-            error = extract.decrypt_backup(db, key, msg_db, crypt, args.showkey)
+            if args.wab:
+                wab = open(args.wab, "rb").read()
+                error_wa = extract.decrypt_backup(wab, key, contact_db, crypt, args.showkey, DbType.CONTACT)
+                key.seek(0)
+            else:
+                error_wa = 0
+            error_message = extract.decrypt_backup(db, key, msg_db, crypt, args.showkey, DbType.MESSAGE)               
+            if error_wa != 0:
+                error = error_wa
+            elif error_message != 0:
+                error = error_message
+            else:
+                error = 0
             if error != 0:
                 if error == 1:
                     print("Dependencies of decrypt_backup and/or extract_encrypted_key"
@@ -258,10 +281,6 @@ def main():
                 else:
                     print("Unknown error occurred.", error)
                     exit(5)
-        if args.wa is None:
-            contact_db = "wa.db"
-        else:
-            contact_db = args.wa
         if args.media is None:
             args.media = "WhatsApp"
 
diff --git a/Whatsapp_Chat_Exporter/extract.py b/Whatsapp_Chat_Exporter/extract.py
index e0ada4e..5cbe5dd 100644
--- a/Whatsapp_Chat_Exporter/extract.py
+++ b/Whatsapp_Chat_Exporter/extract.py
@@ -1,11 +1,7 @@
 #!/usr/bin/python3
 
 import sqlite3
-import json
-import jinja2
 import os
-import shutil
-import re
 import io
 import hmac
 from pathlib import Path
@@ -13,7 +9,7 @@ from mimetypes import MimeTypes
 from hashlib import sha256
 from base64 import b64decode, b64encode
 from Whatsapp_Chat_Exporter.data_model import ChatStore, Message
-from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, determine_metadata, get_status_location
+from Whatsapp_Chat_Exporter.utility import MAX_SIZE, ROW_SIZE, DbType, determine_metadata, get_status_location
 from Whatsapp_Chat_Exporter.utility import rendering, Crypt, Device, get_file_name, setup_template
 from Whatsapp_Chat_Exporter.utility import brute_force_offset, CRYPT14_OFFSETS, JidType
 
@@ -53,7 +49,7 @@ def _extract_encrypted_key(keyfile):
     return _generate_hmac_of_hmac(key_stream)
 
 
-def decrypt_backup(database, key, output, crypt=Crypt.CRYPT14, show_crypt15=False):
+def decrypt_backup(database, key, output, crypt=Crypt.CRYPT14, show_crypt15=False, db_type=DbType.MESSAGE):
     if not support_backup:
         return 1
     if isinstance(key, io.IOBase):
@@ -83,8 +79,12 @@ def decrypt_backup(database, key, output, crypt=Crypt.CRYPT14, show_crypt15=Fals
         if len(database) < 131:
             raise ValueError("The crypt15 file must be at least 131 bytes")
         t1 = t2 = None
-        iv = database[8:24]
-        db_offset = database[0] + 2  # Skip protobuf + protobuf size and backup type
+        if db_type == DbType.MESSAGE:
+            iv = database[8:24]
+            db_offset = database[0] + 2  # Skip protobuf + protobuf size and backup type
+        elif db_type == DbType.CONTACT:
+            iv = database[7:23]
+            db_offset = database[0] + 1  # Skip protobuf + protobuf size
         db_ciphertext = database[db_offset:]
 
     if t1 != t2:
diff --git a/Whatsapp_Chat_Exporter/utility.py b/Whatsapp_Chat_Exporter/utility.py
index 6490081..9e6cf94 100644
--- a/Whatsapp_Chat_Exporter/utility.py
+++ b/Whatsapp_Chat_Exporter/utility.py
@@ -172,6 +172,11 @@ class Crypt(IntEnum):
     CRYPT12 = 12
 
 
+class DbType(StrEnum):
+    MESSAGE = "message"
+    CONTACT = "contact"
+
+
 def brute_force_offset(max_iv=200, max_db=200):
     for iv in range(0, max_iv):
         for db in range(0, max_db):