#!/usr/bin/env python

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

# 导入所需的标准库和自定义模块
import hashlib  # 用于计算哈希值
import os  # 用于文件和路径操作
import sqlite3  # SQLite数据库操作
import threading  # 多线程支持
import time  # 时间相关操作

# 导入自定义工具函数
from lib.core.common import getSafeExString  # 安全地获取异常字符串
from lib.core.common import serializeObject  # 对象序列化
from lib.core.common import singleTimeWarnMessage  # 单次警告消息
from lib.core.common import unserializeObject  # 对象反序列化
from lib.core.compat import xrange  # 兼容Python2/3的range函数
from lib.core.convert import getBytes  # 转换为字节
from lib.core.convert import getUnicode  # 转换为Unicode
from lib.core.data import logger  # 日志记录器
from lib.core.exception import SqlmapConnectionException  # 自定义连接异常
from lib.core.settings import HASHDB_END_TRANSACTION_RETRIES  # 事务结束重试次数
from lib.core.settings import HASHDB_FLUSH_RETRIES  # 刷新重试次数
from lib.core.settings import HASHDB_FLUSH_THRESHOLD  # 刷新阈值
from lib.core.settings import HASHDB_RETRIEVE_RETRIES  # 检索重试次数
from lib.core.threads import getCurrentThreadData  # 获取当前线程数据
from lib.core.threads import getCurrentThreadName  # 获取当前线程名称
from thirdparty import six  # Python 2/3 兼容库

