improve setflags command

pwndbg/commands/__init__.py

@@ -287,7 +287,15 @@ class _ArgparsedCommand(Command):
             self.parser.prog = function.__name__
         else:
             self.parser.prog = command_name
-        self.__doc__ = function.__doc__ = self.parser.description.strip()
+        
+        # TODO/FIXME: Can we also append the generated positional args?
+        # E.g. "-f --flag  This does something"
+        doc = self.parser.description.strip()
+        if self.parser.epilog:
+            doc += '\n' + self.parser.epilog
+
+        self.__doc__ = function.__doc__ = doc
+
         super(_ArgparsedCommand, self).__init__(function, command_name=command_name, *a, **kw)
 
     def split_args(self, argument):

pwndbg/commands/flags.py

@@ -3,37 +3,53 @@ import argparse
 import gdb
 
 import pwndbg.commands
+from argparse import RawTextHelpFormatter
 
-parser = argparse.ArgumentParser(description="Modify register flags",
-                                epilog="setflag ZF 1")
+description="Modify the flags register"
+epilog = """Examples:
+  On X86/X64:
+    setflag ZF 1        -- set zero flag
+    setflag CF 0        -- unset carry flag
+
+  On ARM:
+    setflag Z 0         -- unset the Z cpsr/xpsr flag
+
+  To see flags registers:
+    info reg eflags     -- on x86/x64
+    info reg cspr/xpsr  -- on ARM (specific register may vary)
+
+(This command supports flags registers that are defined for architectures in the pwndbg/regs.py file)
+    """
+
+parser = argparse.ArgumentParser(description=description, epilog=epilog, formatter_class=RawTextHelpFormatter)
 parser.add_argument('flag', type=str,
                     help='Flag for which you want to change the value')
 parser.add_argument('value', type=int,
                     help='Value to which you want to set the flag - only valid options are 0 and 1')
 
-flags = {
-            "CF": 0,
-            "PF": 2,
-            "AF": 4,
-            "ZF": 6,
-            "SF": 7,
-            "TF": 8,
-            "IF": 9,
-            "DF": 10,
-            "OF": 11,
-        }
-
-@pwndbg.commands.ArgparsedCommand(parser, aliases=["flag"])
+@pwndbg.commands.ArgparsedCommand(parser, aliases=["flag"], )
 def setflag(flag, value):
     if value not in [0, 1]:
         print("can only set flag bit to 0 or 1")
         return
 
-    if flag.upper() not in flags.keys():
-        print("%s not a valid flag" % flag)
-        return
+    register_set = pwndbg.regs.arch_to_regs[pwndbg.arch.current]
+
+    flag = flag.upper()
+    for flag_reg, flags in register_set.flags.items():
+        for (flag_name, flag_bit) in flags.items():
+            if flag_name == flag:
+                old_flags_reg_value = pwndbg.regs[flag_reg]
+                bit_value = 1 << flag_bit
+
+                if value == 1:
+                    new_flags_reg_value = old_flags_reg_value | bit_value
+                else:
+                    new_flags_reg_value = old_flags_reg_value & ~bit_value
+
+                setattr(pwndbg.regs, flag_reg, new_flags_reg_value)
+                print("Set flag %s=%d in flag register %s (old val=%#x, new val=%#x)" % (flag, value, flag_reg, old_flags_reg_value, new_flags_reg_value))
+                return
+
+    print("The %s not a valid/recognized flag" % flag)
 
-    if value == 1:
-        gdb.execute("set $eflags |= (1 << %d)" % flags[flag.upper()])
-    elif value == 0:
-        gdb.execute("set $eflags &= ~(1 << %d)" % flags[flag.upper()])