#!/usr/bin/env python

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

# 导入所需的标准库和第三方库
import io
import re
import time
import types

# 导入自定义工具函数
from lib.core.common import getHostHeader
from lib.core.common import getSafeExString
from lib.core.common import logHTTPTraffic
from lib.core.common import readInput
from lib.core.convert import getBytes
from lib.core.convert import getUnicode
from lib.core.data import conf
from lib.core.data import kb
from lib.core.data import logger
from lib.core.enums import CUSTOM_LOGGING
from lib.core.enums import HTTP_HEADER
from lib.core.enums import HTTPMETHOD
from lib.core.enums import REDIRECTION
from lib.core.exception import SqlmapConnectionException
from lib.core.settings import DEFAULT_COOKIE_DELIMITER
from lib.core.settings import MAX_CONNECTION_READ_SIZE
from lib.core.settings import MAX_CONNECTION_TOTAL_SIZE
from lib.core.settings import MAX_SINGLE_URL_REDIRECTIONS
from lib.core.settings import MAX_TOTAL_REDIRECTIONS
from lib.core.threads import getCurrentThreadData
from lib.request.basic import decodePage
from lib.request.basic import parseResponse
from thirdparty import six
from thirdparty.six.moves import urllib as _urllib

class SmartRedirectHandler(_urllib.request.HTTPRedirectHandler):
    """
    智能重定向处理器类,继承自urllib的HTTPRedirectHandler
    用于处理HTTP重定向响应(301/302/303/307等)
    """

    def _get_header_redirect(self, headers):
        """
        从响应头中获取重定向URL
        :param headers: HTTP响应头
        :return: 重定向URL或None
        """
        retVal = None

        if headers:
            if HTTP_HEADER.LOCATION in headers:
                retVal = headers[HTTP_HEADER.LOCATION]
            elif HTTP_HEADER.URI in headers:
                retVal = headers[HTTP_HEADER.URI]

        return retVal

    def _ask_redirect_choice(self, redcode, redurl, method):
        """
        询问用户是否要跟随重定向
        :param redcode: 重定向状态码
        :param redurl: 重定向URL
        :param method: HTTP请求方法
        """
        with kb.locks.redirect:
            # 如果还没有做出重定向选择
            if kb.choices.redirect is None:
                msg = "got a %d redirect to " % redcode
                msg += "'%s'. Do you want to follow? [Y/n] " % redurl

                kb.choices.redirect = REDIRECTION.YES if readInput(msg, default='Y', boolean=True) else REDIRECTION.NO

            # 如果是POST请求导致的重定向,询问是否要重发POST数据
            if kb.choices.redirect == REDIRECTION.YES and method == HTTPMETHOD.POST and kb.resendPostOnRedirect is None:
                msg = "redirect is a result of a "
                msg += "POST request. Do you want to "
                msg += "resend original POST data to a new "
                msg += "location? [%s] " % ("Y/n" if not kb.originalPage else "y/N")

                kb.resendPostOnRedirect = readInput(msg, default=('Y' if not kb.originalPage else 'N'), boolean=True)

            if kb.resendPostOnRedirect:
                self.redirect_request = self._redirect_request

    def _redirect_request(self, req, fp, code, msg, headers, newurl):
        """
        创建重定向请求
        :return: 新的Request对象
        """
        return _urllib.request.Request(newurl.replace(' ', '%20'), data=req.data, headers=req.headers, origin_req_host=req.get_origin_req_host() if hasattr(req, "get_origin_req_host") else req.origin_req_host)

    def http_error_302(self, req, fp, code, msg, headers):
        """
        处理302重定向响应
        :param req: 原始请求
        :param fp: 响应文件对象
        :param code: 状态码
        :param msg: 状态消息
        :param headers: 响应头
        :return: 处理后的响应对象
        """
        start = time.time()
        content = None
        forceRedirect = False
        # 如果未忽略重定向,则获取重定向URL
        redurl = self._get_header_redirect(headers) if not conf.ignoreRedirects else None

        # 读取响应内容
        try:
            content = fp.read(MAX_CONNECTION_TOTAL_SIZE)
        except:  # 处理不完整读取
            content = b""
        finally:
            if content:
                try:  # 尝试将内容写回读取缓冲区以便后续使用
                    fp.fp._rbuf.truncate(0)
                    fp.fp._rbuf.write(content)
                except:
                    pass

        # 解码响应内容
        content = decodePage(content, headers.get(HTTP_HEADER.CONTENT_ENCODING), headers.get(HTTP_HEADER.CONTENT_TYPE))

        # 记录重定向信息
        threadData = getCurrentThreadData()
        threadData.lastRedirectMsg = (threadData.lastRequestUID, content)

        # 构建重定向日志消息
        redirectMsg = "HTTP redirect "
        redirectMsg += "[#%d] (%d %s):\r\n" % (threadData.lastRequestUID, code, getUnicode(msg))

        if headers:
            logHeaders = "\r\n".join("%s: %s" % (getUnicode(key.capitalize() if hasattr(key, "capitalize") else key), getUnicode(value)) for (key, value) in headers.items())
        else:
            logHeaders = ""

        redirectMsg += logHeaders
        if content:
            redirectMsg += "\r\n\r\n%s" % getUnicode(content[:MAX_CONNECTION_READ_SIZE])

        # 记录HTTP流量日志
        logHTTPTraffic(threadData.lastRequestMsg, redirectMsg, start, time.time())
        logger.log(CUSTOM_LOGGING.TRAFFIC_IN, redirectMsg)

        if redurl:
            try:
                # 如果重定向URL是相对路径,转换为绝对路径
                if not _urllib.parse.urlsplit(redurl).netloc:
                    redurl = _urllib.parse.urljoin(req.get_full_url(), redurl)

                # 检查是否存在无限重定向循环
                self._infinite_loop_check(req)
                if conf.scope:
                    # 检查重定向URL是否在指定范围内
                    if not re.search(conf.scope, redurl, re.I):
                        redurl = None
                    else:
                        forceRedirect = True
                else:
                    # 询问用户是否跟随重定向
                    self._ask_redirect_choice(code, redurl, req.get_method())
            except ValueError:
                redurl = None
                result = fp

        if redurl and (kb.choices.redirect == REDIRECTION.YES or forceRedirect):
            # 解析响应
            parseResponse(content, headers)

            # 更新请求头中的Host
            req.headers[HTTP_HEADER.HOST] = getHostHeader(redurl)
            if headers and HTTP_HEADER.SET_COOKIE in headers:
                # 处理Cookie
                cookies = dict()
                delimiter = conf.cookieDel or DEFAULT_COOKIE_DELIMITER
                last = None

                for part in getUnicode(req.headers.get(HTTP_HEADER.COOKIE, "")).split(delimiter) + ([headers[HTTP_HEADER.SET_COOKIE]] if HTTP_HEADER.SET_COOKIE in headers else []):
                    if '=' in part:
                        part = part.strip()
                        key, value = part.split('=', 1)
                        cookies[key] = value
                        last = key
                    elif last:
                        cookies[last] += "%s%s" % (delimiter, part)

                req.headers[HTTP_HEADER.COOKIE] = delimiter.join("%s=%s" % (key, cookies[key]) for key in cookies)

            try:
                # 执行重定向请求
                result = _urllib.request.HTTPRedirectHandler.http_error_302(self, req, fp, code, msg, headers)
            except _urllib.error.HTTPError as ex:
                result = ex

                # 处理特殊情况的hack
                try:
                    hasattr(result, "read")
                except KeyError:
                    class _(object):
                        pass
                    result = _()

                try:
                    result.info()
                except AttributeError:
                    def _(self):
                        return getattr(self, "hdrs", {})

                    result.info = types.MethodType(_, result)

                if not hasattr(result, "read"):
                    def _(self, length=None):
                        try:
                            retVal = getSafeExString(ex)
                        except:
                            retVal = ""
                        return getBytes(retVal)

                    result.read = types.MethodType(_, result)

                if not getattr(result, "url", None):
                    result.url = redurl

                if not getattr(result, "code", None):
                    result.code = 999
            except:
                redurl = None
                result = fp
                fp.read = io.BytesIO(b"").read
        else:
            result = fp

        # 记录最后的重定向URL
        threadData.lastRedirectURL = (threadData.lastRequestUID, redurl)

        # 设置结果属性
        result.redcode = code
        result.redurl = getUnicode(redurl) if six.PY3 else redurl
        return result

    # 其他重定向状态码使用相同的处理方法
    http_error_301 = http_error_303 = http_error_307 = http_error_302

    def _infinite_loop_check(self, req):
        """
        检查是否存在无限重定向循环
        :param req: 请求对象
        :raises: SqlmapConnectionException 如果检测到无限循环
        """
        if hasattr(req, 'redirect_dict') and (req.redirect_dict.get(req.get_full_url(), 0) >= MAX_SINGLE_URL_REDIRECTIONS or len(req.redirect_dict) >= MAX_TOTAL_REDIRECTIONS):
            errMsg = "infinite redirect loop detected (%s). " % ", ".join(item for item in req.redirect_dict.keys())
            errMsg += "Please check all provided parameters and/or provide missing ones"
            raise SqlmapConnectionException(errMsg)
