#!/usr/bin/env python

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

import hashlib
import os
import sqlite3
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
from lib.core.convert import getBytes
from lib.core.convert import getUnicode
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

class HashDB(object):
 def __init__(self, filepath):
 self.filepath = filepath
 self._write_cache = {}
 self._cache_lock = threading.Lock()

 def _get_cursor(self):
 threadData = getCurrentThreadData()

 if threadData.hashDBCursor is None:
 try:
 connection = sqlite3.connect(self.filepath, timeout=3, isolation_level=None)
 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):
 threadData = getCurrentThreadData()
 threadData.hashDBCursor = cursor

 cursor = property(_get_cursor, _set_cursor)

 def close(self):
 threadData = getCurrentThreadData()
 try:
 if threadData.hashDBCursor:
 threadData.hashDBCursor.close()
 threadData.hashDBCursor.connection.close()
 threadData.hashDBCursor = None
 except:
 pass

 @staticmethod
 def hashKey(key):
 key = getBytes(key if isinstance(key, six.text_type) else repr(key))
 retVal = int(hashlib.md5(key).hexdigest(), 16) & 0x7fffffffffffffff # Reference: http://stackoverflow.com/a/4448400
 return retVal

 def retrieve(self, key, unserialize=False):
 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.warn(warnMsg)

 return retVal

 def write(self, key, value, serialize=False):
 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):
 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: # e.g. surrogates not allowed (Issue #3851)
 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.warn(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:
 # Reference: http://stackoverflow.com/a/25245731
 self.cursor.close()
 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
 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