From 9370f96a67469bf2ecd6a66d0df41f732647eb0d Mon Sep 17 00:00:00 2001 From: Bernardo Damele Date: Sun, 3 Feb 2013 22:09:33 +0000 Subject: [PATCH] step by step getting there to partial output presentation to restful API (issue #297), not quite yet though.. --- lib/core/common.py | 7 +++++-- lib/techniques/blind/inference.py | 7 ++++++- lib/techniques/error/use.py | 2 +- lib/techniques/union/use.py | 2 +- lib/utils/api.py | 17 +++++++++++------ 5 files changed, 24 insertions(+), 11 deletions(-) diff --git a/lib/core/common.py b/lib/core/common.py index e02b76745..50acbfedf 100644 --- a/lib/core/common.py +++ b/lib/core/common.py @@ -1771,7 +1771,7 @@ def goGoodSamaritan(prevValue, originalCharset): else: return None, None, None, originalCharset -def getPartRun(): +def getPartRun(alias=True): """ Goes through call stack and finds constructs matching conf.dbmsHandler.*. Returns it or its alias used in txt/common-outputs.txt @@ -1803,7 +1803,10 @@ def getPartRun(): pass # Return the INI tag to consider for common outputs (e.g. 'Databases') - return commonPartsDict[retVal][1] if isinstance(commonPartsDict.get(retVal), tuple) else retVal + if alias: + return commonPartsDict[retVal][1] if isinstance(commonPartsDict.get(retVal), tuple) else retVal + else: + return retVal def getUnicode(value, encoding=None, system=False, noneToNull=False): """ diff --git a/lib/techniques/blind/inference.py b/lib/techniques/blind/inference.py index 43e4e486b..9a69cc90c 100644 --- a/lib/techniques/blind/inference.py +++ b/lib/techniques/blind/inference.py @@ -89,7 +89,12 @@ def bisection(payload, expression, length=None, charsetType=None, firstChar=None try: # Set kb.partRun in case "common prediction" feature (a.k.a. "good # samaritan") is used or the engine is called from the API - kb.partRun = getPartRun() if conf.predictOutput or hasattr(conf, "api") else None + if conf.predictOutput: + kb.partRun = getPartRun() + elif hasattr(conf, "api"): + kb.partRun = getPartRun(alias=False) + else: + kb.partRun = None if partialValue: firstChar = len(partialValue) diff --git a/lib/techniques/error/use.py b/lib/techniques/error/use.py index c3558a101..8d1066c8d 100644 --- a/lib/techniques/error/use.py +++ b/lib/techniques/error/use.py @@ -245,7 +245,7 @@ def errorUse(expression, dump=False): _, _, _, _, _, expressionFieldsList, expressionFields, _ = agent.getFields(expression) # Set kb.partRun in case the engine is called from the API - kb.partRun = getPartRun() if hasattr(conf, "api") else None + kb.partRun = getPartRun(alias=False) if hasattr(conf, "api") else None # We have to check if the SQL query might return multiple entries # and in such case forge the SQL limiting the query output one diff --git a/lib/techniques/union/use.py b/lib/techniques/union/use.py index 2062cf410..fadc8b604 100644 --- a/lib/techniques/union/use.py +++ b/lib/techniques/union/use.py @@ -165,7 +165,7 @@ def unionUse(expression, unpack=True, dump=False): _, _, _, _, _, expressionFieldsList, expressionFields, _ = agent.getFields(origExpr) # Set kb.partRun in case the engine is called from the API - kb.partRun = getPartRun() if hasattr(conf, "api") else None + kb.partRun = getPartRun(alias=False) if hasattr(conf, "api") else None if expressionFieldsList and len(expressionFieldsList) > 1 and "ORDER BY" in expression.upper(): # Removed ORDER BY clause because UNION does not play well with it diff --git a/lib/utils/api.py b/lib/utils/api.py index 20d141708..9681eb921 100644 --- a/lib/utils/api.py +++ b/lib/utils/api.py @@ -178,18 +178,26 @@ class StdDbOut(object): if content_type is None: content_type = 99 + output = conf.database_cursor.execute("SELECT id, value FROM data WHERE taskid = ? AND status = ? AND content_type = ? LIMIT 0,1", + (self.taskid, status, content_type)) + if status == CONTENT_STATUS.IN_PROGRESS: - output = conf.database_cursor.execute("SELECT id, value FROM data WHERE taskid = ? AND status = ? AND content_type = ? LIMIT 0,1", - (self.taskid, status, content_type)) + # Ignore all non-relevant messages + if kb.partRun is None: + return if len(output) == 0: conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)", (self.taskid, status, content_type, jsonize(value))) else: - new_value = "%s%s" % (output[0][1], value) + new_value = "%s%s" % (dejsonize(output[0][1]), value) conf.database_cursor.execute("UPDATE data SET value = ? WHERE id = ?", (jsonize(new_value), output[0][0])) else: + if len(output) > 0: + conf.database_cursor.execute("DELETE FROM data WHERE taskid = ? AND status = %s AND content_type = ?" % CONTENT_STATUS.IN_PROGRESS, + (self.taskid, content_type)) + conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)", (self.taskid, status, content_type, jsonize(value))) else: @@ -217,9 +225,6 @@ class LogRecorder(logging.StreamHandler): def setRestAPILog(): if hasattr(conf, "api"): - #conf.database_connection = sqlite3.connect(conf.database, timeout=1, isolation_level=None) - #conf.database_cursor = conf.database_connection.cursor() - conf.database_cursor = Database(conf.database) conf.database_cursor.connect("client")