feat: multi process (GH-6)

This commit is contained in:
本光
2025-04-06 00:40:47 +08:00
committed by GitHub
parent 352d14ec82
commit 4363ff705c

View File

@@ -1,14 +1,15 @@
import argparse
from Crypto.Cipher import AES
import logging
import multiprocessing
import os
import subprocess
from multiprocessing import Pool
from typing import Dict, List, Tuple
from Crypto.Cipher import AES
from detect import detect_process
from runtime import RuntimeInfo
SUBPROCESS_TIMEOUT = 30
@@ -17,13 +18,11 @@ def general_aes_ctr_decrypt(data: bytes, key: bytes, nonce: bytes) -> bytes:
return cipher.decrypt(data)
def decrypt_process(runtimes: Dict[str, RuntimeInfo], sequences: List[Tuple[str, bytes]], args):
def decrypt_single_file(args_tuple):
path, data, runtime, args, output_dir = args_tuple
logger = logging.getLogger('shot')
output_dir: str = args.output_dir or args.directory
for path, data in sequences:
try:
serial_number = data[2:8].decode('utf-8')
runtime = runtimes[serial_number]
logger.info(f'Decrypting: {serial_number} ({path})')
dest_path = os.path.join(output_dir, path) if output_dir else path
@@ -43,14 +42,12 @@ def decrypt_process(runtimes: Dict[str, RuntimeInfo], sequences: List[Tuple[str,
f.write(b'\xa2' + runtime.mix_str_aes_nonce())
f.write(b'\xf0\xff')
f.write(data[:cipher_text_offset])
f.write(general_aes_ctr_decrypt(
data[cipher_text_offset:cipher_text_offset+cipher_text_length], runtime.runtime_aes_key, nonce))
f.write(general_aes_ctr_decrypt(data[cipher_text_offset : cipher_text_offset + cipher_text_length], runtime.runtime_aes_key, nonce))
f.write(data[cipher_text_offset + cipher_text_length :])
exe_name = 'pyarmor-1shot.exe' if os.name == 'nt' else 'pyarmor-1shot'
exe_path = os.path.join(
os.path.dirname(os.path.abspath(__file__)), exe_name)
# TODO: multi process
exe_path = os.path.join(os.path.dirname(os.path.abspath(__file__)), exe_name)
sp = subprocess.run(
[
exe_path,
@@ -65,11 +62,13 @@ def decrypt_process(runtimes: Dict[str, RuntimeInfo], sequences: List[Tuple[str,
for line in stdout:
logger.warning(f'PYCDC: {line} ({path})')
for line in stderr:
if line.startswith((
if line.startswith(
(
'Warning: Stack history is empty',
'Warning: Stack history is not empty',
'Warning: block stack is not empty',
)):
)
):
if args.show_warn_stack or args.show_all:
logger.warning(f'PYCDC: {line} ({path})')
elif line.startswith('Unsupported opcode:'):
@@ -82,10 +81,31 @@ def decrypt_process(runtimes: Dict[str, RuntimeInfo], sequences: List[Tuple[str,
logger.error(f'PYCDC: {line} ({path})')
if sp.returncode != 0:
logger.warning(f'PYCDC returned {sp.returncode} ({path})')
continue
return False
return True
except Exception as e:
logger.error(f'Decrypt failed: {e} ({path})')
continue
return False
def decrypt_process(runtimes: Dict[str, RuntimeInfo], sequences: List[Tuple[str, bytes]], args):
logger = logging.getLogger('shot')
output_dir: str = args.output_dir or args.directory
process_args = []
for path, data in sequences:
try:
serial_number = data[2:8].decode('utf-8')
runtime = runtimes[serial_number]
process_args.append((path, data, runtime, args, output_dir))
except Exception as e:
logger.error(f'Failed to prepare file for decryption: {e} ({path})')
num_processes = min(args.processes if hasattr(args, 'processes') else multiprocessing.cpu_count(), len(process_args))
logger.info(f'Starting decryption with {num_processes} processes for {len(process_args)} files')
with Pool(processes=num_processes) as pool:
results = pool.map(decrypt_single_file, process_args)
success_count = sum(1 for result in results if result)
logger.info(f'Decryption completed: {success_count} succeeded, {len(process_args) - success_count} failed')
def parse_args():