Improve brute-force offsets with process pool

Refactored the brute-force offset search in `_decrypt_crypt14` to use `ProcessPoolExecutor` for better parallelism and performance. Improved progress reporting and clean shutdown on success or interruption.
This commit is contained in:
KnugiHK
2026-01-17 14:43:51 +08:00
parent c05e76569b
commit 68dcc6abe0

View File

@@ -6,6 +6,7 @@ import concurrent.futures
from tqdm import tqdm from tqdm import tqdm
from typing import Tuple, Union from typing import Tuple, Union
from hashlib import sha256 from hashlib import sha256
from functools import partial
from Whatsapp_Chat_Exporter.utility import CLEAR_LINE, CRYPT14_OFFSETS, Crypt, DbType from Whatsapp_Chat_Exporter.utility import CLEAR_LINE, CRYPT14_OFFSETS, Crypt, DbType
try: try:
@@ -175,38 +176,43 @@ def _decrypt_crypt14(database: bytes, main_key: bytes, max_worker: int = 10) ->
) )
return decrypted_db # Successful decryption return decrypted_db # Successful decryption
logger.info(f"Common offsets failed. Will attempt to brute-force{CLEAR_LINE}")
offset_max = 200 offset_max = 200
logger.info(f"Common offsets failed. Attempt to brute-force...{CLEAR_LINE}") workers = max_worker
with tqdm(total=offset_max ** 2, desc="Brute-forcing offsets", unit="trial", leave=False) as pbar: check_offset = partial(_attempt_decrypt_task, database=database, main_key=main_key)
with concurrent.futures.ThreadPoolExecutor(max_worker) as executor: all_offsets = list(brute_force_offset(offset_max, offset_max))
# Map futures to their offsets executor = concurrent.futures.ProcessPoolExecutor(max_workers=workers)
future_to_offset = { try:
executor.submit(_attempt_decrypt_task, offset, database, main_key): offset with tqdm(total=len(all_offsets), desc="Brute-forcing offsets", unit="trial", leave=False) as pbar:
for offset in brute_force_offset(offset_max, offset_max) results = executor.map(check_offset, all_offsets, chunksize=8)
} found = False
for offset_info, result in zip(all_offsets, results):
pbar.update(1)
if result:
start_iv, _, start_db = offset_info
# Clean shutdown on success
executor.shutdown(wait=False, cancel_futures=True)
found = True
break
if found:
logger.info(
f"The offsets of your IV and database are {start_iv} and {start_db}, respectively.{CLEAR_LINE}"
)
logger.info(
f"To include your offsets in the expoter, please report it in the discussion thread on GitHub:{CLEAR_LINE}"
)
logger.info(f"https://github.com/KnugiHK/Whatsapp-Chat-Exporter/discussions/47{CLEAR_LINE}")
return result
try: except KeyboardInterrupt:
for future in concurrent.futures.as_completed(future_to_offset): executor.shutdown(wait=False, cancel_futures=True)
pbar.update(1) print("\n")
result = future.result() raise KeyboardInterrupt(
f"Brute force interrupted by user (Ctrl+C). Shutting down gracefully...{CLEAR_LINE}"
if result is not None: )
# Success! Shutdown other tasks immediately
executor.shutdown(wait=False, cancel_futures=True) finally:
executor.shutdown(wait=False)
start_iv, _, start_db = future_to_offset[future]
logger.info(
f"The offsets of your IV and database are {start_iv} and "
f"{start_db}, respectively. To include your offsets in the "
"program, please report it by creating an issue on GitHub: "
"https://github.com/KnugiHK/Whatsapp-Chat-Exporter/discussions/47"
f"\nShutting down other threads...{CLEAR_LINE}"
)
return result
except KeyboardInterrupt:
executor.shutdown(wait=False, cancel_futures=True)
raise KeyboardInterrupt("Brute force interrupted by user (Ctrl+C). Shutting down gracefully...{CLEAR_LINE}")
raise OffsetNotFoundError("Could not find the correct offsets for decryption.") raise OffsetNotFoundError("Could not find the correct offsets for decryption.")