Add Convenience Utility to readuntil()

readuntil() and readlineuntil() will now automatically bind() a
predicate and given arguments to produce the single function predicate
required.
The 'until' module will provide convenience utilities for use with
readuntil() and readlineuntil().  For now, it contains functools.partial
renamed as bind(), lastline() which can call a predicate with the last
element of the array of lines given from readlineuntil(), and simplified
versions of re.search and re.fullmatch renamed as contains and equals.
These allow us to write powerful and legible statements like:
comm.readlineuntil(lastline,contains,b'Enter')

Signed-off-by: dusoleil <howcansocksbereal@gmail.com>

3 files changed, 20 insertions, 3 deletions

diff --git a/sploit/__init__.py b/sploit/__init__.py
index c7d2c93..c90b980 100644
--- a/sploit/__init__.py
+++ b/sploit/__init__.py
@@ -1 +1 @@
-__all__ = ["log","comm"]
+__all__ = ["log","comm","until"]
diff --git a/sploit/comm.py b/sploit/comm.py
index f7d2f5d..0ee786b 100644
--- a/sploit/comm.py
+++ b/sploit/comm.py
@@ -7,6 +7,7 @@ import select
 import signal
 
 from sploit.log import log
+from sploit.until import bind
 
 class Comm:
     def __init__(self, backend):
@@ -26,8 +27,9 @@ class Comm:
         log(data)
         return data
 
-    def readuntil(self, pred):
+    def readuntil(self, pred, /, *args, **kwargs):
         data = b''
+        pred = bind(pred, *args, **kwargs)
         while(True):
             data += self.back.stdin.read(1)
             if(pred(data)):
@@ -35,8 +37,9 @@ class Comm:
         log(data)
         return data
 
-    def readlineuntil(self, pred):
+    def readlineuntil(self, pred, /, *args, **kwargs):
         dataarr = []
+        pred = bind(pred, *args, **kwargs)
         while(True):
             data = self.back.stdin.readline()
             log(data)
diff --git a/sploit/until.py b/sploit/until.py
new file mode 100644
index 0000000..b4f390f
--- /dev/null
+++ b/sploit/until.py
@@ -0,0 +1,14 @@
+from functools import partial as bind
+import re
+
+def lastline(pred, /, *args, **kwargs):
+    s = args[-1]
+    args = args[:-1]
+    p = bind(pred, *args, **kwargs)
+    return p(s[-1])
+
+def contains(regex, s):
+    return re.search(regex, s)
+
+def equals(regex, s):
+    return re.fullmatch(regex, s)