1#!/usr/bin/env python 2 3""" 4Copyright (c) 2006-2019 sqlmap developers (http://sqlmap.org/) 5See the file 'LICENSE' for copying permission 6""" 7 8import hashlib 9import os 10import sqlite3 11import threading 12import time 13 14from lib.core.common import getSafeExString 15from lib.core.common import serializeObject 16from lib.core.common import singleTimeWarnMessage 17from lib.core.common import unserializeObject 18from lib.core.compat import xrange 19from lib.core.convert import getBytes 20from lib.core.convert import getUnicode 21from lib.core.data import logger 22from lib.core.exception import SqlmapConnectionException 23from lib.core.settings import HASHDB_END_TRANSACTION_RETRIES 24from lib.core.settings import HASHDB_FLUSH_RETRIES 25from lib.core.settings import HASHDB_FLUSH_THRESHOLD 26from lib.core.settings import HASHDB_RETRIEVE_RETRIES 27from lib.core.threads import getCurrentThreadData 28from lib.core.threads import getCurrentThreadName 29from thirdparty import six 30 31class HashDB(object): 32 def __init__(self, filepath): 33 self.filepath = filepath 34 self._write_cache = {} 35 self._cache_lock = threading.Lock() 36 37 def _get_cursor(self): 38 threadData = getCurrentThreadData() 39 40 if threadData.hashDBCursor is None: 41 try: 42 connection = sqlite3.connect(self.filepath, timeout=3, isolation_level=None) 43 threadData.hashDBCursor = connection.cursor() 44 threadData.hashDBCursor.execute("CREATE TABLE IF NOT EXISTS storage (id INTEGER PRIMARY KEY, value TEXT)") 45 connection.commit() 46 except Exception as ex: 47 errMsg = "error occurred while opening a session " 48 errMsg += "file '%s' ('%s')" % (self.filepath, getSafeExString(ex)) 49 raise SqlmapConnectionException(errMsg) 50 51 return threadData.hashDBCursor 52 53 def _set_cursor(self, cursor): 54 threadData = getCurrentThreadData() 55 threadData.hashDBCursor = cursor 56 57 cursor = property(_get_cursor, _set_cursor) 58 59 def close(self): 60 threadData = getCurrentThreadData() 61 try: 62 if threadData.hashDBCursor: 63 threadData.hashDBCursor.close() 64 threadData.hashDBCursor.connection.close() 65 threadData.hashDBCursor = None 66 except: 67 pass 68 69 @staticmethod 70 def hashKey(key): 71 key = getBytes(key if isinstance(key, six.text_type) else repr(key)) 72 retVal = int(hashlib.md5(key).hexdigest(), 16) & 0x7fffffffffffffff # Reference: http://stackoverflow.com/a/4448400 73 return retVal 74 75 def retrieve(self, key, unserialize=False): 76 retVal = None 77 78 if key and (self._write_cache or os.path.isfile(self.filepath)): 79 hash_ = HashDB.hashKey(key) 80 retVal = self._write_cache.get(hash_) 81 if not retVal: 82 for _ in xrange(HASHDB_RETRIEVE_RETRIES): 83 try: 84 for row in self.cursor.execute("SELECT value FROM storage WHERE id=?", (hash_,)): 85 retVal = row[0] 86 except (sqlite3.OperationalError, sqlite3.DatabaseError) as ex: 87 if any(_ in getSafeExString(ex) for _ in ("locked", "no such table")): 88 warnMsg = "problem occurred while accessing session file '%s' ('%s')" % (self.filepath, getSafeExString(ex)) 89 singleTimeWarnMessage(warnMsg) 90 elif "Could not decode" in getSafeExString(ex): 91 break 92 else: 93 errMsg = "error occurred while accessing session file '%s' ('%s'). " % (self.filepath, getSafeExString(ex)) 94 errMsg += "If the problem persists please rerun with '--flush-session'" 95 raise SqlmapConnectionException(errMsg) 96 else: 97 break 98 99 time.sleep(1) 100 101 if retVal and unserialize: 102 try: 103 retVal = unserializeObject(retVal) 104 except: 105 retVal = None 106 warnMsg = "error occurred while unserializing value for session key '%s'. " % key 107 warnMsg += "If the problem persists please rerun with '--flush-session'" 108 logger.warn(warnMsg) 109 110 return retVal 111 112 def write(self, key, value, serialize=False): 113 if key: 114 hash_ = HashDB.hashKey(key) 115 self._cache_lock.acquire() 116 self._write_cache[hash_] = getUnicode(value) if not serialize else serializeObject(value) 117 self._cache_lock.release() 118 119 if getCurrentThreadName() in ('0', 'MainThread'): 120 self.flush() 121 122 def flush(self, forced=False): 123 if not self._write_cache: 124 return 125 126 if not forced and len(self._write_cache) < HASHDB_FLUSH_THRESHOLD: 127 return 128 129 self._cache_lock.acquire() 130 _ = self._write_cache 131 self._write_cache = {} 132 self._cache_lock.release() 133 134 try: 135 self.beginTransaction() 136 for hash_, value in _.items(): 137 retries = 0 138 while True: 139 try: 140 try: 141 self.cursor.execute("INSERT INTO storage VALUES (?, ?)", (hash_, value,)) 142 except sqlite3.IntegrityError: 143 self.cursor.execute("UPDATE storage SET value=? WHERE id=?", (value, hash_,)) 144 except UnicodeError: # e.g. surrogates not allowed (Issue #3851) 145 break 146 except sqlite3.DatabaseError as ex: 147 if not os.path.exists(self.filepath): 148 debugMsg = "session file '%s' does not exist" % self.filepath 149 logger.debug(debugMsg) 150 break 151 152 if retries == 0: 153 warnMsg = "there has been a problem while writing to " 154 warnMsg += "the session file ('%s')" % getSafeExString(ex) 155 logger.warn(warnMsg) 156 157 if retries >= HASHDB_FLUSH_RETRIES: 158 return 159 else: 160 retries += 1 161 time.sleep(1) 162 else: 163 break 164 finally: 165 self.endTransaction() 166 167 def beginTransaction(self): 168 threadData = getCurrentThreadData() 169 if not threadData.inTransaction: 170 try: 171 self.cursor.execute("BEGIN TRANSACTION") 172 except: 173 # Reference: http://stackoverflow.com/a/25245731 174 self.cursor.close() 175 threadData.hashDBCursor = None 176 self.cursor.execute("BEGIN TRANSACTION") 177 finally: 178 threadData.inTransaction = True 179 180 def endTransaction(self): 181 threadData = getCurrentThreadData() 182 if threadData.inTransaction: 183 retries = 0 184 while retries < HASHDB_END_TRANSACTION_RETRIES: 185 try: 186 self.cursor.execute("END TRANSACTION") 187 threadData.inTransaction = False 188 except sqlite3.OperationalError: 189 pass 190 else: 191 return 192 193 retries += 1 194 time.sleep(1) 195 196 try: 197 self.cursor.execute("ROLLBACK TRANSACTION") 198 except sqlite3.OperationalError: 199 self.cursor.close() 200 self.cursor = None 201 finally: 202 threadData.inTransaction = False 203