Repository: Gallopsled/pwntools

Issue:
ROP should also use __setattr__ to automatically set register values with setRegisters
Currently, if you have a ROP gadget that has a `pop eax; ret` gadget, you can access `ROP.eax` and get the gadget information.

This should work in the other direction, so that setting e.g. `rop.rax = 0xdeadbeef` adds the correct data to the stack.


Discussion:
Hi,
I am interested in working on this.
The only thing I am unsure of at the moment is how we would handle setting multiple registers "at once" i.e. with the minimum number of rop chains. I see that we are using `_` delimited registers in `__getattr__`:
```
        #
        # Check for a '_'-delimited list of registers
        #
        x86_suffixes = ['ax', 'bx', 'cx', 'dx', 'bp', 'sp', 'di', 'si',
                        'r8', 'r9', '10', '11', '12', '13', '14', '15']

        if all(map(lambda x: x[-2:] in x86_suffixes, attr.split('_'))):
            return self.search(regs=attr.split('_'), order='regs')
```
But what should go on the right hand side of the assignment? A tuple? e.g.
```
rop.rax_rdi_rsi = (0xdead, 0xbeef, 0xcafe)
```
This looks a bit ugly to me, but I'm not sure what else we can do.
Ideally something like this would work, but would require each `__setattr__` call to check previous ones:
```
rop.rax = 0xdead
rop.rdi = 0xbeef
rop.rsi = 0xcafe
rop.chain()  # minimum rop chain for setting all of the above, not three separate chains (if possible)
```

Thanks,
Mariusz
Yes, this is a good option, but I don't think this would be required. Most will actually use it with a single register name.

A pretty pretty option would be to allow something like this:
```py
rop = ROP(context.binary)
rop(rax=0xdead, rdi=0xbeef, rsi=0xcafe)
```
Indeed the `__call__` interface @Arusekk proposes appears quite pleasant. Any other opinions?
I like the __call__ interface, can we make it so that it takes kwargs *or*
a dict?  I know we can just do **dict but not everybody knows that.

On Sun, Sep 27, 2020 at 11:51 PM Mariusz Skoneczko <notifications@github.com>
wrote:

>
>
> Indeed the __call__ interface @Arusekk <https://github.com/Arusekk>
> proposes appears quite pleasant. Any other opinions?
>
>
>
>
> —
> You are receiving this because you are subscribed to this thread.
> Reply to this email directly, view it on GitHub
> <https://github.com/Gallopsled/pwntools/issues/1636#issuecomment-699768035>,
> or unsubscribe
> <https://github.com/notifications/unsubscribe-auth/AAA3IGEJ4YERBGGJSKW24V3SIAI6DANCNFSM4OWJKL2Q>
> .
>
>
> --

*Zach Riggle*

See #1681

Resolving patch:
diff --git a/pwnlib/rop/rop.py b/pwnlib/rop/rop.py
--- a/pwnlib/rop/rop.py
+++ b/pwnlib/rop/rop.py
@@ -523,6 +523,14 @@ class ROP(object):
     >>> r.ret is None
     True
     """
+    BAD_ATTRS = [
+        'trait_names',          # ipython tab-complete
+        'download',             # frequent typo
+        'upload',               # frequent typo
+    ]
+    X86_SUFFIXES = ['ax', 'bx', 'cx', 'dx', 'bp', 'sp', 'di', 'si',
+                    'r8', 'r9', '10', '11', '12', '13', '14', '15']
+
     def __init__(self, elfs, base = None, badchars = b'', **kwargs):
         """
         Arguments:
@@ -661,6 +669,43 @@ def setRegisters(self, registers):
 
         return stack
 
+    def __call__(self, *args, **kwargs):
+        """Set the given register(s)' by constructing a rop chain.
+
+        This is a thin wrapper around :meth:`setRegisters` which
+        actually executes the rop chain.
+
+        You can call this :class:`ROP` instance and provide keyword arguments,
+        or a dictionary.
+
+        Arguments:
+            regs(dict): Mapping of registers to values.
+                        Can instead provide ``kwargs``.
+
+        >>> context.clear(arch='amd64')
+        >>> assembly = 'pop rax; pop rdi; pop rsi; ret; pop rax; ret;'
+        >>> e = ELF.from_assembly(assembly)
+        >>> r = ROP(e)
+        >>> r(rax=0xdead, rdi=0xbeef, rsi=0xcafe)
+        >>> print(r.dump())
+        0x0000:       0x10000000
+        0x0008:           0xdead
+        0x0010:           0xbeef
+        0x0018:           0xcafe
+        >>> r = ROP(e)
+        >>> r({'rax': 0xdead, 'rdi': 0xbeef, 'rsi': 0xcafe})
+        >>> print(r.dump())
+        0x0000:       0x10000000
+        0x0008:           0xdead
+        0x0010:           0xbeef
+        0x0018:           0xcafe
+        """
+        if len(args) == 1 and isinstance(args[0], dict):
+            for value, _ in self.setRegisters(args[0]):
+                self.raw(value)
+        else:
+            self(kwargs)
+
     def resolve(self, resolvable):
         """Resolves a symbol to an address
 
@@ -703,6 +748,7 @@ def generatePadding(self, offset, count):
         """
         Generates padding to be inserted into the ROP stack.
 
