Wrap cmd to return output as a string instead of printing

pybag/dbgeng/callbacks.py

@@ -374,14 +374,26 @@ class DbgEngCallbacks(CoClass):
  comtypes.errorinfo.ISupportErrorInfo,
  comtypes.connectionpoints.IConnectionPointContainer]
 
- def __init__(self, event_handler, outfn):
+ def __init__(self, event_handler, stdout):
  super().__init__()
- self.outfn = outfn
+ self._stdout_orig = stdout
+ self.stdout = stdout
  self._ev = event_handler
 
+ @property
+ def stdout(self):
+ return self._stdout
+
+ @stdout.setter
+ def stdout(self, stdout):
+ self._stdout = stdout
+
+ def reset_stdout(self):
+ self.stdout = self._stdout_orig
+
  def IDebugOutputCallbacks_Output(self, mask, text):
  #print("OutputCallbacks Output called {}".format(mask))
- self.outfn(text.decode('utf-8')) # XXX - can just be decode? no utf-8
+ self.stdout.write(text.decode('utf-8')) # XXX - can just be decode? no utf-8
 
  #
  # The Event callbacks need a 'this' parameter. It is special to comtypes and lets us

pybag/pydbg.py

@@ -1,3 +1,4 @@
+import io
 import itertools
 import queue
 import struct
@@ -33,7 +34,7 @@ def InitComObjects(Dbg):
  Dbg.mod = Modules(Dbg._dataspaces, Dbg._symbols)
  Dbg.events = EventHandler(Dbg)
  Dbg.breakpoints = Breakpoints(Dbg._control)
- Dbg.callbacks = DbgEngCallbacks(Dbg.events, sys.stdout.write)
+ Dbg.callbacks = DbgEngCallbacks(Dbg.events, sys.stdout)
 
  Dbg.events.breakpoint(Dbg.breakpoints)
 
@@ -240,7 +241,12 @@ def wait(self, timeout=DbgEng.WAIT_INFINITE):
 
  def cmd(self, cmdline):
  """cmd(cmdline) -> execute a windbg console command"""
+ buffer = io.StringIO()
+ self.callbacks.stdout = buffer
  self._control.Execute(cmdline)
+ self.callbacks.reset_stdout()
+ buffer.seek(0)
+ return buffer.read()
 
  def bitness(self):
  """bitness() -> Return target bitness"""