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