#!/usr/bin/env python

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

# 导入所需的Python标准库
import glob  # 用于文件路径模式匹配
import os    # 提供与操作系统交互的功能
import re    # 正则表达式模块
import shutil  # 提供高级文件操作功能
import subprocess  # 用于创建子进程
import time   # 时间相关功能
import zipfile  # ZIP文件操作

# 从自定义库中导入所需函数
from lib.core.common import dataToStdout  # 输出数据到标准输出
from lib.core.common import extractRegexResult  # 提取正则表达式结果
from lib.core.common import getLatestRevision  # 获取最新版本号
from lib.core.common import getSafeExString  # 安全地获取异常字符串
from lib.core.common import openFile  # 打开文件
from lib.core.common import pollProcess  # 轮询进程
from lib.core.common import readInput  # 读取用户输入
from lib.core.convert import getText  # 文本转换
from lib.core.data import conf  # 配置数据
from lib.core.data import logger  # 日志记录器
from lib.core.data import paths  # 路径信息
from lib.core.revision import getRevisionNumber  # 获取版本号
from lib.core.settings import GIT_REPOSITORY  # Git仓库地址
from lib.core.settings import IS_WIN  # 是否为Windows系统
from lib.core.settings import VERSION  # 版本信息
from lib.core.settings import TYPE  # 安装类型
from lib.core.settings import ZIPBALL_PAGE  # ZIP包下载页面
from thirdparty.six.moves import urllib as _urllib  # URL处理

