Thanks to visit codestin.com
Credit goes to github.com

Skip to content

Commit f8bc747

Browse files
bdamelebdamele
committed
improvement to restful API to store to IPC database partial entries, not yet functional (issue #297)
1 parent a92f1fb commit f8bc747

8 files changed

Lines changed: 61 additions & 37 deletions

File tree

‎_sqlmap.py‎

Lines changed: 1 addition & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -129,8 +129,7 @@ def main():
129129

130130
if hasattr(conf, "api"):
131131
try:
132-
conf.database_cursor.close()
133-
conf.database_connection.close()
132+
conf.database_cursor.disconnect()
134133
except KeyboardInterrupt:
135134
pass
136135

‎lib/core/common.py‎

Lines changed: 3 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -58,6 +58,7 @@
5858
from lib.core.dicts import DEPRECATED_OPTIONS
5959
from lib.core.dicts import SQL_STATEMENTS
6060
from lib.core.enums import ADJUST_TIME_DELAY
61+
from lib.core.enums import CONTENT_STATUS
6162
from lib.core.enums import CHARSET_TYPE
6263
from lib.core.enums import DBMS
6364
from lib.core.enums import EXPECTED
@@ -744,7 +745,7 @@ def setColor(message, bold=False):
744745

745746
return retVal
746747

747-
def dataToStdout(data, forceOutput=False, bold=False, content_type=None, status=None):
748+
def dataToStdout(data, forceOutput=False, bold=False, content_type=None, status=CONTENT_STATUS.IN_PROGRESS):
748749
"""
749750
Writes text to the stdout (console) stream
750751
"""
@@ -762,8 +763,7 @@ def dataToStdout(data, forceOutput=False, bold=False, content_type=None, status=
762763
message = data
763764

764765
if hasattr(conf, "api"):
765-
if content_type and status:
766-
sys.stdout.write(message, status, content_type)
766+
sys.stdout.write(message, status, content_type)
767767
else:
768768
sys.stdout.write(setColor(message, bold))
769769

‎lib/core/dump.py‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -26,7 +26,7 @@
2626
from lib.core.data import kb
2727
from lib.core.data import logger
2828
from lib.core.dicts import DUMP_REPLACEMENTS
29-
from lib.core.enums import API_CONTENT_STATUS
29+
from lib.core.enums import CONTENT_STATUS
3030
from lib.core.enums import CONTENT_TYPE
3131
from lib.core.enums import DBMS
3232
from lib.core.enums import DUMP_FORMAT
@@ -55,7 +55,7 @@ def __init__(self):
5555

5656
def _write(self, data, newline=True, console=True, content_type=None):
5757
if hasattr(conf, "api"):
58-
dataToStdout(data, content_type=content_type, status=API_CONTENT_STATUS.COMPLETE)
58+
dataToStdout(data, content_type=content_type, status=CONTENT_STATUS.COMPLETE)
5959
return
6060

6161
text = "%s%s" % (data, "\n" if newline else " ")

‎lib/core/enums.py‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -271,6 +271,6 @@ class CONTENT_TYPE:
271271
OS_CMD = 23
272272
REG_READ = 24
273273

274-
class API_CONTENT_STATUS:
274+
class CONTENT_STATUS:
275275
IN_PROGRESS = 0
276276
COMPLETE = 1

‎lib/techniques/blind/inference.py‎

Lines changed: 4 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -88,8 +88,8 @@ def bisection(payload, expression, length=None, charsetType=None, firstChar=None
8888

8989
try:
9090
# Set kb.partRun in case "common prediction" feature (a.k.a. "good
91-
# samaritan") is used
92-
kb.partRun = getPartRun() if conf.predictOutput else None
91+
# samaritan") is used or the engine is called from the API
92+
kb.partRun = getPartRun() if conf.predictOutput or hasattr(conf, "api") else None
9393

9494
if partialValue:
9595
firstChar = len(partialValue)
@@ -486,7 +486,7 @@ def blindThread():
486486
if result:
487487
if showEta:
488488
etaProgressUpdate(time.time() - charStart, len(commonValue))
489-
elif conf.verbose in (1, 2):
489+
elif conf.verbose in (1, 2) or hasattr(conf, "api"):
490490
dataToStdout(filterControlChars(commonValue[index - 1:]))
491491

492492
finalValue = commonValue
@@ -534,7 +534,7 @@ def blindThread():
534534

535535
if showEta:
536536
etaProgressUpdate(time.time() - charStart, index)
537-
elif conf.verbose in (1, 2):
537+
elif conf.verbose in (1, 2) or hasattr(conf, "api"):
538538
dataToStdout(filterControlChars(val))
539539

540540
# some DBMSes (e.g. Firebird, DB2, etc.) have issues with trailing spaces

‎lib/techniques/error/use.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -16,6 +16,7 @@
1616
from lib.core.common import dataToStdout
1717
from lib.core.common import decodeHexValue
1818
from lib.core.common import extractRegexResult
19+
from lib.core.common import getPartRun
1920
from lib.core.common import getUnicode
2021
from lib.core.common import hashDBRetrieve
2122
from lib.core.common import hashDBWrite
@@ -243,6 +244,9 @@ def errorUse(expression, dump=False):
243244

244245
_, _, _, _, _, expressionFieldsList, expressionFields, _ = agent.getFields(expression)
245246

247+
# Set kb.partRun in case the engine is called from the API
248+
kb.partRun = getPartRun() if hasattr(conf, "api") else None
249+
246250
# We have to check if the SQL query might return multiple entries
247251
# and in such case forge the SQL limiting the query output one
248252
# entry at a time

‎lib/techniques/union/use.py‎

Lines changed: 4 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,7 @@
1919
from lib.core.common import extractRegexResult
2020
from lib.core.common import flattenValue
2121
from lib.core.common import getConsoleWidth
22+
from lib.core.common import getPartRun
2223
from lib.core.common import getUnicode
2324
from lib.core.common import hashDBRetrieve
2425
from lib.core.common import hashDBWrite
@@ -163,6 +164,9 @@ def unionUse(expression, unpack=True, dump=False):
163164

164165
_, _, _, _, _, expressionFieldsList, expressionFields, _ = agent.getFields(origExpr)
165166

167+
# Set kb.partRun in case the engine is called from the API
168+
kb.partRun = getPartRun() if hasattr(conf, "api") else None
169+
166170
if expressionFieldsList and len(expressionFieldsList) > 1 and "ORDER BY" in expression.upper():
167171
# Removed ORDER BY clause because UNION does not play well with it
168172
expression = re.sub("\s*ORDER BY\s+[\w,]+", "", expression, re.I)

‎lib/utils/api.py‎

Lines changed: 42 additions & 25 deletions
Original file line numberDiff line numberDiff line change
@@ -17,15 +17,16 @@
1717

1818
from lib.core.common import unArrayizeValue
1919
from lib.core.convert import base64pickle
20-
from lib.core.convert import base64unpickle
2120
from lib.core.convert import hexencode
2221
from lib.core.convert import dejsonize
2322
from lib.core.convert import jsonize
2423
from lib.core.data import conf
24+
from lib.core.data import kb
2525
from lib.core.data import paths
2626
from lib.core.data import logger
2727
from lib.core.datatype import AttribDict
2828
from lib.core.defaults import _defaults
29+
from lib.core.enums import CONTENT_STATUS
2930
from lib.core.log import LOGGER_HANDLER
3031
from lib.core.optiondict import optDict
3132
from lib.core.subprocessng import Popen
@@ -47,24 +48,27 @@
4748
# Local global variables
4849
adminid = ""
4950
db = None
51+
db_filepath = tempfile.mkstemp(prefix="sqlmapipc-", text=False)[1]
5052
tasks = dict()
5153

5254
# API objects
5355
class Database(object):
56+
global db_filepath
57+
5458
LOGS_TABLE = "CREATE TABLE logs(id INTEGER PRIMARY KEY AUTOINCREMENT, taskid INTEGER, time TEXT, level TEXT, message TEXT)"
5559
DATA_TABLE = "CREATE TABLE data(id INTEGER PRIMARY KEY AUTOINCREMENT, taskid INTEGER, status INTEGER, content_type INTEGER, value TEXT)"
5660
ERRORS_TABLE = "CREATE TABLE errors(id INTEGER PRIMARY KEY AUTOINCREMENT, taskid INTEGER, error TEXT)"
5761

58-
def __init__(self):
59-
pass
60-
61-
def create(self):
62-
_, self.database = tempfile.mkstemp(prefix="sqlmapipc-", text=False)
63-
logger.debug("IPC database: %s" % self.database)
62+
def __init__(self, database=None):
63+
if database:
64+
self.database = database
65+
else:
66+
self.database = db_filepath
6467

65-
def connect(self):
68+
def connect(self, who="server"):
6669
self.connection = sqlite3.connect(self.database, timeout=3, isolation_level=None)
6770
self.cursor = self.connection.cursor()
71+
logger.debug("REST-JSON API %s connected to IPC database" % who)
6872

6973
def disconnect(self):
7074
self.cursor.close()
@@ -79,18 +83,13 @@ def execute(self, statement, arguments=None):
7983
if statement.lstrip().upper().startswith("SELECT"):
8084
return self.cursor.fetchall()
8185

82-
def initialize(self):
83-
self.create()
84-
self.connect()
86+
def init(self):
8587
self.execute(self.LOGS_TABLE)
8688
self.execute(self.DATA_TABLE)
8789
self.execute(self.ERRORS_TABLE)
8890

89-
def get_filepath(self):
90-
return self.database
91-
9291
class Task(object):
93-
global db
92+
global db_filepath
9493

9594
def __init__(self, taskid):
9695
self.process = None
@@ -109,7 +108,7 @@ def initialize_options(self, taskid):
109108
# Let sqlmap engine knows it is getting called by the API, the task ID and the file path of the IPC database
110109
self.options.api = True
111110
self.options.taskid = taskid
112-
self.options.database = db.get_filepath()
111+
self.options.database = db_filepath
113112

114113
# Enforce batch mode and disable coloring
115114
self.options.batch = True
@@ -174,12 +173,25 @@ def __init__(self, taskid, messagetype="stdout"):
174173
else:
175174
sys.stderr = self
176175

177-
def write(self, value, status=None, content_type=None):
176+
def write(self, value, status=CONTENT_STATUS.IN_PROGRESS, content_type=None):
178177
if self.messagetype == "stdout":
179-
#conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
180-
# (self.taskid, status, content_type, base64pickle(value)))
181-
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
182-
(self.taskid, status, content_type, jsonize(value)))
178+
if content_type is None:
179+
content_type = 99
180+
181+
if status == CONTENT_STATUS.IN_PROGRESS:
182+
output = conf.database_cursor.execute("SELECT id, value FROM data WHERE taskid = ? AND status = ? AND content_type = ? LIMIT 0,1",
183+
(self.taskid, status, content_type))
184+
185+
if len(output) == 0:
186+
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
187+
(self.taskid, status, content_type, jsonize(value)))
188+
else:
189+
new_value = "%s%s" % (output[0][1], value)
190+
conf.database_cursor.execute("UPDATE data SET value = ? WHERE id = ?",
191+
(jsonize(new_value), output[0][0]))
192+
else:
193+
conf.database_cursor.execute("INSERT INTO data VALUES(NULL, ?, ?, ?, ?)",
194+
(self.taskid, status, content_type, jsonize(value)))
183195
else:
184196
conf.database_cursor.execute("INSERT INTO errors VALUES(NULL, ?, ?)",
185197
(self.taskid, str(value) if value else ""))
@@ -205,8 +217,11 @@ def emit(self, record):
205217

206218
def setRestAPILog():
207219
if hasattr(conf, "api"):
208-
conf.database_connection = sqlite3.connect(conf.database, timeout=1, isolation_level=None)
209-
conf.database_cursor = conf.database_connection.cursor()
220+
#conf.database_connection = sqlite3.connect(conf.database, timeout=1, isolation_level=None)
221+
#conf.database_cursor = conf.database_connection.cursor()
222+
223+
conf.database_cursor = Database(conf.database)
224+
conf.database_cursor.connect("client")
210225

211226
# Set a logging handler that writes log messages to a IPC database
212227
logger.removeHandler(LOGGER_HANDLER)
@@ -455,7 +470,6 @@ def scan_data(taskid):
455470

456471
# Read all data from the IPC database for the taskid
457472
for status, content_type, value in db.execute("SELECT status, content_type, value FROM data WHERE taskid = ? ORDER BY id ASC", (taskid,)):
458-
#json_data_message.append({"status": status, "type": content_type, "value": base64unpickle(value)})
459473
json_data_message.append({"status": status, "type": content_type, "value": dejsonize(value)})
460474

461475
# Read all error messages from the IPC database
@@ -536,15 +550,18 @@ def server(host="0.0.0.0", port=RESTAPI_SERVER_PORT):
536550
"""
537551
global adminid
538552
global db
553+
global db_filepath
539554

540555
adminid = hexencode(os.urandom(16))
541556

542557
logger.info("Running REST-JSON API server at '%s:%d'.." % (host, port))
543558
logger.info("Admin ID: %s" % adminid)
559+
logger.debug("IPC database: %s" % db_filepath)
544560

545561
# Initialize IPC database
546562
db = Database()
547-
db.initialize()
563+
db.connect()
564+
db.init()
548565

549566
# Run RESTful API
550567
run(host=host, port=port, quiet=True, debug=False)

0 commit comments

Comments
 (0)