first case where partial output is retrievable via RESTful API - issue #297

bdamele · bdamele · commit f7d826fee168 · 2013-02-05T14:43:03.000Z
diff --git a/lib/core/enums.py b/lib/core/enums.py
@@ -271,6 +271,34 @@ class CONTENT_TYPE:
     OS_CMD = 23
     REG_READ = 24
 
+PART_RUN_CONTENT_TYPES = {
+    "checkDbms": CONTENT_TYPE.TECHNIQUES,
+    "getFingerprint": CONTENT_TYPE.DBMS_FINGERPRINT,
+    "getBanner": CONTENT_TYPE.BANNER,
+    "getCurrentUser": CONTENT_TYPE.CURRENT_USER,
+    "getCurrentDb": CONTENT_TYPE.CURRENT_DB,
+    "getHostname": CONTENT_TYPE.HOSTNAME,
+    "isDba": CONTENT_TYPE.IS_DBA,
+    "getUsers": CONTENT_TYPE.USERS,
+    "getPasswordHashes": CONTENT_TYPE.PASSWORDS,
+    "getPrivileges": CONTENT_TYPE.PRIVILEGES,
+    "getRoles": CONTENT_TYPE.ROLES,
+    "getDbs": CONTENT_TYPE.DBS,
+    "getTables": CONTENT_TYPE.TABLES,
+    "getColumns": CONTENT_TYPE.COLUMNS,
+    "getSchema": CONTENT_TYPE.SCHEMA,
+    "getCount": CONTENT_TYPE.COUNT,
+    "dumpTable": CONTENT_TYPE.DUMP_TABLE,
+    "search": CONTENT_TYPE.SEARCH,
+    "sqlQuery": CONTENT_TYPE.SQL_QUERY,
+    "tableExists": CONTENT_TYPE.COMMON_TABLES,
+    "columnExists": CONTENT_TYPE.COMMON_COLUMNS,
+    "readFile": CONTENT_TYPE.FILE_READ,
+    "writeFile": CONTENT_TYPE.FILE_WRITE,
+    "osCmd": CONTENT_TYPE.OS_CMD,
+    "regRead": CONTENT_TYPE.REG_READ
+}
+
 class CONTENT_STATUS:
     IN_PROGRESS = 0
     COMPLETE = 1
diff --git a/lib/utils/api.py b/lib/utils/api.py
@@ -27,6 +27,7 @@
 from lib.core.datatype import AttribDict
 from lib.core.defaults import _defaults
 from lib.core.enums import CONTENT_STATUS
+from lib.core.enums import PART_RUN_CONTENT_TYPES
 from lib.core.log import LOGGER_HANDLER
 from lib.core.optiondict import optDict
 from lib.core.subprocessng import Popen
@@ -176,16 +177,16 @@ def __init__(self, taskid, messagetype="stdout"):
     def write(self, value, status=CONTENT_STATUS.IN_PROGRESS, content_type=None):
         if self.messagetype == "stdout":
             if content_type is None:
-                content_type = 99
+                if kb.partRun is not None:
+                    content_type = PART_RUN_CONTENT_TYPES.get(kb.partRun)
+                else:
+                    # Ignore all non-relevant messages
+                    return
 
             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:
-                # 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)))
