Added tests. Rolled back some code changes.

analyzer/windows/analyzer.py

@@ -233,8 +233,7 @@ def __init__(self):
         self.LASTINJECT_TIME = None
         self.NUM_INJECTED = 0
 
-    @staticmethod
-    def get_pipe_path(name):
+    def get_pipe_path(self, name):
         """Return \\\\.\\PIPE on Windows XP and \\??\\PIPE elsewhere."""
         version = sys.getwindowsversion()
         if version.major == 5 and version.minor == 1:
@@ -351,34 +350,6 @@ def complete(self):
     def get_completion_key(self):
         return getattr(self.config, "completion_key", "")
 
-    def monitor_dcom(self):
-        """Add the 'DcomLaunch' service to critical process list."""
-        if not self.MONITORED_DCOM:
-            self.MONITORED_DCOM = True
-            self.critical_process_by_name("DcomLaunch")
-
-    def monitor_wmi(self):
-        """Add the 'winmgmt' service to critical process list."""
-        if not self.MONITORED_WMI:
-            self.MONITORED_WMI = True
-            self.critical_process_by_name("winmgmt")
-
-    def critical_process_by_name(self, process_name):
-        """Add the named process to the CRITICAL_PROCESS_LIST."""
-        process_id = pid_from_service_name(process_name)
-        self.critical_process(process_id)
-
-    def critical_process(self, process_id, sleep_msec=2000):
-        """Add this process_id to the CRITICAL_PROCESS_LIST."""
-        if process_id:
-            servproc = Process(options=self.options, config=self.config, pid=process_id)
-            self.CRITICAL_PROCESS_LIST.append(int(process_id))
-            filepath = servproc.get_filepath()
-            servproc.inject(interest=filepath, nosleepskip=True)
-            self.LASTINJECT_TIME = timeit.default_timer()
-            servproc.close()
-            KERNEL32.Sleep(sleep_msec)
-
     def run(self):
         """Run analysis.
         @return: operation status.
@@ -975,6 +946,7 @@ def remove_pid(self, pid):
 
 class CommandPipeHandler:
     """Pipe Handler.