+        >>> context.clear(arch='i386')
         >>> rop = ROP([])
         >>> val = rop.generatePadding(5,15)
         >>> cyclic_find(val[:4])
@@ -1061,6 +1107,7 @@ def raw(self, value):
         Arguments:
             data(int/str): The raw value to put onto the rop chain.
 
+        >>> context.clear(arch='i386')
         >>> rop = ROP([])
         >>> rop.raw('AAAAAAAA')
         >>> rop.raw('BBBBBBBB')
@@ -1350,6 +1397,7 @@ def __getattr__(self, attr):
         Also provides a shorthand for ``.call()``:
             ``rop.function(args)`` is equivalent to ``rop.call(function, args)``
 
+        >>> context.clear(arch='i386')
         >>> elf=ELF(which('bash'))
         >>> rop=ROP([elf])
         >>> rop.rdi     == rop.search(regs=['rdi'], order = 'regs')
@@ -1366,14 +1414,9 @@ def __getattr__(self, attr):
         True
         """
         gadget = collections.namedtuple('gadget', ['address', 'details'])
-        bad_attrs = [
-            'trait_names',          # ipython tab-complete
-            'download',             # frequent typo
-            'upload',               # frequent typo
-        ]
 
         if attr in self.__dict__ \
-        or attr in bad_attrs \
+        or attr in self.BAD_ATTRS \
         or attr.startswith('_'):
             raise AttributeError('ROP instance has no attribute %r' % attr)
 
@@ -1417,10 +1460,7 @@ def __getattr__(self, attr):
         #
         # Check for a '_'-delimited list of registers
         #
-        x86_suffixes = ['ax', 'bx', 'cx', 'dx', 'bp', 'sp', 'di', 'si',
-                        'r8', 'r9', '10', '11', '12', '13', '14', '15']
-
-        if all(map(lambda x: x[-2:] in x86_suffixes, attr.split('_'))):
+        if all(map(lambda x: x[-2:] in self.X86_SUFFIXES, attr.split('_'))):
             return self.search(regs=attr.split('_'), order='regs')
 
         #
@@ -1430,3 +1470,42 @@ def call(*args):
             return self.call(attr, args)
 
         return call
+
+    def __setattr__(self, attr, value):
+        """Helper for setting registers.
+
+        This convenience feature allows one to set the values of registers
+        with simple python assignment syntax.
+
+        Warning:
+            Only one register is set at a time (one per rop chain).
+            This may lead to some previously set to registers be overwritten!
+
+        Note:
+            If you would like to set multiple registers in as few rop chains
+            as possible, see :meth:`__call__`.
+
+        >>> context.clear(arch='amd64')
+        >>> assembly = 'pop rax; pop rdi; pop rsi; ret; pop rax; ret;'
+        >>> e = ELF.from_assembly(assembly)
+        >>> r = ROP(e)
+        >>> r.rax = 0xdead
+        >>> r.rdi = 0xbeef
+        >>> r.rsi = 0xcafe
+        >>> print(r.dump())
+        0x0000:       0x10000004 pop rax; ret
+        0x0008:           0xdead
+        0x0010:       0x10000001 pop rdi; pop rsi; ret
+        0x0018:           0xbeef
+        0x0020:      b'iaaajaaa' <pad rsi>
+        0x0028:       0x10000002 pop rsi; ret
+        0x0030:           0xcafe
+        """
+        if attr in self.BAD_ATTRS:
+            raise AttributeError('ROP instance has no attribute %r' % attr)
+
+        if attr[-2:] in self.X86_SUFFIXES:  # handle setting registers
+            self({attr: value})
+
+        # Otherwise, perform usual setting
+        self.__dict__[attr] = value


Test patch: