Skip to content

Commit

Permalink
step by step getting there to partial output presentation to restful …
Browse files Browse the repository at this point in the history
…API (issue #297), not quite yet though..
  • Loading branch information
bdamele committed Feb 3, 2013
1 parent b55555e commit 9370f96
Show file tree
Hide file tree
Showing 5 changed files with 24 additions and 11 deletions.
7 changes: 5 additions & 2 deletions lib/core/common.py
Expand Up @@ -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
Expand Down Expand Up @@ -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):
"""
Expand Down
7 changes: 6 additions & 1 deletion lib/techniques/blind/inference.py
Expand Up @@ -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)
Expand Down
2 changes: 1 addition & 1 deletion lib/techniques/error/use.py
Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion lib/techniques/union/use.py
Expand Up @@ -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
Expand Down
17 changes: 11 additions & 6 deletions lib/utils/api.py
Expand Up @@ -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:
Expand Down Expand Up @@ -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")

Expand Down

0 comments on commit 9370f96

Please sign in to comment.