From 6ffc5665d0df2e6f7f3126d52e453e393219c758 Mon Sep 17 00:00:00 2001
From: Miroslav Stampar <miroslav.stampar@gmail.com>
Date: Fri, 27 Jul 2012 16:29:33 +0200
Subject: [PATCH] Update for Issue #28

---
 lib/core/option.py         | 18 +++++++++++++++++-
 lib/takeover/metasploit.py |  7 +++++++
 2 files changed, 24 insertions(+), 1 deletion(-)

diff --git a/lib/core/option.py b/lib/core/option.py
index 2cbdc0a7e..ace52e04d 100644
--- a/lib/core/option.py
+++ b/lib/core/option.py
@@ -588,6 +588,22 @@ def __setMetasploit():
     msfEnvPathExists = False
 
     if IS_WIN:
+        if not conf.msfPath:
+            def _(key, value):
+                retVal = None
+
+                try:
+                    from  _winreg import ConnectRegistry, OpenKey, QueryValueEx, HKEY_LOCAL_MACHINE
+                    _ = ConnectRegistry(None, HKEY_LOCAL_MACHINE)
+                    _ = OpenKey(_, key)
+                    retval = QueryValueEx(_, value)[0]
+                except:
+                    pass
+
+                return retVal
+
+            conf.msfPath = _(r"SOFTWARE\Rapid7\Metasploit", "Location")
+
         warnMsg = "some sqlmap takeover functionalities are not yet "
         warnMsg += "supported on Windows. Please use Linux in a virtual "
         warnMsg += "machine for out-of-band features."
@@ -607,7 +623,7 @@ def __setMetasploit():
             raise sqlmapMissingPrivileges, errMsg
 
     if conf.msfPath:
-        for path in (conf.msfPath, os.path.join(conf.msfPath, 'bin')):
+        for path in (conf.msfPath, os.path.join(conf.msfPath, "bin")):
             if all(os.path.exists(normalizePath(os.path.join(path, _))) for _ in ("", "msfcli", "msfconsole", "msfencode", "msfpayload")):
                 msfEnvPathExists = True
                 conf.msfPath = path
diff --git a/lib/takeover/metasploit.py b/lib/takeover/metasploit.py
index 71980b415..a8c302aad 100644
--- a/lib/takeover/metasploit.py
+++ b/lib/takeover/metasploit.py
@@ -32,6 +32,7 @@ from lib.core.enums import DBMS
 from lib.core.enums import OS
 from lib.core.exception import sqlmapDataException
 from lib.core.exception import sqlmapFilePathException
+from lib.core.settings import IS_WIN
 from lib.core.settings import UNICODE_ENCODING
 from lib.core.subprocessng import blockingReadFromFD
 from lib.core.subprocessng import blockingWriteToFD
@@ -58,6 +59,12 @@ class Metasploit:
         self.__msfEncode = normalizePath(os.path.join(conf.msfPath, "msfencode"))
         self.__msfPayload = normalizePath(os.path.join(conf.msfPath, "msfpayload"))
 
+        if IS_WIN:
+            _ = normalizePath(os.path.join(conf.msfPath, "..", "scripts", "setenv.bat"))
+            for attr in dir(self):
+                if "__msf" in attr:
+                    setattr(self, attr, "%s & %s" % (_, getattr(self, attr)))
+
         self.__msfPayloadsList = {
                                       "windows": {
                                                    1: ( "Meterpreter (default)", "windows/meterpreter" ),