Updating Process.stdout to allow expecting things

docs/overview/quickstart.rst

@@ -30,6 +30,7 @@ Just Show Me
 
     # Interact with the process
     p.stdout(12) # Read 12 bytes of stdout
+    p.stdout("something") # Read until "something" is in output
     p.stderr(12) # Read 12 bytes of stderr
     p.stdin(b"hello!\n") # Write to stdin
     p.interactive() # Quasi-interactive shell. ctrl-c to exit.
@@ -42,7 +43,7 @@ A Little Deeper
 The two cent starting guide is that everything in ``revenge`` hangs off the
 core class called ``Process``.
 
-:meth:`revenge.Process`
+:class:`revenge.process.Process`
 
 This has traditionally been the starting point for opening applications,
 however in some cases (Android for the moment) it has become necessary to add a

revenge/engines/frida/process.py

@@ -101,20 +101,29 @@ def stderr(self, n):
         self.__stderr = self.__stderr[n:]
         return ret
 
-    @common.validate_argument_types(n=(int, str))
-    def stdout(self, n):
-
-        if isinstance(n, str):
-            if n.lower() == "all":
-                ret = self.__stdout
-                self.__stdout = b""
-                return ret
-            else:
-                raise RevengeInvalidArgumentType("Only valid string is 'all'")
+    @common.validate_argument_types(n=(int, str, bytes))
+    def stdout(self, n=0):
 
-        ret = self.__stdout[:n]
-        self.__stdout = self.__stdout[n:]
-        return ret
+        if n == 0:
+            ret = self.__stdout
+            self.__stdout = b""
+            return ret
+
+        # String acts as an expect
+        if isinstance(n, (str, bytes)):
+            n = common.auto_bytes(n)
+            # TODO: Might be more efficient to use try/except...
+            while n not in self.__stdout: sleep(0.01)
+            index = self.__stdout.index(n) + len(n)
+            ret = self.__stdout[:index]
+            self.__stdout = self.__stdout[index:]
+            return ret
+
+        else:
+            # n is an int. take that much
+            ret = self.__stdout[:n]
+            self.__stdout = self.__stdout[n:]
+            return ret
 
     def interactive(self):
         old_stdout_echo = self._stdout_echo
@@ -125,7 +134,8 @@ def interactive(self):
         self._stderr_echo = True
 
         # Flush out stdout buffer
-        print(self.stdout('all').decode('utf-8'), end="", flush=True)
+        #print(self.stdout('all').decode('utf-8'), end="", flush=True)
+        print(self.stdout().decode('utf-8'), end="", flush=True)
 
         # TODO: Maybe change this to single char get and send at some point?
         while True:
@@ -138,6 +148,7 @@ def interactive(self):
         self._stdout_echo = old_stdout_echo
         self._stderr_echo = old_stderr_echo
 
+from time import sleep
 import prompt_toolkit
 from ...exceptions import *
 from ...native_exception import NativeException

revenge/process.py

@@ -29,9 +29,9 @@
 class Process(object):
 
     def __init__(self, target, resume=False, verbose=False, load_symbols=None,
-            envp=None, engine=None):
+                 envp=None, engine=None):
         """Represents a process.
-
+        
         Args:
             target (str, int, list): File name or pid to attach to. If target
                 is a list, it will be set as argv.
@@ -64,6 +64,9 @@ def __init__(self, target, resume=False, verbose=False, load_symbols=None,
                 # Read from stdout
                 p.stdout(16)
 
+                # Read up to expected output in stdout
+                p.stdout("expected")
+
                 # Interact like a shell
                 p.interactive()
         """
@@ -173,10 +176,15 @@ def target_type(self, x):
 
     @common.implement_in_engine()
     def stdout(self, n):
-        """bytes: Read n bytes from stdout.
+        """Read n bytes from stdout.
         
         Args:
-            n (int, str): Number of bytes to read. Or 'all' to read everything
+            n (int, str, bytes): Number of bytes to read or string to expect.
+                If no value is given, it's presumed you are trying to read 
+                all currently queued output.
+
+        Returns:
+            bytes: Output of stdout
         """
         pass
 

tests/linux/engines/frida/test_frida_process.py

@@ -290,7 +290,9 @@ def stderr_expect(process, thing):
 def test_frida_process_stdio(capsys):
     process = revenge.Process("/bin/cat", verbose=False, resume=True)
     process.stdin("hello world\n")
-    assert b"hello world" in stdout_expect(process, b"world")
+    #assert b"hello world" in stdout_expect(process, b"world")
+    assert b"hello" == process.stdout(b"lo")
+    assert b" world" == process.stdout("ld")
     process.quit()
 
     process = revenge.Process("/bin/cat", verbose=False, resume=False)