Use __enter__() and __exit__() for StdoutCapture

monkey/infection_monkey/exploit/zerologon.py

@@ -303,8 +303,7 @@ class ZerologonExploiter(HostExploiter):
 
         remote_shell = wmiexec.get_remote_shell()
         if remote_shell:
-            output_captor = StdoutCapture()
+            with StdoutCapture() as output_captor:
-            output_captor.capture_stdout_output()
                 try:
                     # Save HKLM keys on victim.
                     remote_shell.onecmd('reg save HKLM\\SYSTEM system.save && ' +

monkey/infection_monkey/exploit/zerologon_utils/dump_secrets.py

@@ -98,9 +98,7 @@ class DumpSecrets:
             self.__username, self.__password, self.__domain, self.__lmhash, self.__nthash)
 
     def dump(self):
-        output_captor = StdoutCapture()
+        with StdoutCapture() as output_captor:
-        output_captor.capture_stdout_output()
-
             dumped_secrets = ''
 
             try:

monkey/infection_monkey/utils/capture_output.py

@@ -3,20 +3,16 @@ import sys
 
 
 class StdoutCapture:
-    def __init__(self):
+    def __enter__(self) -> None:
-        _orig_stdout = None
-        _new_stdout = None
-
-    def capture_stdout_output(self) -> None:
         self._orig_stdout = sys.stdout
         self._new_stdout = io.StringIO()
         sys.stdout = self._new_stdout
+        return self
 
     def get_captured_stdout_output(self) -> str:
-        self._reset_stdout_to_original()
         self._new_stdout.seek(0)
-        info = self._new_stdout.read()
+        output = self._new_stdout.read()
-        return info
+        return output
 
-    def _reset_stdout_to_original(self) -> None:
+    def __exit__(self, _, __, ___) -> None:
         sys.stdout = self._orig_stdout