#!/usr/bin/env python

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

# 导入所需的Python标准库
import importlib  # 用于动态导入模块
import logging   # 用于日志记录
import os        # 用于操作系统相关功能
import re        # 用于正则表达式操作
import sys       # 用于系统相关功能
import traceback # 用于异常追踪
import warnings  # 用于警告控制

# 尝试导入SQLAlchemy
_path = list(sys.path)  # 保存原始系统路径
_sqlalchemy = None
try:
    sys.path = sys.path[1:]  # 修改系统路径以避免命名冲突
    module = importlib.import_module("sqlalchemy")  # 动态导入SQLAlchemy
    if hasattr(module, "dialects"):
        _sqlalchemy = module
        warnings.simplefilter(action="ignore", category=_sqlalchemy.exc.SAWarning)  # 忽略SQLAlchemy的警告
except:
    pass
finally:
    sys.path = _path  # 恢复原始系统路径

# 尝试导入MySQL-python驱动
try:
    import MySQLdb  # SQLAlchemy在使用MySQL时需要
    warnings.filterwarnings("error", category=MySQLdb.Warning)  # 将MySQL警告转换为错误
except (ImportError, AttributeError):
    pass

# 导入项目内部模块
from lib.core.data import conf
from lib.core.data import logger
from lib.core.exception import SqlmapConnectionException
from lib.core.exception import SqlmapFilePathException
from lib.core.exception import SqlmapMissingDependence
from plugins.generic.connector import Connector as GenericConnector
from thirdparty import six
from thirdparty.six.moves import urllib as _urllib

def getSafeExString(ex, encoding=None):  # 交叉引用的函数
    raise NotImplementedError

class SQLAlchemy(GenericConnector):
    """SQLAlchemy连接器类，用于处理数据库连接和查询"""
    
    def __init__(self, dialect=None):
        """初始化SQLAlchemy连接器
        
        Args:
            dialect: 数据库方言(如mysql, postgresql等)
        """
        GenericConnector.__init__(self)

        self.dialect = dialect
        self.address = conf.direct  # 从配置获取数据库连接地址

        # 处理数据库用户名中的特殊字符
        if conf.dbmsUser:
            self.address = self.address.replace("'%s':" % conf.dbmsUser, "%s:" % _urllib.parse.quote(conf.dbmsUser))
            self.address = self.address.replace("%s:" % conf.dbmsUser, "%s:" % _urllib.parse.quote(conf.dbmsUser))

        # 处理数据库密码中的特殊字符
        if conf.dbmsPass:
            self.address = self.address.replace(":'%s'@" % conf.dbmsPass, ":%s@" % _urllib.parse.quote(conf.dbmsPass))
            self.address = self.address.replace(":%s@" % conf.dbmsPass, ":%s@" % _urllib.parse.quote(conf.dbmsPass))

        # 设置数据库方言
        if self.dialect:
            self.address = re.sub(r"\A.+://", "%s://" % self.dialect, self.address)

    def connect(self):
        """建立数据库连接"""
        if _sqlalchemy:
            self.initConnection()

            try:
                # 处理SQLite数据库文件路径
                if not self.port and self.db:
                    if not os.path.exists(self.db):
                        raise SqlmapFilePathException("the provided database file '%s' does not exist" % self.db)

                    _ = self.address.split("//", 1)
                    self.address = "%s////%s" % (_[0], os.path.abspath(self.db))

                # 根据不同数据库类型创建引擎
                if self.dialect == "sqlite":
                    engine = _sqlalchemy.create_engine(self.address, connect_args={"check_same_thread": False})
                elif self.dialect == "oracle":
                    engine = _sqlalchemy.create_engine(self.address)
                else:
                    engine = _sqlalchemy.create_engine(self.address, connect_args={})

                self.connector = engine.connect()  # 建立连接
            except (TypeError, ValueError):
                # 处理特定的连接错误
                if "_get_server_version_info" in traceback.format_exc():
                    try:
                        import pymssql
                        if int(pymssql.__version__[0]) < 2:
                            raise SqlmapConnectionException("SQLAlchemy connection issue (obsolete version of pymssql ('%s') is causing problems)" % pymssql.__version__)
                    except ImportError:
                        pass
                elif "invalid literal for int() with base 10: '0b" in traceback.format_exc():
                    raise SqlmapConnectionException("SQLAlchemy connection issue ('https://bitbucket.org/zzzeek/sqlalchemy/issues/3975')")
                else:
                    pass
            except SqlmapFilePathException:
                raise
            except Exception as ex:
                raise SqlmapConnectionException("SQLAlchemy connection issue ('%s')" % getSafeExString(ex))

            self.printConnected()
        else:
            raise SqlmapMissingDependence("SQLAlchemy not available (e.g. 'pip%s install SQLAlchemy')" % ('3' if six.PY3 else ""))

    def fetchall(self):
        """获取所有查询结果
        
        Returns:
            list: 查询结果列表，每个元素为一个元组
        """
        try:
            retVal = []
            for row in self.cursor.fetchall():
                retVal.append(tuple(row))
            return retVal
        except _sqlalchemy.exc.ProgrammingError as ex:
            logger.log(logging.WARN if conf.dbmsHandler else logging.DEBUG, "(remote) %s" % getSafeExString(ex))
            return None

    def execute(self, query):
        """执行SQL查询
        
        Args:
            query: SQL查询语句
            
        Returns:
            bool: 查询是否执行成功
        """
        retVal = False

        # 将查询转换为SQLAlchemy的text对象
        if hasattr(_sqlalchemy, "text"):
            query = _sqlalchemy.text(query)

        try:
            self.cursor = self.connector.execute(query)
            retVal = True
        except (_sqlalchemy.exc.OperationalError, _sqlalchemy.exc.ProgrammingError) as ex:
            logger.log(logging.WARN if conf.dbmsHandler else logging.DEBUG, "(remote) %s" % getSafeExString(ex))
        except _sqlalchemy.exc.InternalError as ex:
            raise SqlmapConnectionException(getSafeExString(ex))

        return retVal

    def select(self, query):
        """执行SELECT查询
        
        Args:
            query: SELECT查询语句
            
        Returns:
            list: 查询结果列表，如果查询失败返回None
        """
        retVal = None

        if self.execute(query):
            retVal = self.fetchall()

        return retVal
