Repository: Gallopsled/pwntools

Issue:
recv_regex but it returns captures
A pattern I do a lot when CTFing is "read a blob of data, apply regex, extract data". It's not exactly a huge issue but its still a few extra lines and I feel like it's a common enough pattern to be worth implementing. I'm envisioning a parameter for the recvregex family that will return a tuple of the captures instead of the matched string but i'm not super picky.  I'll see about implementing it myself and making a PR when I can find some time. 


Discussion:
We already have recvpred or recvline_pred or what it is called for some use cases.

Email z niedzieli, 12 września 2021 od Elmy:
> Hey I think thats a great idea.
> 
> As of know p.recvuntil() only receives until a matched string, but I think it would be a pretty cool and nice feature to have a p.recvregex() that only receives a certain capture group and all the other content prior to the recvregex(). 
> 
> As of know that would take an extra line
> ```py
> # example to receive CTF{.+}
> p.recvuntil(until string before "CTF")
> p.recvuntil(b"}")
> # or we could have
> p.recvregex("CTF{.+}") # stops receiving upon hitting first match
> ```
> 
> But as for this:
> 
> > return a tuple of the captures instead of the matched string
> 
> I'm not sure about receiving tuple of the captures, how do you envision it?
> how would pwntools know how many captures to receive before terminating?
> Are you looking at a p.recvline(), p.recvall() or p.recvuntil()?
> 
> Perhaps we could add an argument to all of the existing `recv`, where it first receives the data and regexes it for capture groups.
> 
> i.e.
> ```py
> p.recvall("CTF{.+}", regex=True)
> ```
> 
> Great idea though and I look forward to its implementation :) 
> 
> -- 
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly or view it on GitHub:
> https://github.com/Gallopsled/pwntools/issues/1962#issuecomment-91756454

-- 
Wysłane z mojego urządzenia Sailfish
I don't think the predicate receive functions are quite what I'm looking for because they still just return the entire matched bytes. I want to flatten the workflow of "receive data, apply regex, extract captures" into "provided a recv function a regex with capture groups and receive the captured data once it matches" and I thought it'd be generally convenient enough to be worth contributing back. 
Thanks for the idea but we're way ahead of you!  Check out `recvregex` and its friends here:

https://docs.pwntools.com/en/latest/tubes.html?highlight=recvregex#pwnlib.tubes.tube.tube.recvregex
I was under the impression that recvregex has essentially the same behavior as recvuntil except it applies a regex instead of a string comparison. What i'm looking for is something like this. 

`The address is 11234123`

And I want something like `r.recvregex("([0-9]+)", capture=True)` which returns `11234123` instead of `The address is 11234123`
Besides is it just me or is the regex slightly broken? `+` and `*` doesnt seem to be working for me 🤔 

```py
>>> from pwn import *
>>> t = tube()
>>> t.recv_raw = lambda n: b"The address is 11234123"
>>> t.recvregex("([0-9]+)")
<stdin>:1: BytesWarning: Text is not bytes; assuming ASCII, no guarantees. See https://docs.pwntools.com/#bytes
b'The address is 1'
```
That looks correct to me; recv operates byte by byte so as soon as you have one number it matches and returns. You can probably solve it by making it match something after the number. In this case that's difficult but if you had something afterwards you could match that. 
I think it would definitely be useful to have something that does group capture as-a-feature.  

I think we have a few options, the easiest is to check the type of the input to be a `re.Pattern` object (we already transform the input into one) and if the original input was one, return the result of the `Pattern` on the data.

Then the user can just use `.groups(...)` on the result.

This realistically only saves like a line or two vs:

```py
pattern = rb'(0x[A-Fa-f0-9]+)'
data = io.recvregex(pattern)
match = re.search(pattern, data)
```

vs

```py
pattern = re.compile(rb'(0x[A-Fa-f0-9]+)')
match = io.recvregex(pattern)
```

Since the default argument for `exact=False`, the predicate used is `re.search` not `re.match` so `match.group()` should be the full match and `match.groups()` would have any user-defined capture groups.


Resolving patch:
diff --git a/pwnlib/tubes/tube.py b/pwnlib/tubes/tube.py
--- a/pwnlib/tubes/tube.py
+++ b/pwnlib/tubes/tube.py
@@ -633,17 +633,30 @@ def recvline_endswith(self, delims, keepends=False, timeout=default):
                                   keepends=keepends,
                                   timeout=timeout)
 
-    def recvregex(self, regex, exact=False, timeout=default):
-        """recvregex(regex, exact=False, timeout=default) -> bytes
+    def recvregex(self, regex, exact=False, timeout=default, capture=False):
+        r"""recvregex(regex, exact=False, timeout=default, capture=False) -> bytes
 
         Wrapper around :func:`recvpred`, which will return when a regex
         matches the string in the buffer.
 
+        Returns all received data up until the regex matched. If `capture` is
+        set to True, a :class:`re.Match` object is returned instead.
+
         By default :func:`re.RegexObject.search` is used, but if `exact` is
         set to True, then :func:`re.RegexObject.match` will be used instead.
 
         If the request is not satisfied before ``timeout`` seconds pass,
         all data is buffered and an empty string (``''``) is returned.
+
+        Examples:
+
+            >>> t = tube()
+            >>> t.recv_raw = lambda n: b'The lucky number is 1337 as always\nBla blubb blargh\n'
+            >>> m = t.recvregex(br'number is ([0-9]+) as always\n', capture=True)
+            >>> m.group(1)
+            b'1337'
+            >>> t.recvregex(br'Bla .* blargh\n')
+            b'Bla blubb blargh\n'
         """
 
         if isinstance(regex, (bytes, bytearray, six.text_type)):
@@ -655,7 +668,10 @@ def recvregex(self, regex, exact=False, timeout=default):
         else:
             pred = regex.search
 
-        return self.recvpred(pred, timeout = timeout)
+        if capture:
+            return pred(self.recvpred(pred, timeout = timeout))
+        else:
+            return self.recvpred(pred, timeout = timeout)
 
     def recvline_regex(self, regex, exact=False, keepends=False, timeout=default):
         """recvline_regex(regex, exact=False, keepends=False, timeout=default) -> bytes


Test patch: