#!/usr/bin/env python

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

# 导入所需的Python标准库
import atexit  # 用于注册程序退出时的回调函数
import os     # 提供与操作系统交互的功能

# 导入自定义模块
from lib.core import readlineng as readline  # 导入readline模块用于命令行输入处理
from lib.core.common import getSafeExString  # 用于安全地获取异常的字符串表示
from lib.core.data import logger            # 日志记录器
from lib.core.data import paths             # 存储各种路径信息
from lib.core.enums import AUTOCOMPLETE_TYPE # 自动完成类型的枚举
from lib.core.enums import OS               # 操作系统类型的枚举
from lib.core.settings import IS_WIN        # 判断是否为Windows系统
from lib.core.settings import MAX_HISTORY_LENGTH # 历史记录最大长度

try:
    # 尝试导入rlcompleter模块用于命令补全
    import rlcompleter

    class CompleterNG(rlcompleter.Completer):
        """自定义的命令补全器类，继承自rlcompleter.Completer"""
        def global_matches(self, text):
            """
            计算简单名称的匹配项
            参数:
                text: 要匹配的文本
            返回:
                匹配的命令列表
            """
            matches = []
            n = len(text)

            for ns in (self.namespace,):
                for word in ns:
                    if word[:n] == text:
                        matches.append(word)

            return matches
except:
    # 如果导入失败，禁用readline功能
    readline._readline = None

def readlineAvailable():
    """
    检查readline模块是否可用
    在Windows系统的Python默认安装中通常不可用
    """
    return readline._readline is not None

def clearHistory():
    """清除命令行历史记录"""
    if not readlineAvailable():
        return

    readline.clear_history()

def saveHistory(completion=None):
    """
    保存命令行历史记录到文件
    参数:
        completion: 自动完成类型，决定历史记录保存的位置
    """
    try:
        if not readlineAvailable():
            return

        # 根据不同的自动完成类型选择不同的历史记录文件路径
        if completion == AUTOCOMPLETE_TYPE.SQL:
            historyPath = paths.SQL_SHELL_HISTORY
        elif completion == AUTOCOMPLETE_TYPE.OS:
            historyPath = paths.OS_SHELL_HISTORY
        elif completion == AUTOCOMPLETE_TYPE.API:
            historyPath = paths.API_SHELL_HISTORY
        else:
            historyPath = paths.SQLMAP_SHELL_HISTORY

        # 创建历史记录文件
        try:
            with open(historyPath, "w+"):
                pass
        except:
            pass

        # 设置历史记录最大长度并写入文件
        readline.set_history_length(MAX_HISTORY_LENGTH)
        try:
            readline.write_history_file(historyPath)
        except IOError as ex:
            warnMsg = "there was a problem writing the history file '%s' (%s)" % (historyPath, getSafeExString(ex))
            logger.warning(warnMsg)
    except KeyboardInterrupt:
        pass

def loadHistory(completion=None):
    """
    从文件加载命令行历史记录
    参数:
        completion: 自动完成类型，决定从哪个文件加载历史记录
    """
    if not readlineAvailable():
        return

    clearHistory()

    # 根据自动完成类型选择历史记录文件路径
    if completion == AUTOCOMPLETE_TYPE.SQL:
        historyPath = paths.SQL_SHELL_HISTORY
    elif completion == AUTOCOMPLETE_TYPE.OS:
        historyPath = paths.OS_SHELL_HISTORY
    elif completion == AUTOCOMPLETE_TYPE.API:
        historyPath = paths.API_SHELL_HISTORY
    else:
        historyPath = paths.SQLMAP_SHELL_HISTORY

    # 如果历史记录文件存在，尝试加载它
    if os.path.exists(historyPath):
        try:
            readline.read_history_file(historyPath)
        except IOError as ex:
            warnMsg = "there was a problem loading the history file '%s' (%s)" % (historyPath, getSafeExString(ex))
            logger.warning(warnMsg)
        except UnicodeError:
            if IS_WIN:
                warnMsg = "there was a problem loading the history file '%s'. " % historyPath
                warnMsg += "More info can be found at 'https://github.com/pyreadline/pyreadline/issues/30'"
                logger.warning(warnMsg)

def autoCompletion(completion=None, os=None, commands=None):
    """
    设置命令行自动完成功能
    参数:
        completion: 自动完成类型
        os: 操作系统类型
        commands: 自定义命令列表
    """
    if not readlineAvailable():
        return

    if completion == AUTOCOMPLETE_TYPE.OS:
        if os == OS.WINDOWS:
            # Windows系统的常用命令
            completer = CompleterNG({
                "attrib": None, "copy": None, "del": None,
                "dir": None, "echo": None, "fc": None,
                "label": None, "md": None, "mem": None,
                "move": None, "net": None, "netstat -na": None,
                "tree": None, "truename": None, "type": None,
                "ver": None, "vol": None, "xcopy": None,
            })

        else:
            # Unix/Linux系统的常用命令
            completer = CompleterNG({
                "cat": None, "chmod": None, "chown": None,
                "cp": None, "cut": None, "date": None, "df": None,
                "diff": None, "du": None, "echo": None, "env": None,
                "file": None, "find": None, "free": None, "grep": None,
                "id": None, "ifconfig": None, "ls": None, "mkdir": None,
                "mv": None, "netstat": None, "pwd": None, "rm": None,
                "uname": None, "whoami": None,
            })

        # 设置命令补全器
        readline.set_completer(completer.complete)
        readline.parse_and_bind("tab: complete")

    elif commands:
        # 使用自定义命令列表设置补全器
        completer = CompleterNG(dict(((_, None) for _ in commands)))
        readline.set_completer_delims(' ')
        readline.set_completer(completer.complete)
        readline.parse_and_bind("tab: complete")

    # 加载历史记录并注册退出时保存历史记录
    loadHistory(completion)
    atexit.register(saveHistory, completion)
