Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Added multiprocessing.Manager to fix MacOS Issue #8

Open
wants to merge 1 commit into
base: main
Choose a base branch
from
Open
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
16 changes: 8 additions & 8 deletions ntdissector/ntds/ntds.py
Original file line number Diff line number Diff line change
Expand Up @@ -22,6 +22,8 @@
from json import loads
from tqdm import tqdm
import multiprocessing as mp
import threading
import queue
from binascii import hexlify, unhexlify
from base64 import b64encode
from ntdissector.utils.crypto import PEK_LIST, format_asn1_to_pem
Expand Down Expand Up @@ -708,7 +710,7 @@ def __translate(self, name: str) -> str:
return name

# multiprocessing with a queue
def __dumpObjectWorker(self, workersQ: mp.Queue, workerLock: mp.Lock) -> None:
def __dumpObjectWorker(self, workersQ: queue.Queue, workerLock: threading.Lock) -> None:
wid = mp.current_process().name.split("-")[1]
results = dict()
logging.debug(f"Worker-{wid} started ")
Expand Down Expand Up @@ -759,8 +761,9 @@ def printStats(stats: dict):
def __reduce(record):
return {"page_num": record._node.tag.page.num, "page_buf": bytes(record._node.tag.page.buf), "node_num": record._node.num}

workersQ = mp.Queue()
workerLock = mp.Lock()
mpManager = mp.Manager()
workersQ = mpManager.Queue()
workerLock = mpManager.Lock()
pool = mp.get_context("fork").Pool(self.workers, self.__dumpObjectWorker, (workersQ, workerLock))
logging.info("Filtering records with this list of object classes : %s" % classList)
if self.__skipDel:
Expand Down Expand Up @@ -813,16 +816,13 @@ def __reduce(record):
workersQ.put(None)

try:
workersQ.close()
workersQ.join_thread()
logging.debug("Worker queue closed successfully")
pool.close()
pool.join()
logging.debug("Process pool closed successfully")
except BrokenPipeError as e:
logging.error("BrokenPipeError while closing pool - %e" % e)
logging.error("BrokenPipeError while closing pool - %s" % e)
except Exception as e:
logging.error("Error while closing pool - %e" % e)
logging.error("Error while closing pool - %s" % e)
raise e

def getClasses(self) -> list:
Expand Down