Merge pull request #50 from doronz88/feature/new-commands

Feature/new commands

fa/commands/deref_data.py

@@ -0,0 +1,44 @@
+from typing import List
+
+from fa import context, utils
+
+try:
+    import idc
+except ImportError:
+    pass
+
+DESCRIPTION = '''Dereference pointer as integer data type.
+
+Note that the data is assumed to be stored in little endian format.
+
+Example #1:
+    0x00000000: LDR R1, [SP, #0x34]
+    
+    results = [0]
+    -> deref-data -l 4
+    results = [0xe5d1034]
+
+Example #2:
+    0x00000000: LDR R1, [SP, #0x34]
+    
+    results = [0]
+    -> deref-data -l 2
+    results = [0x1034]
+'''
+
+
+def get_parser():
+    p = utils.ArgumentParserNoExit('deref-data',
+                                   description=DESCRIPTION)
+    p.add_argument('-l', '--len', type=int, required=True,
+                   help='length of the data in bytes')
+    return p
+
+
+@context.ida_context
+def deref_data(addresses: List[int], data_len: int) -> List[int]:
+    return [int.from_bytes(idc.get_bytes(ea, data_len), 'little') for ea in addresses]
+
+
+def run(segments, args, addresses, interpreter=None, **kwargs):
+    return deref_data(addresses, args.len)

fa/commands/if-not.py

@@ -0,0 +1,39 @@
+from argparse import RawTextHelpFormatter
+
+from fa import utils
+
+DESCRIPTION = '''perform an 'if not' statement to create conditional branches
+using an FA command
+
+EXAMPLE:
+    results = [0, 4, 8]
+
+    -> if-not 'verify-single' a_is_single_label
+
+    set-name a_is_single
+    b end
+
+    label a_is_not_single_label
+    set-name a_is_not_single
+
+    label end
+'''
+
+
+def get_parser():
+    p = utils.ArgumentParserNoExit('if-not',
+                                   description=DESCRIPTION,
+                                   formatter_class=RawTextHelpFormatter)
+    p.add_argument('cond', help='condition as an FA command')
+    p.add_argument('label', help='label to jump to if condition is false')
+    return p
+
+
+def run(segments, args, addresses, interpreter=None, **kwargs):
+    if len(interpreter.find_from_instructions_list([args.cond],
+                                                   addresses=addresses[:])) == 0:
+        interpreter.set_pc(args.label)
+
+        # pc is incremented by 1, after each instruction
+        interpreter.dec_pc()
+    return addresses

fa/commands/if.py

@@ -31,7 +31,7 @@ def get_parser():
 
 def run(segments, args, addresses, interpreter=None, **kwargs):
     if len(interpreter.find_from_instructions_list([args.cond],
-                                                   addresses=addresses)):
+                                                   addresses=addresses[:])):
         interpreter.set_pc(args.label)
 
         # pc is incremented by 1, after each instruction

fa/commands/intersect.py

@@ -1,4 +1,5 @@
 from argparse import RawTextHelpFormatter
+from typing import List
 
 from fa import utils
 
@@ -21,14 +22,19 @@ def get_parser():
                                    description=DESCRIPTION,
                                    formatter_class=RawTextHelpFormatter)
     p.add_argument('variables', nargs='+', help='variable names')
+    p.add_argument('--piped', '-p', action='store_true')
     return p
 
 
-def run(segments, args, addresses, interpreter=None, **kwargs):
-    first_var = args.variables[0]
-    results = set(interpreter.get_variable(first_var))
+def run(segments, args, addresses: List[int], interpreter=None, **kwargs):
+    if args.piped:
+        first_var = addresses
+    else:
+        first_var = interpreter.get_variable(args.variables.pop(0))
 
-    for c in args.variables[1:]:
+    results = set(first_var)
+
+    for c in args.variables:
         results.intersection_update(interpreter.get_variable(c))
 
     return list(results)

fa/commands/next_instruction.py

@@ -0,0 +1,92 @@
+from argparse import ArgumentParser, RawTextHelpFormatter
+from typing import Iterable, List, Optional, Tuple
+from fa import context, utils
+
+try:
+    import idautils
+    import idc
+except ImportError:
+    pass
+
+DESCRIPTION = '''Map the resultset to the next instruction of a given pattern. The instruction is searched for linearly.
+
+Example #1:
+    0x00000000: mov r0, r1
+    0x00000004: mov r1, r2
+    0x00000008: push {r4}
+    0x0000000c: mov r2, r3
+    
+    results = [0, 4, 8]
+    -> next-instruction mov
+    results = [0, 4, 12]
+
+Example #2:
+    0x00000000: mov r0, r1
+    0x00000004: mov r1, r2
+    0x00000008: push {r4}
+    0x0000000c: mov r2, r3
+    
+    results = [0, 4, 8]
+    -> next-instruction mov --op 2
+    results = [12, 12, 12]
+'''
+
+
+def get_parser() -> ArgumentParser:
+    p = utils.ArgumentParserNoExit('next-instruction',
+                                   description=DESCRIPTION,
+                                   formatter_class=RawTextHelpFormatter)
+
+    p.add_argument('mnem', nargs='+')
+    p.add_argument('--limit', type=int, help='Number of instructions to search per address', default=None)
+    p.add_argument('--back', action='store_true', help='Search backwards instead of forwards')
+    utils.add_operand_args(p)
+    return p
+
+
+def _find_next_instruction(mnems: Iterable[str],
+                           regs_description: Iterable[Tuple[int, Iterable[str]]],
+                           address: int,
+                           backwards: bool = False,
+                           limit: Optional[int] = None) -> Optional[int]:
+    instructions = list(idautils.FuncItems(address))
+
+    if backwards:
+        instructions = [ea for ea in instructions if ea <= address][::-1]
+    else:
+        instructions = [ea for ea in instructions if ea >= address]
+
+    if limit is not None:
+        instructions = instructions[:limit]
+
+    for ea in instructions:
+        current_mnem = idc.print_insn_mnem(ea).lower()
+        if current_mnem in mnems:
+            if not regs_description:
+                return ea
+
+            for description in regs_description:
+                index, values = description
+                if not utils.compare_operand(ea, index, values):
+                    break
+            else:
+                return ea
+
+    return None
+
+
+@context.ida_context
+def next_instruction(addresses: List[int],
+                     mnem: str,
+                     regs_description: Iterable[Tuple[int, Iterable[str]]],
+                     backwards: bool = False,
+                     limit: Optional[int] = None) -> List[int]:
+    for address in addresses:
+        r = _find_next_instruction(mnem, regs_description, address, backwards, limit)
+        if r is not None:
+            yield r
+
+
+def run(segments, args, addresses: List[int], interpreter=None, **kwargs):
+    regs_description = utils.create_regs_description_from_args(args)
+    return list(next_instruction(addresses, args.mnem, regs_description, args.address, args.back, args.limit))

