<commit_before>#!/usr/bin/env python3

from cme.logger import setup_logger, setup_debug_logger, CMEAdapter
from cme.helpers.logger import highlight
from cme.helpers.misc import identify_target_file
from cme.parsers.ip import parse_targets
from cme.parsers.nmap import parse_nmap_xml
from cme.parsers.nessus import parse_nessus_file
from cme.cli import gen_cli_args
from cme.loaders.protocol_loader import protocol_loader
from cme.loaders.module_loader import module_loader
from cme.servers.http import CMEServer
from cme.first_run import first_run_setup
from cme.context import Context
from concurrent.futures import ThreadPoolExecutor
from pprint import pformat
from decimal import Decimal
import time
import asyncio
import aioconsole
import functools
import configparser
import cme.helpers.powershell as powershell
import cme
import shutil
import webbrowser
import sqlite3
import random
import os
import sys
import logging

setup_logger()
logger = CMEAdapter()

async def monitor_threadpool(pool, targets):
    logging.debug('Started thread poller')

    while True:
        try:
            text = await aioconsole.ainput("")
            if text == "":
                pool_size = pool._work_queue.qsize()
                finished_threads = len(targets) - pool_size
                percentage = Decimal(finished_threads) / Decimal(len(targets)) * Decimal(100)
                logger.info(f"completed: {percentage:.2f}% ({finished_threads}/{len(targets)})")
        except asyncio.CancelledError:
            logging.debug("Stopped thread poller")
            break

async def run_protocol(loop, protocol_obj, args, db, target, jitter):
    try:
        if jitter:
            value = random.choice(range(jitter[0], jitter[1]))
            logging.debug(f"Doin' the jitterbug for {value} second(s)")
            await asyncio.sleep(value)

        thread = loop.run_in_executor(
            None,
            functools.partial(
                protocol_obj,
                args,
                db,
                str(target)
            )
        )

        await asyncio.wait_for(
            thread,
            timeout=args.timeout
        )

    except asyncio.TimeoutError:
        logging.debug("Thread exceeded timeout")
    except asyncio.CancelledError:
        logging.debug("Stopping thread")
        thread.cancel()

async def start_threadpool(protocol_obj, args, db, targets, jitter):
    pool = ThreadPoolExecutor(max_workers=args.threads + 1)
    loop = asyncio.get_running_loop()
    loop.set_default_executor(pool)

    monitor_task = asyncio.create_task(
        monitor_threadpool(pool, targets)
    )

    jobs = [
        run_protocol(
            loop,
            protocol_obj,
            args,
            db,<commit_msg>Stop crackmapexec from crashing completely when concurrency issues cause lock-contention on the database - in SMB-mode.<commit_after>#!/usr/bin/env python3

from cme.logger import setup_logger, setup_debug_logger, CMEAdapter
from cme.helpers.logger import highlight
from cme.helpers.misc import identify_target_file
from cme.parsers.ip import parse_targets
from cme.parsers.nmap import parse_nmap_xml
from cme.parsers.nessus import parse_nessus_file
from cme.cli import gen_cli_args
from cme.loaders.protocol_loader import protocol_loader
from cme.loaders.module_loader import module_loader
from cme.servers.http import CMEServer
from cme.first_run import first_run_setup
from cme.context import Context
from concurrent.futures import ThreadPoolExecutor
from pprint import pformat
from decimal import Decimal
import time
import asyncio
import aioconsole
import functools
import configparser
import cme.helpers.powershell as powershell
import cme
import shutil
import webbrowser
import sqlite3
import random
import os
import sys
import logging

setup_logger()
logger = CMEAdapter()

async def monitor_threadpool(pool, targets):
    logging.debug('Started thread poller')

    while True:
        try:
            text = await aioconsole.ainput("")
            if text == "":
                pool_size = pool._work_queue.qsize()
                finished_threads = len(targets) - pool_size
                percentage = Decimal(finished_threads) / Decimal(len(targets)) * Decimal(100)
                logger.info(f"completed: {percentage:.2f}% ({finished_threads}/{len(targets)})")
        except asyncio.CancelledError:
            logging.debug("Stopped thread poller")
            break

async def run_protocol(loop, protocol_obj, args, db, target, jitter):
    try:
        if jitter:
            value = random.choice(range(jitter[0], jitter[1]))
            logging.debug(f"Doin' the jitterbug for {value} second(s)")
            await asyncio.sleep(value)

        thread = loop.run_in_executor(
            None,
            functools.partial(
                protocol_obj,
                args,
                db,
                str(target)
            )
        )

        await asyncio.wait_for(
            thread,
            timeout=args.timeout
        )

    except asyncio.TimeoutError:
        logging.debug("Thread exceeded timeout")
    except asyncio.CancelledError:
        logging.debug("Stopping thread")
        thread.cancel()
    except sqlite3.OperationalError as e:
        logging.debug("Sqlite error - sqlite3.operationalError - {}".format(str(e)))

async def start_threadpool(protocol_obj, args, db, targets, jitter):
    pool = ThreadPoolExecutor(max_workers=args.threads + 1)
    loop = asyncio.get_running_loop()
    loop.set_default_executor(pool)

    monitor_task = asyncio.create_task(
        monitor_threadpool(pool, targets)
    )

    jobs = [
        run_protocol(
            loop,
            protocol_obj,
            args,
            db,
{'fraction_non_alphanumeric': 0.07724498692240628, 'fraction_numerical': 0.0034873583260680036, 'mean_word_length': 2.744125326370757, 'pattern_counts': {'":': 2, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 2}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4183518', 'n_tokens_mistral': 1679, 'n_tokens_neox': 1625, 'n_words': 415}