Use __enter__() and __exit__() for StdoutCapture

This commit is contained in:
Shreya 2021-02-17 18:55:14 +05:30
parent e0ae8381ba
commit 6c9ce028e0
3 changed files with 135 additions and 142 deletions

View File

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

View File

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

View File

@ -3,20 +3,16 @@ import sys
class StdoutCapture:
def __init__(self):
_orig_stdout = None
_new_stdout = None
def capture_stdout_output(self) -> None:
def __enter__(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()
return info
output = self._new_stdout.read()
return output
def _reset_stdout_to_original(self) -> None:
def __exit__(self, _, __, ___) -> None:
sys.stdout = self._orig_stdout