add info draft

infection_monkey/system_info/windows_info_collector.py

@@ -2,9 +2,12 @@ import os
 import logging
 import traceback
 
+import sys
+sys.coinit_flags = 0 # needed for proper destruction of the wmi python module
+import wmi
 import _winreg
-from wmi import WMI
-#from mimikatz_collector import MimikatzCollector
+
+from mimikatz_collector import MimikatzCollector
 from . import InfoCollector
 
 LOG = logging.getLogger(__name__)
@@ -17,10 +20,59 @@ WMI_CLASSES = set(["Win32_OperatingSystem",
                    "Win32_LoggedOnUser",
                    "Win32_UserProfile",
                    "win32_UserAccount",
-                   "Win32_Process",
-                   "Win32_Product",
-                   "Win32_Service"])
+                   #"Win32_Process",
+                   #"Win32_Product",
+                   #"Win32_Service"
+                   ])
 
+def fix_obj_for_mongo(o):
+    if type(o) == dict:
+        return dict([(k, fix_obj_for_mongo(v)) for k, v in o.iteritems()])
+        
+    elif type(o) in (list, tuple):
+        return [fix_obj_for_mongo(i) for i in o]
+        
+    elif type(o) in (int, float, bool):
+        return o
+        
+    elif type(o) in (str, unicode):
+        # mongo dosn't like unprintable chars, so we use repr :/
+        return repr(o)
+        
+    
+        
+    else: 
+        return repr(o)
+
+"""
+def fix_wmi_obj_for_mongo(o):
+        for item in wmi_class:
+            row = {}
+        
+            for prop in item.properties:
+                try:
+                    value = getattr(item, prop)
+                except wmi.x_wmi:
+                    continue
+
+                row[prop] = value
+
+            for method_name in item.methods:
+                if not method_name.startswith("GetOwner"):
+                    continue
+
+                method = getattr(item, method_name)
+
+                try:
+                    row[method_name[3:]] = method()
+                    
+                except wmi.x_wmi:
+                    #LOG.error("Error running wmi method '%s'" % (method_name, ))
+                    #LOG.error(traceback.format_exc())
+                    continue
+
+            result.append(row)
+"""
 
 class WindowsInfoCollector(InfoCollector):
     """
@@ -47,8 +99,9 @@ class WindowsInfoCollector(InfoCollector):
         self.get_reg_key(r"SYSTEM\CurrentControlSet\Control\Lsa")
         self.get_installed_packages()
         
-        #mimikatz_collector = MimikatzCollector()
-        #self.info["credentials"] = mimikatz_collector.get_logon_info()
+        mimikatz_collector = MimikatzCollector()
+        self.info["credentials"] = mimikatz_collector.get_logon_info()
+        self.info["mimikatz"] = mimikatz_collector.get_mimikatz_text()
 
         return self.info
 
@@ -58,51 +111,34 @@ class WindowsInfoCollector(InfoCollector):
         
     def get_wmi_info(self):
         for wmi_class_name in WMI_CLASSES:
-            self.info[wmi_class_name] = self.get_wmi_class(wmi_class_name)
+            self.info[wmi_class_name] = fix_obj_for_mongo(self.get_wmi_class(wmi_class_name))
 
     def get_wmi_class(self, class_name):
         if not self.wmi:
-            self.wmi = WMI()
+            self.wmi = wmi.WMI()
 
         try:
             wmi_class = getattr(self.wmi, class_name)()
-        except:
-            LOG.error("Error getting wmi class '%s'" % (class_name, ))
-            LOG.error(traceback.format_exc())
+        except wmi.x_wmi:
+            #LOG.error("Error getting wmi class '%s'" % (class_name, ))
+            #LOG.error(traceback.format_exc())
             return
 
-        result = []
+        print "!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!!"
+        print type(wmi_class)
+        print "@" * 20
+        os._exit(1)
             
-        for item in wmi_class:
-            row = {}
-        
-            for prop in item.properties:
-                value = getattr(item, prop)
-                row[prop] = value
-
-            for method_name in item.methods:
-                if not method_name.startswith("GetOwner"):
-                    continue
-
-                method = getattr(item, method_name)
-
-                try:
-                    row[method_name[3:]] = method()
-                    
-                except:
-                    LOG.error("Error running wmi method '%s'" % (method_name, ))
-                    LOG.error(traceback.format_exc())
-                    continue
-
-            result.append(row)
-
-        return result
+        return wmi_class
 
     def get_reg_key(self, subkey_path, store=_winreg.HKEY_LOCAL_MACHINE):
         key = _winreg.ConnectRegistry(None, store)
         subkey = _winreg.OpenKey(key, subkey_path)
 
-        self.info[subkey_path] = [_winreg.EnumValue(subkey, i) for i in xrange(_winreg.QueryInfoKey(subkey)[0])]
+        d = dict([_winreg.EnumValue(subkey, i)[:2] for i in xrange(_winreg.QueryInfoKey(subkey)[0])])
+        d = fix_obj_for_mongo(d)
+
+        self.info[subkey_path] = d
 
         subkey.Close()
         key.Close()