class HashDB(object):
    """
    哈希数据库类,用于管理SQLite数据库中的键值存储
    """
    def __init__(self, filepath):
        """
        初始化哈希数据库
        @param filepath: 数据库文件路径
        """
        self.filepath = filepath  # 数据库文件路径
        self._write_cache = {}  # 写入缓存字典
        self._cache_lock = threading.Lock()  # 缓存锁,用于线程同步
        self._connections = []  # 数据库连接列表

    def _get_cursor(self):
        """
        获取数据库游标
        @return: SQLite游标对象
        """
        threadData = getCurrentThreadData()  # 获取当前线程数据

        # 如果当前线程没有游标,则创建新的连接和游标
        if threadData.hashDBCursor is None:
            try:
                # 创建SQLite连接,禁用事务自动提交
                connection = sqlite3.connect(self.filepath, timeout=3, isolation_level=None)
                self._connections.append(connection)
                threadData.hashDBCursor = connection.cursor()
                # 创建存储表(如果不存在)
                threadData.hashDBCursor.execute("CREATE TABLE IF NOT EXISTS storage (id INTEGER PRIMARY KEY, value TEXT)")
                connection.commit()
            except Exception as ex:
                errMsg = "error occurred while opening a session "
                errMsg += "file '%s' ('%s')" % (self.filepath, getSafeExString(ex))
                raise SqlmapConnectionException(errMsg)

        return threadData.hashDBCursor

    def _set_cursor(self, cursor):
        """
        设置数据库游标
        @param cursor: SQLite游标对象
        """
        threadData = getCurrentThreadData()
        threadData.hashDBCursor = cursor

    # 游标属性
    cursor = property(_get_cursor, _set_cursor)

    def close(self):
        """
        关闭当前线程的数据库连接
        """
        threadData = getCurrentThreadData()
        try:
            if threadData.hashDBCursor:
                threadData.hashDBCursor.connection.commit()
                threadData.hashDBCursor.close()
                threadData.hashDBCursor.connection.close()
                threadData.hashDBCursor = None
        except:
            pass

    def closeAll(self):
        """
        关闭所有数据库连接
        """
        for connection in self._connections:
            try:
                connection.commit()
                connection.close()
            except:
                pass

    @staticmethod
    def hashKey(key):
        """
        计算键的哈希值
        @param key: 要哈希的键
        @return: 64位整数哈希值
        """
        key = getBytes(key if isinstance(key, six.text_type) else repr(key), errors="xmlcharrefreplace")
        retVal = int(hashlib.md5(key).hexdigest(), 16) & 0x7fffffffffffffff  # 确保返回64位正整数
        return retVal

    def retrieve(self, key, unserialize=False):
        """
        从数据库检索值
        @param key: 要检索的键
        @param unserialize: 是否需要反序列化
        @return: 检索到的值
        """
        retVal = None

        if key and (self._write_cache or os.path.isfile(self.filepath)):
            hash_ = HashDB.hashKey(key)
            # 首先检查缓存
            retVal = self._write_cache.get(hash_)
            if not retVal:
                # 多次尝试从数据库检索
                for _ in xrange(HASHDB_RETRIEVE_RETRIES):
                    try:
                        for row in self.cursor.execute("SELECT value FROM storage WHERE id=?", (hash_,)):
                            retVal = row[0]
                    except (sqlite3.OperationalError, sqlite3.DatabaseError) as ex:
                        if any(_ in getSafeExString(ex) for _ in ("locked", "no such table")):
                            warnMsg = "problem occurred while accessing session file '%s' ('%s')" % (self.filepath, getSafeExString(ex))
                            singleTimeWarnMessage(warnMsg)
                        elif "Could not decode" in getSafeExString(ex):
                            break
                        else:
                            errMsg = "error occurred while accessing session file '%s' ('%s'). " % (self.filepath, getSafeExString(ex))
                            errMsg += "If the problem persists please rerun with '--flush-session'"
                            raise SqlmapConnectionException(errMsg)
                    else:
                        break

                    time.sleep(1)

        # 如果需要反序列化
        if retVal and unserialize:
            try:
                retVal = unserializeObject(retVal)
            except:
                retVal = None
                warnMsg = "error occurred while unserializing value for session key '%s'. " % key
                warnMsg += "If the problem persists please rerun with '--flush-session'"
                logger.warning(warnMsg)

        return retVal

    def write(self, key, value, serialize=False):
        """
        写入值到数据库
        @param key: 键
        @param value: 值
        @param serialize: 是否需要序列化
        """
        if key:
            hash_ = HashDB.hashKey(key)
            self._cache_lock.acquire()
            self._write_cache[hash_] = getUnicode(value) if not serialize else serializeObject(value)
            self._cache_lock.release()

        # 主线程自动刷新缓存
        if getCurrentThreadName() in ('0', "MainThread"):
            self.flush()

    def flush(self, forced=False):
        """
        将缓存刷新到数据库
        @param forced: 是否强制刷新
        """
        if not self._write_cache:
            return

        # 如果未强制刷新且缓存未达到阈值,则不刷新
        if not forced and len(self._write_cache) < HASHDB_FLUSH_THRESHOLD:
            return

        # 获取并清空缓存
        self._cache_lock.acquire()
        _ = self._write_cache
        self._write_cache = {}
        self._cache_lock.release()

        try:
            self.beginTransaction()
            # 遍历缓存项写入数据库
            for hash_, value in _.items():
                retries = 0
                while True:
                    try:
                        try:
                            self.cursor.execute("INSERT INTO storage VALUES (?, ?)", (hash_, value,))
                        except sqlite3.IntegrityError:
                            self.cursor.execute("UPDATE storage SET value=? WHERE id=?", (value, hash_,))
                    except (UnicodeError, OverflowError):  # 处理编码错误
                        break
                    except sqlite3.DatabaseError as ex:
                        if not os.path.exists(self.filepath):
                            debugMsg = "session file '%s' does not exist" % self.filepath
                            logger.debug(debugMsg)
                            break

                        if retries == 0:
                            warnMsg = "there has been a problem while writing to "
                            warnMsg += "the session file ('%s')" % getSafeExString(ex)
                            logger.warning(warnMsg)

                        if retries >= HASHDB_FLUSH_RETRIES:
                            return
                        else:
                            retries += 1
                            time.sleep(1)
                    else:
                        break
        finally:
            self.endTransaction()

    def beginTransaction(self):
        """
        开始数据库事务
        """
        threadData = getCurrentThreadData()
        if not threadData.inTransaction:
            try:
                self.cursor.execute("BEGIN TRANSACTION")
            except:
                try:
                    self.cursor.close()
                except sqlite3.ProgrammingError:
                    pass
                threadData.hashDBCursor = None
                self.cursor.execute("BEGIN TRANSACTION")
            finally:
                threadData.inTransaction = True

    def endTransaction(self):
        """
        结束数据库事务
        """
        threadData = getCurrentThreadData()
        if threadData.inTransaction:
            retries = 0
            while retries < HASHDB_END_TRANSACTION_RETRIES:
                try:
                    self.cursor.execute("END TRANSACTION")
                    threadData.inTransaction = False
                except sqlite3.OperationalError:
                    pass
                except sqlite3.ProgrammingError:
                    self.cursor = None
                    threadData.inTransaction = False
                    return
                else:
                    return

                retries += 1
                time.sleep(1)

            try:
                self.cursor.execute("ROLLBACK TRANSACTION")
            except sqlite3.OperationalError:
                self.cursor.close()
                self.cursor = None
            finally:
                threadData.inTransaction = False