Added resume functionality to -d and fixed logging with -d

bdamele · bdamele · commit eecee3b27499 · 2010-04-12T09:35:20.000Z
diff --git a/lib/core/convert.py b/lib/core/convert.py
@@ -27,7 +27,8 @@
 except:
     import md5
     import sha
-    
+
+import pickle
 import sys
 import struct
 import urllib
@@ -38,7 +39,13 @@ def base64decode(string):
     return string.decode("base64")
 
 def base64encode(string):
-    return string.encode("base64")[:-1]
+    return string.encode("base64")[:-1].replace("\n", "")
+
+def base64pickle(string):
+    return base64encode(pickle.dumps(string))
+
+def base64unpickle(string):
+    return pickle.loads(base64decode(string))
 
 def hexdecode(string):
     string = string.lower()
diff --git a/lib/core/target.py b/lib/core/target.py
@@ -123,14 +123,14 @@ def __setOutputResume():
         if not conf.flushSession:
             readSessionFP = open(conf.sessionFile, "r")
             lines = readSessionFP.readlines()
-    
+
             for line in lines:
                 if line.count("][") == 4:
                     line = line.split("][")
-    
+
                     if len(line) != 5:
                         continue
-    
+
                     url, _, _, expression, value = line
     
                     if not value:
@@ -141,10 +141,10 @@ def __setOutputResume():
     
                     if value[-1] == "\n":
                         value = value[:-1]
-    
-                    if url != conf.url:
+
+                    if url not in ( conf.url, conf.hostname ):
                         continue
-    
+
                     if url not in kb.resumedQueries.keys():
                         kb.resumedQueries[url] = {}
                         kb.resumedQueries[url][expression] = value
diff --git a/lib/request/direct.py b/lib/request/direct.py
@@ -23,8 +23,12 @@
 """
 
 from lib.core.agent import agent
+from lib.core.common import dataToSessionFile
+from lib.core.convert import base64pickle
+from lib.core.convert import base64unpickle
 from lib.core.data import conf
 from lib.core.data import kb
+from lib.core.data import logger
 from lib.core.settings import SQL_STATEMENTS
 from lib.utils.timeout import timeout
 
@@ -42,14 +46,25 @@ def direct(query, content=True):
                 select = True
                 break
 
-    if select:
-        output = timeout(func=conf.dbmsConnector.select, args=(query,), duration=conf.timeout, default=None)
-    else:
+    logger.log(9, query)
+
+    if not select:
         output = timeout(func=conf.dbmsConnector.execute, args=(query,), duration=conf.timeout, default=None)
+    elif conf.hostname in kb.resumedQueries and query in kb.resumedQueries[conf.hostname]:
+        output = base64unpickle(kb.resumedQueries[conf.hostname][query][:-1])
+
+        infoMsg  = "resumed from file '%s': " % conf.sessionFile
+        infoMsg += "%s..." % str(output)[:20]
+        logger.info(infoMsg)
+    elif select:
+        output = timeout(func=conf.dbmsConnector.select, args=(query,), duration=conf.timeout, default=None)
 
     if output is None or len(output) == 0:
         return None
     elif content:
+        if conf.hostname not in kb.resumedQueries or ( conf.hostname in kb.resumedQueries and query not in kb.resumedQueries[conf.hostname] ):
+            dataToSessionFile("[%s][%s][%s][%s][%s]\n" % (conf.hostname, kb.injPlace, conf.parameters[kb.injPlace], query, base64pickle(output)))
+
         if len(output) == 1:
             if len(output[0]) == 1:
                 return str(list(output)[0][0])
diff --git a/plugins/dbms/access/connector.py b/plugins/dbms/access/connector.py
@@ -67,8 +67,6 @@ def fetchall(self):
             return None
 
     def execute(self, query):
-        logger.debug(query)
-
         try:
             self.cursor.execute(query)
         except (pyodbc.OperationalError, pyodbc.ProgrammingError), msg:
diff --git a/plugins/dbms/firebird/connector.py b/plugins/dbms/firebird/connector.py
@@ -66,8 +66,6 @@ def fetchall(self):
             return None
 
     def execute(self, query):
-        logger.debug(query)
-
         try:
             self.cursor.execute(query)
         except kinterbasdb.OperationalError, msg:
diff --git a/plugins/dbms/mssqlserver/connector.py b/plugins/dbms/mssqlserver/connector.py
@@ -70,8 +70,6 @@ def fetchall(self):
             return None
 
     def execute(self, query):
-        logger.debug(query)
-
         try:
             self.cursor.execute(query)
         except (pymssql.OperationalError, pymssql.ProgrammingError), msg:
diff --git a/plugins/dbms/mysql/connector.py b/plugins/dbms/mysql/connector.py
@@ -66,8 +66,6 @@ def fetchall(self):
             return None
 
     def execute(self, query):
-        logger.debug(query)
-
         try:
             self.cursor.execute(query)
         except (MySQLdb.OperationalError, MySQLdb.ProgrammingError), msg:
diff --git a/plugins/dbms/oracle/connector.py b/plugins/dbms/oracle/connector.py
@@ -67,8 +67,6 @@ def fetchall(self):
             return None
 
     def execute(self, query):
-        logger.debug(query)
-
         try:
             self.cursor.execute(query)
         except (cx_Oracle.DatabaseError), msg:
diff --git a/plugins/dbms/postgresql/connector.py b/plugins/dbms/postgresql/connector.py
@@ -65,8 +65,6 @@ def fetchall(self):
             return None
 
     def execute(self, query):
-        logger.debug(query)
-
         try:
             self.cursor.execute(query)
         except (psycopg2.OperationalError, psycopg2.ProgrammingError), msg:
diff --git a/plugins/dbms/sqlite/connector.py b/plugins/dbms/sqlite/connector.py
@@ -67,8 +67,6 @@ def fetchall(self):
             return None
 
     def execute(self, query):
-        logger.debug(query)
-
         try:
             self.cursor.execute(query)
         except sqlite3.OperationalError, msg:
