From 21cf786d51ad604282bc430ff305c5039a6b09ea Mon Sep 17 00:00:00 2001
From: Oran Nadler <orann>
Date: Tue, 10 Apr 2018 17:51:18 +0300
Subject: [PATCH]  add more wmi classeS

---
 infection_monkey/system_info/windows_info_collector.py | 7 ++++---
 1 file changed, 4 insertions(+), 3 deletions(-)

diff --git a/infection_monkey/system_info/windows_info_collector.py b/infection_monkey/system_info/windows_info_collector.py
index cf821fa45..0b888ec2b 100644
--- a/infection_monkey/system_info/windows_info_collector.py
+++ b/infection_monkey/system_info/windows_info_collector.py
@@ -21,9 +21,10 @@ WMI_CLASSES = set(["Win32_OperatingSystem",
                    "Win32_UserProfile",
                    "Win32_Group",
                    "Win32_GroupUser",
-                   #"Win32_Product",
+                   "Win32_Product",
+                   "Win32_Service",
+                   "Win32_OptionalFeature",
                    #"Win32_Process",
-                   #"Win32_Service"
                    ])
 
 def fix_obj_for_mongo(o):
@@ -131,7 +132,7 @@ class WindowsInfoCollector(InfoCollector):
         key = _winreg.ConnectRegistry(None, store)
         subkey = _winreg.OpenKey(key, subkey_path)
 
-        d = dict([_winreg.EnumValue(subkey, i)[:2] for i in xrange(_winreg.QueryInfoKey(subkey)[1])])
+        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