You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
281 lines
10 KiB
281 lines
10 KiB
#!/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
|