#!/usr/bin/env python

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

# 导入urllib库并重命名为_urllib
from thirdparty.six.moves import urllib as _urllib

class SmartHTTPBasicAuthHandler(_urllib.request.HTTPBasicAuthHandler):
    """
    参考: http://selenic.com/hg/rev/6c51a5056020
    修复Bug: http://bugs.python.org/issue8797
    
    这是一个处理HTTP基础认证的智能处理器类,继承自HTTPBasicAuthHandler
    主要用于处理认证重试的逻辑,避免无限循环重试的问题
    """

    def __init__(self, *args, **kwargs):
        # 调用父类的初始化方法
        _urllib.request.HTTPBasicAuthHandler.__init__(self, *args, **kwargs)
        # 创建一个集合用于存储已重试过的请求
        self.retried_req = set()
        # 重试计数器初始化为0
        self.retried_count = 0

    def reset_retry_count(self):
        """
        重置重试计数的方法
        Python 2.6.5在遇到401或407错误时会调用此方法,可能导致无限循环
        因此这里禁用了重置功能,改为在http_error_auth_reqed中进行重置
        """
        pass

    def http_error_auth_reqed(self, auth_header, host, req, headers):
        """
        处理需要认证的HTTP错误
        
        参数说明:
        auth_header: 认证头信息
        host: 目标主机
        req: 请求对象
        headers: 请求头
        
        处理逻辑:
        1. 对每个新请求重置重试计数器
        2. 限制最大重试次数为5次
        3. 超过重试次数则抛出HTTP 401错误
        """
        # 如果是新的请求(通过hash判断)
        if hash(req) not in self.retried_req:
            # 将请求添加到已重试集合中
            self.retried_req.add(hash(req))
            # 重置重试计数
            self.retried_count = 0
        else:
            # 如果重试次数超过5次
            if self.retried_count > 5:
                # 抛出HTTP 401认证失败错误
                raise _urllib.error.HTTPError(req.get_full_url(), 401, "basic auth failed", headers, None)
            else:
                # 增加重试计数
                self.retried_count += 1

        # 调用父类的错误处理方法
        return _urllib.request.HTTPBasicAuthHandler.http_error_auth_reqed(self, auth_header, host, req, headers)