+
     This class handles the notifications received through the Pipe Server and
     decides what to do with them.
     """
@@ -1037,16 +1009,22 @@ def _handle_getpids(self, data):
         hidepids.update([self.analyzer.pid, self.analyzer.ppid])
         return struct.pack("%dI" % len(hidepids), *hidepids)
 
-    # remove pid from process list because we received a notification
-    # from kernel land
     def _handle_kterminate(self, data):
+        """Handle terminate notification.
+
+        Remove pid from process list because we received a notification
+        from kernel land
+        """
         process_id = int(data)
         if process_id and process_id in self.analyzer.process_list.pids:
             self.analyzer.process_list.remove_pid(process_id)
 
-    # same than below but we don't want to inject any DLLs because
-    # it's a kernel analysis
     def _handle_kprocess(self, data):
+        """Handle process notification.
+
+        Same than below but we don't want to inject any DLLs because
+        it's a kernel analysis
+        """
         self.analyzer.process_lock.acquire()
         process_id = int(data)
         thread_id = None
@@ -1071,15 +1049,51 @@ def _handle_ksubvert(self, data):
     def _handle_shell(self, data):
         explorer_pid = get_explorer_pid()
         if explorer_pid:
-            self.analyzer.critical_process(explorer_pid)
+            explorer = Process(options=self.analyzer.options, config=self.analyzer.config, pid=explorer_pid)
+            self.analyzer.CRITICAL_PROCESS_LIST.append(int(explorer_pid))
+            filepath = explorer.get_filepath()
+            explorer.inject(interest=filepath, nosleepskip=True)
+            self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+            explorer.close()
+            KERNEL32.Sleep(2000)
 
     def _handle_interop(self, data):
-        self.analyzer.monitor_dcom()
+        if not self.analyzer.MONITORED_DCOM:
+            self.analyzer.MONITORED_DCOM = True
+            dcom_pid = pid_from_service_name("DcomLaunch")
+            if dcom_pid:
+                servproc = Process(options=self.analyzer.options, config=self.analyzer.config, pid=dcom_pid)
+                self.analyzer.CRITICAL_PROCESS_LIST.append(int(dcom_pid))
+                filepath = servproc.get_filepath()
+                servproc.inject(interest=filepath, nosleepskip=True)
+                self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+                servproc.close()
+                KERNEL32.Sleep(2000)
 
     def _handle_wmi(self, data):
-        if not ANALYSIS_TIMED_OUT:
-            self.analyzer.monitor_wmi()
-            self.analyzer.monitor_dcom()
+        if not self.analyzer.MONITORED_WMI and not ANALYSIS_TIMED_OUT:
+            self.analyzer.MONITORED_WMI = True
+            if not self.analyzer.MONITORED_DCOM:
+                self.analyzer.MONITORED_DCOM = True
+                dcom_pid = pid_from_service_name("DcomLaunch")
+                if dcom_pid:
+                    servproc = Process(options=self.analyzer.options, config=self.analyzer.config, pid=dcom_pid)
+                    self.analyzer.CRITICAL_PROCESS_LIST.append(int(dcom_pid))
+                    filepath = servproc.get_filepath()
+                    servproc.inject(interest=filepath, nosleepskip=True)
+                    self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+                    servproc.close()
+                    KERNEL32.Sleep(2000)
+
+            wmi_pid = pid_from_service_name("winmgmt")
+            if wmi_pid:
+                servproc = Process(options=self.analyzer.options, config=self.analyzer.config, pid=wmi_pid)
+                self.analyzer.CRITICAL_PROCESS_LIST.append(int(wmi_pid))
+                filepath = servproc.get_filepath()
+                servproc.inject(interest=filepath, nosleepskip=True)
+                self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+                servproc.close()
+                KERNEL32.Sleep(2000)
 
     def _handle_tasksched(self, data):
         if not self.analyzer.MONITORED_TASKSCHED and not ANALYSIS_TIMED_OUT:
@@ -1098,7 +1112,15 @@ def _handle_tasksched(self, data):
             subprocess.call("net start schedule", startupinfo=si)
             log.info("Started Task Scheduler Service")
 
-            self.analyzer.critical_process_by_name("schedule")
+            sched_pid = pid_from_service_name("schedule")
+            if sched_pid:
+                servproc = Process(options=self.analyzer.options, config=self.analyzer.config, pid=sched_pid)
+                self.analyzer.CRITICAL_PROCESS_LIST.append(int(sched_pid))
+                filepath = servproc.get_filepath()
+                servproc.inject(interest=filepath, nosleepskip=True)
+                self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+                servproc.close()
+                KERNEL32.Sleep(2000)
 
     def _handle_bits(self, data):
         if not self.analyzer.MONITORED_BITS and not ANALYSIS_TIMED_OUT:
@@ -1114,12 +1136,31 @@ def _handle_bits(self, data):
             log.info("Stopped BITS Service")
             subprocess.call("sc config BITS type= own", startupinfo=si)
 
-            self.analyzer.monitor_dcom()
+            if not self.analyzer.MONITORED_DCOM:
+                self.analyzer.MONITORED_DCOM = True
+                dcom_pid = pid_from_service_name("DcomLaunch")
+                if dcom_pid:
+                    servproc = Process(options=self.analyzer.options, config=self.analyzer.config, pid=dcom_pid)
+
+                    self.analyzer.CRITICAL_PROCESS_LIST.append(int(dcom_pid))
+                    filepath = servproc.get_filepath()
+                    servproc.inject(interest=filepath, nosleepskip=True)
+                    self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+                    servproc.close()
+                    KERNEL32.Sleep(2000)
 
             log.info("Starting BITS Service")
             subprocess.call("net start BITS", startupinfo=si)
             log.info("Started BITS Service")
-            self.analyzer.critical_process_by_name("BITS")
+            bits_pid = pid_from_service_name("BITS")
+            if bits_pid:
+                servproc = Process(options=self.analyzer.options, config=self.analyzer.config, pid=bits_pid)
+                self.analyzer.CRITICAL_PROCESS_LIST.append(int(bits_pid))
+                filepath = servproc.get_filepath()
+                servproc.inject(interest=filepath, nosleepskip=True)
+                self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+                servproc.close()
+                KERNEL32.Sleep(2000)
 
     # Handle case of a service being started by a monitored process
     # Switch the service type to own process behind its back so we
@@ -1138,18 +1179,26 @@ def _handle_service(self, servname):
                 # if tasklist previously failed to get the services.exe PID we'll be
                 # unable to inject
                 if self.analyzer.SERVICES_PID:
+                    servproc = Process(options=self.analyzer.options, config=self.analyzer.config, pid=self.analyzer.SERVICES_PID)
+                    self.analyzer.CRITICAL_PROCESS_LIST.append(int(self.analyzer.SERVICES_PID))
+                    filepath = servproc.get_filepath()
+                    servproc.inject(interest=filepath, nosleepskip=True)
+                    self.analyzer.LASTINJECT_TIME = timeit.default_timer()
+                    servproc.close()
+                    KERNEL32.Sleep(1000)
                     self.analyzer.MONITORED_SERVICES = True
-                    self.analyzer.critical_process(self.analyzer.SERVICES_PID, sleep_msec=1000)
                 else:
                     log.error("Unable to monitor service %s", servname)
 
     def _handle_resume(self, data):
         # RESUME:2560,3728'
         self.analyzer.LASTINJECT_TIME = timeit.default_timer()
 
-    # Handle attempted shutdowns/restarts -- flush logs for all monitored processes
-    # additional handling can be added later
     def _handle_shutdown(self, data):
+        """Handle attempted shutdowns/restarts.
+
+        Flush logs for all monitored processes. Additional handling can be added later.
+        """
         log.info("Received shutdown request")
         self.analyzer.process_lock.acquire()
         for process_id in self.analyzer.process_list.pids:
@@ -1161,9 +1210,11 @@ def _handle_shutdown(self, data):
                 self.analyzer.files.dump_files()
         self.analyzer.process_lock.release()
 
-    # Handle case of malware terminating a process -- notify the target
-    # ahead of time so that it can flush its log buffer
     def _handle_kill(self, data):
+        """Handle case of malware terminating a process.
+
+        Notify the target ahead of time so that it can flush its log buffer.
+        """
         self.analyzer.process_lock.acquire()
 
         process_id = int(data)