#!/usr/bin/env python

"""
Copyright (c) 2006-2024 sqlmap developers (https://sqlmap.org/)
See the file 'LICENSE' for copying permission
"""

# 导入处理zip文件的模块
import zipfile

# 导入一些工具函数和异常类
from lib.core.common import getSafeExString  # 用于安全地获取异常信息的字符串表示
from lib.core.common import isZipFile  # 用于判断文件是否为zip格式
from lib.core.exception import SqlmapDataException  # sqlmap数据相关异常
from lib.core.exception import SqlmapInstallationException  # sqlmap安装相关异常
from thirdparty import six  # Python 2/3 兼容性库

class Wordlist(six.Iterator):
    """
    用于遍历大型字典文件的迭代器类。
    这个类可以处理普通文本文件和zip压缩文件中的字典。
    支持多进程并行处理,可以将工作负载分配给不同进程。
    
    >>> from lib.core.option import paths
    >>> isinstance(next(Wordlist(paths.SMALL_DICT)), six.binary_type)
    True
    >>> isinstance(next(Wordlist(paths.WORDLIST)), six.binary_type)
    True
    """

    def __init__(self, filenames, proc_id=None, proc_count=None, custom=None):
        """
        初始化函数
        :param filenames: 字典文件名(可以是单个字符串或字符串列表)
        :param proc_id: 当前进程的ID(用于多进程并行处理)
        :param proc_count: 总进程数
        :param custom: 自定义的额外词列表(可选)
        """
        # 确保filenames始终是列表格式
        self.filenames = [filenames] if isinstance(filenames, six.string_types) else filenames
        self.fp = None  # 当前打开的文件对象
        self.index = 0  # 当前正在处理的文件索引
        self.counter = -1  # 已处理的行数计数器(从-1开始)
        self.current = None  # 当前正在处理的文件名
        self.iter = None  # 当前文件的迭代器对象
        self.custom = custom or []  # 存储自定义词列表,如果没有则为空列表
        self.proc_id = proc_id  # 当前进程的ID
        self.proc_count = proc_count  # 总进程数
        self.adjust()  # 初始化完成后,调整文件指针和迭代器状态

    def __iter__(self):
        """
        实现迭代器协议的__iter__方法
        使得这个类的实例可以在for循环中使用
        """
        return self

    def adjust(self):
        """
        调整文件指针和迭代器的状态
        在切换到新文件或重置迭代器时使用
        """
        self.closeFP()  # 首先关闭当前打开的文件(如果有)
        
        if self.index > len(self.filenames):
            return  # 如果已经处理完所有文件,直接返回
        elif self.index == len(self.filenames):
            self.iter = iter(self.custom)  # 如果处理完所有文件,切换到自定义词列表
        else:
            self.current = self.filenames[self.index]  # 获取当前要处理的文件名
            if isZipFile(self.current):  # 判断是否为zip文件
                try:
                    _ = zipfile.ZipFile(self.current, 'r')  # 尝试打开zip文件
                except zipfile.error as ex:
                    errMsg = "something appears to be wrong with "
                    errMsg += "the file '%s' ('%s'). Please make " % (self.current, getSafeExString(ex))
                    errMsg += "sure that you haven't made any changes to it"
                    raise SqlmapInstallationException(errMsg)
                if len(_.namelist()) == 0:  # 检查zip文件是否为空
                    errMsg = "no file(s) inside '%s'" % self.current
                    raise SqlmapDataException(errMsg)
                self.fp = _.open(_.namelist()[0])  # 打开zip中的第一个文件
            else:
                self.fp = open(self.current, "rb")  # 以二进制模式打开普通文件
            self.iter = iter(self.fp)  # 创建文件内容的迭代器

        self.index += 1  # 更新文件索引,为处理下一个文件做准备

    def closeFP(self):
        """
        关闭当前打开的文件
        防止资源泄露
        """
        if self.fp:
            self.fp.close()
            self.fp = None

    def __next__(self):
        """
        实现迭代器的next方法
        返回字典中的下一个词
        支持多进程处理时的任务分配
        """
        retVal = None
        while True:
            self.counter += 1  # 增加处理行数计数
            try:
                retVal = next(self.iter).rstrip()  # 获取下一行并去除末尾空白字符
            except zipfile.error as ex:
                errMsg = "something appears to be wrong with "
                errMsg += "the file '%s' ('%s'). Please make " % (self.current, getSafeExString(ex))
                errMsg += "sure that you haven't made any changes to it"
                raise SqlmapInstallationException(errMsg)
            except StopIteration:  # 当前文件处理完毕
                self.adjust()  # 切换到下一个文件
                retVal = next(self.iter).rstrip()
            # 在多进程模式下,根据进程ID筛选要处理的行
            if not self.proc_count or self.counter % self.proc_count == self.proc_id:
                break
        return retVal

    def rewind(self):
        """
        重置迭代器到开始位置
        允许重新遍历字典
        """
        self.index = 0  # 重置文件索引
        self.adjust()  # 重新调整文件指针和迭代器
