step by step getting there to partial output presentation to restful API (issue #297), not quite yet though..

bdamele · bdamele · commit 9370f96a6746 · 2013-02-03T22:09:33.000Z
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -178,18 +178,26 @@ def write(self, value, status=CONTENT_STATUS.IN_PROGRESS, content_type=None):
             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 @@ def emit(self, record):
 
 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")
 
