mirror of
https://github.com/KnugiHK/WhatsApp-Chat-Exporter.git
synced 2026-02-10 19:22:24 +00:00
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:
@@ -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.")
|
||||||
|
|
||||||
|
|||||||
Reference in New Issue
Block a user