From a49f69b1227d3179d6b00ed5c996ee965f5208fe Mon Sep 17 00:00:00 2001 From: Maxime Meignan Date: Thu, 5 Oct 2023 18:01:19 +0200 Subject: [PATCH] ExtractOffsets.py: adds an optionnal control on number of threads --- Offsets/ExtractOffsets.py | 6 +++--- 1 file changed, 3 insertions(+), 3 deletions(-) diff --git a/Offsets/ExtractOffsets.py b/Offsets/ExtractOffsets.py index b0e7590..bc6efc0 100644 --- a/Offsets/ExtractOffsets.py +++ b/Offsets/ExtractOffsets.py @@ -10,7 +10,7 @@ import subprocess from concurrent.futures import ThreadPoolExecutor, as_completed import threading - +THREADS_LIMIT = None CSVLock = threading.Lock() machineType = dict(x86=332, x64=34404) @@ -102,7 +102,7 @@ def downloadPEFileFromMS(pe_basename, pe_ext, knownPEVersions, output_folder): i = 0 futures = set() lock = threading.Lock() - with ThreadPoolExecutor() as executor: + with ThreadPoolExecutor(max_workers=THREADS_LIMIT) as executor: for pe_hash in pe_list: entry = pe_list[pe_hash] futures.add(executor.submit(downloadSpecificFile, entry, pe_basename, pe_ext, knownPEVersions, output_folder, lock)) @@ -227,7 +227,7 @@ def extractOffsets(input_file, output_file, mode): elif os.path.isdir(input_file): print(f'[*] Processing folder: {input_file}') - with ThreadPoolExecutor() as extractorPool: + with ThreadPoolExecutor(max_workers=THREADS_LIMIT) as extractorPool: args = [(os.path.join(input_file, file), output_file, mode) for file in os.listdir(input_file)] for (i, res) in enumerate(extractorPool.map(extractOffsets, *zip(*args))): print(f"{i + 1}/{len(args)}", end="\r")