diff --git a/lib/core/enums.py b/lib/core/enums.py index e1bd0eb3a..60e3bfd9c 100644 --- 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 index 9681eb921..22b1629da 100644 --- a/lib/utils/api.py +++ b/lib/utils/api.py @@ -27,6 +27,7 @@ from lib.core.data import logger 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 @@ class StdDbOut(object): 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)))