def update():
    """
    更新sqlmap的主函数
    """
    # 如果未启用更新全部选项,直接返回
    if not conf.updateAll:
        return

    success = False  # 更新是否成功的标志

    # 如果是通过pip安装的
    if TYPE == "pip":
        infoMsg = "updating sqlmap to the latest stable version from the "
        infoMsg += "PyPI repository"
        logger.info(infoMsg)  # 记录更新信息

        debugMsg = "sqlmap will try to update itself using 'pip' command"
        logger.debug(debugMsg)  # 记录调试信息

        dataToStdout("\r[%s] [INFO] update in progress" % time.strftime("%X"))  # 显示更新进度

        output = ""
        try:
            # 执行pip更新命令
            process = subprocess.Popen("pip install -U sqlmap", shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=paths.SQLMAP_ROOT_PATH)
            pollProcess(process, True)  # 轮询进程
            output, _ = process.communicate()  # 获取输出
            success = not process.returncode  # 检查返回码
        except Exception as ex:
            success = False
            output = getSafeExString(ex)
        finally:
            output = getText(output)

        # 根据更新结果输出相应信息
        if success:
            logger.info("%s the latest revision '%s'" % ("already at" if "already up-to-date" in output else "updated to", extractRegexResult(r"\binstalled sqlmap-(?P<result>\d+\.\d+\.\d+)", output) or extractRegexResult(r"\((?P<result>\d+\.\d+\.\d+)\)", output)))
        else:
            logger.error("update could not be completed ('%s')" % re.sub(r"[^a-z0-9:/\\]+", " ", output).strip())

    # 如果不是Git仓库
    elif not os.path.exists(os.path.join(paths.SQLMAP_ROOT_PATH, ".git")):
        warnMsg = "not a git repository. It is recommended to clone the 'sqlmapproject/sqlmap' repository "
        warnMsg += "from GitHub (e.g. 'git clone --depth 1 %s sqlmap')" % GIT_REPOSITORY
        logger.warning(warnMsg)  # 提示用户使用git克隆

        # 检查是否已是最新版本
        if VERSION == getLatestRevision():
            logger.info("already at the latest revision '%s'" % (getRevisionNumber() or VERSION))
            return

        # 询问用户是否尝试下载ZIP包更新
        message = "do you want to try to fetch the latest 'zipball' from repository and extract it (experimental) ? [y/N]"
        if readInput(message, default='N', boolean=True):
            directory = os.path.abspath(paths.SQLMAP_ROOT_PATH)

            try:
                # 尝试创建/更新主程序文件
                open(os.path.join(directory, "sqlmap.py"), "w+b")
            except Exception as ex:
                errMsg = "unable to update content of directory '%s' ('%s')" % (directory, getSafeExString(ex))
                logger.error(errMsg)
            else:
                # 保存原文件属性
                attrs = os.stat(os.path.join(directory, "sqlmap.py")).st_mode
                # 清理目录内容
                for wildcard in ('*', ".*"):
                    for _ in glob.glob(os.path.join(directory, wildcard)):
                        try:
                            if os.path.isdir(_):
                                shutil.rmtree(_)
                            else:
                                os.remove(_)
                        except:
                            pass

                # 检查目录是否清空
                if glob.glob(os.path.join(directory, '*')):
                    errMsg = "unable to clear the content of directory '%s'" % directory
                    logger.error(errMsg)
                else:
                    try:
                        # 下载并解压最新的ZIP包
                        archive = _urllib.request.urlretrieve(ZIPBALL_PAGE)[0]

                        with zipfile.ZipFile(archive) as f:
                            for info in f.infolist():
                                info.filename = re.sub(r"\Asqlmap[^/]+", "", info.filename)
                                if info.filename:
                                    f.extract(info, directory)

                        # 获取并显示新版本信息
                        filepath = os.path.join(paths.SQLMAP_ROOT_PATH, "lib", "core", "settings.py")
                        if os.path.isfile(filepath):
                            with openFile(filepath, "rb") as f:
                                version = re.search(r"(?m)^VERSION\s*=\s*['\"]([^'\"]+)", f.read()).group(1)
                                logger.info("updated to the latest version '%s#dev'" % version)
                                success = True
                    except Exception as ex:
                        logger.error("update could not be completed ('%s')" % getSafeExString(ex))
                    else:
                        if not success:
                            logger.error("update could not be completed")
                        else:
                            try:
                                # 恢复文件属性
                                os.chmod(os.path.join(directory, "sqlmap.py"), attrs)
                            except OSError:
                                logger.warning("could not set the file attributes of '%s'" % os.path.join(directory, "sqlmap.py"))

    # 如果是Git仓库
    else:
        infoMsg = "updating sqlmap to the latest development revision from the "
        infoMsg += "GitHub repository"
        logger.info(infoMsg)

        debugMsg = "sqlmap will try to update itself using 'git' command"
        logger.debug(debugMsg)

        dataToStdout("\r[%s] [INFO] update in progress" % time.strftime("%X"))

        output = ""
        try:
            # 执行git更新命令
            process = subprocess.Popen("git checkout . && git pull %s HEAD" % GIT_REPOSITORY, shell=True, stdout=subprocess.PIPE, stderr=subprocess.STDOUT, cwd=paths.SQLMAP_ROOT_PATH)
            pollProcess(process, True)
            output, _ = process.communicate()
            success = not process.returncode
        except Exception as ex:
            success = False
            output = getSafeExString(ex)
        finally:
            output = getText(output)

        # 根据git更新结果输出信息
        if success:
            logger.info("%s the latest revision '%s'" % ("already at" if "Already" in output else "updated to", getRevisionNumber()))
        else:
            if "Not a git repository" in output:
                errMsg = "not a valid git repository. Please checkout the 'sqlmapproject/sqlmap' repository "
                errMsg += "from GitHub (e.g. 'git clone --depth 1 %s sqlmap')" % GIT_REPOSITORY
                logger.error(errMsg)
            else:
                logger.error("update could not be completed ('%s')" % re.sub(r"\W+", " ", output).strip())

    # 如果更新失败,根据操作系统给出建议
    if not success:
        if IS_WIN:
            infoMsg = "for Windows platform it's recommended "
            infoMsg += "to use a GitHub for Windows client for updating "
            infoMsg += "purposes (https://desktop.github.com/) or just "
            infoMsg += "download the latest snapshot from "
            infoMsg += "https://github.com/sqlmapproject/sqlmap/downloads"
        else:
            infoMsg = "for Linux platform it's recommended "
            infoMsg += "to install a standard 'git' package (e.g.: 'apt install git')"

        logger.info(infoMsg)