Mirrors/wavestone-cdt-edrsandblast
1
0
Fork 0
mirror of https://github.com/wavestone-cdt/EDRSandblast.git synced 2026-06-11 01:41:20 +00:00
Code Issues Packages Projects Releases Wiki Activity

ExtractOffsets.py: adds an optionnal control on number of threads

Browse Source
This commit is contained in:
Maxime Meignan
2023-10-05 18:01:19 +02:00
parent e479bef8f5
commit a49f69b122
1 changed files with 3 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
Offsets/ExtractOffsets.py
+3 -3
View File
@@ -10,7 +10,7 @@ import subprocess
from concurrent.futures import ThreadPoolExecutor, as_completed from concurrent.futures import ThreadPoolExecutor, as_completed
import threading import threading
THREADS_LIMIT = None
CSVLock = threading.Lock() CSVLock = threading.Lock()
machineType = dict(x86=332, x64=34404) machineType = dict(x86=332, x64=34404)
@@ -102,7 +102,7 @@ def downloadPEFileFromMS(pe_basename, pe_ext, knownPEVersions, output_folder):
i = 0 i = 0
futures = set() futures = set()
lock = threading.Lock() lock = threading.Lock()
with ThreadPoolExecutor() as executor: with ThreadPoolExecutor(max_workers=THREADS_LIMIT) as executor:
for pe_hash in pe_list: for pe_hash in pe_list:
entry = pe_list[pe_hash] entry = pe_list[pe_hash]
futures.add(executor.submit(downloadSpecificFile, entry, pe_basename, pe_ext, knownPEVersions, output_folder, lock)) 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): elif os.path.isdir(input_file):
print(f'[*] Processing folder: {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)] 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))): for (i, res) in enumerate(extractorPool.map(extractOffsets, *zip(*args))):
print(f"{i + 1}/{len(args)}", end="\r") print(f"{i + 1}/{len(args)}", end="\r")