fa/commands/store.py

@@ -26,5 +26,5 @@ def get_parser():
 
 
 def run(segments, args, addresses, interpreter=None, **kwargs):
-    interpreter.set_variable(args.name, addresses)
+    interpreter.set_variable(args.name, addresses[:])
     return addresses

fa/commands/verify_opcode.py

@@ -0,0 +1,62 @@
+from argparse import RawTextHelpFormatter
+from typing import Generator, Iterable
+
+from fa import context, utils
+
+try:
+    import idc
+except ImportError:
+    pass
+
+DESCRIPTION = '''reduce the result-set to those matching the given instruction
+
+EXAMPLE #1:
+    0x00000000: mov r0, r1
+    0x00000004: mov r1, r2
+    0x00000008: push {r4}
+
+    results = [0, 2, 4, 6, 8]
+    -> verify-opcode mov
+    results = [0, 4]
+
+EXAMPLE #2:
+    0x00000000: mov r0, r1
+    0x00000004: mov r1, r2
+    0x00000008: push {r4}
+
+    results = [0, 2, 4, 6, 8]
+    -> verify-opcode mov --op1 r2
+    results = [4]
+'''
+
+
+def get_parser():
+    p = utils.ArgumentParserNoExit('verify-opcode',
+                                   description=DESCRIPTION,
+                                   formatter_class=RawTextHelpFormatter)
+    p.add_argument('mnem', nargs='+')
+    utils.add_operand_args(p)
+    return p
+
+
+@context.ida_context
+@utils.yield_unique
+def verify_opcode(addresses: Iterable[int], mnem: str, regs_description) -> Generator[int, None, None]:
+    for ea in addresses:
+        current_mnem = idc.print_insn_mnem(ea).lower()
+        if current_mnem == mnem:
+            if not regs_description:
+                yield ea
+                continue
+
+            for description in regs_description:
+                index, values = description
+                if not utils.compare_operand(ea, index, values):
+                    break
+            else:
+                yield ea
+
+
+def run(segments, args, addresses: Iterable[int], interpreter=None, **kwargs):
+    regs_description = utils.create_create_regs_description(args)
+    return list(verify_opcode(addresses, args.mnem, regs_description))

fa/commands/verify_segment.py

@@ -1,4 +1,7 @@
-from argparse import RawTextHelpFormatter
+import re
+from argparse import ArgumentParser, RawTextHelpFormatter
+from collections.abc import Iterable
+from typing import Generator, List
 
 try:
     import idc
@@ -23,21 +26,32 @@
 
 
 @context.ida_context
-def verify_segment(addresses, segment_name):
+def verify_segment(addresses: Iterable[int], segment_name: str, is_regex: bool = False) -> Generator[int, None, None]:
+    if is_regex:
+        matcher = re.compile(segment_name)
+
+        def match(n) -> bool:
+            return bool(matcher.match(n))
+    else:
+        def match(n) -> bool:
+            return segment_name == n
+
     for ea in addresses:
-        if segment_name == idc.get_segm_name(ea):
+        real_seg_name = idc.get_segm_name(ea)
+        if match(real_seg_name):
             yield ea
 
 
-def get_parser():
+def get_parser() -> ArgumentParser:
     p = utils.ArgumentParserNoExit()
     p.add_argument('name', help='segment name')
+    p.add_argument('--regex', help='interpret name as a regex', action='store_true')
 
     p.prog = 'verify-segment'
     p.description = DESCRIPTION
     p.formatter_class = RawTextHelpFormatter
     return p
 
 
-def run(segments, args, addresses, interpreter=None, **kwargs):
-    return list(verify_segment(addresses, args.name))
+def run(segments, args, addresses: Iterable[int], interpreter=None, **kwargs) -> List[int]:
+    return list(verify_segment(addresses, args.name, args.regex))

fa/fa_types.py

@@ -59,7 +59,7 @@ def update_idb(self):
         sorted(keys)
 
         for k in keys:
-            idc.add_enum_member(id, self._values[k], k, 0xffffffff)
+            idc.add_enum_member(id, self._values[k], k, 0xffffffffffffffff)
 
 
 class FaStruct(FaType):