Merge pull request #52 from doronz88/feature/union

Feature/union

commands.md

@@ -49,6 +49,7 @@ Below is the list of available commands:
 - [store](#store)
 - [symdiff](#symdiff)
 - [trace](#trace)
+- [union](#union)
 - [unique](#unique)
 - [verify-aligned](#verify-aligned)
 - [verify-bytes](#verify-bytes)
@@ -181,8 +182,9 @@ options:
 ```
 usage: deref-data [-h] -l LEN
 
-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]
+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]
 
 options:
   -h, --help         show this help message and exit
@@ -522,7 +524,7 @@ options:
 ```
 ## locate
 ```
-usage: locate [-h] name
+usage: locate [-h] name [name ...]
 
 goto symbol by name
 
@@ -635,7 +637,8 @@ options:
 ```
 ## next-instruction
 ```
-usage: next-instruction [-h] [--limit LIMIT] [--back] [--op0 OP0] [--op1 OP1] [--op2 OP2] [--op3 OP3] [--op4 OP4] [--op5 OP5] mnem [mnem ...]
+usage: next-instruction [-h] [--limit LIMIT] [--back] [--op0 OP0] [--op1 OP1] [--op2 OP2] [--op3 OP3] [--op4 OP4] [--op5 OP5]
+                        mnem [mnem ...]
 
 Map the resultset to the next instruction of a given pattern. The instruction is searched for linearly.
 
@@ -930,6 +933,29 @@ sets a pdb breakpoint
 options:
   -h, --help  show this help message and exit
 ```
+## union
+```
+usage: union [-h] [--piped] variables [variables ...]
+
+union two or more variables
+
+EXAMPLE:
+    results = [0, 4, 8]
+    store a
+    ...
+    results = [0, 12, 20]
+    store b
+
+    -> union a b
+    results = [0, 4, 8, 12, 20]
+
+positional arguments:
+  variables    variable names
+
+options:
+  -h, --help   show this help message and exit
+  --piped, -p
+```
 ## unique
 ```
 usage: unique [-h]

fa/commands/locate.py

@@ -1,4 +1,5 @@
 from argparse import RawTextHelpFormatter
+from typing import Iterable, List
 
 from fa import context, utils
 
@@ -25,15 +26,23 @@ def get_parser():
     p = utils.ArgumentParserNoExit('locate',
                                    description=DESCRIPTION,
                                    formatter_class=RawTextHelpFormatter)
-    p.add_argument('name')
+    p.add_argument('name', nargs='+')
     return p
 
 
 @context.ida_context
-def locate(name):
+def locate_single(name) -> int:
     return idc.get_name_ea_simple(name)
 
 
-def run(segments, args, addresses, interpreter=None, **kwargs):
-    address = locate(args.name)
-    return [address] if address != idc.BADADDR else []
+def locate(names: Iterable[str]) -> List[int]:
+    result = []
+    for n in names:
+        located = locate_single(n)
+        if located != idc.BADADDR:
+            result.append(located)
+    return result
+
+
+def run(segments, args, addresses: Iterable[int], interpreter=None, **kwargs) -> List[int]:
+    return locate(args.name)

fa/commands/set_name.py

@@ -1,18 +1,43 @@
+from typing import List
+
+from fa.commands.locate import locate_single
+from fa.fa_types import IDA_MODULE
 from fa.utils import ArgumentParserNoExit
 
+try:
+    import ida_bytes
+    import ida_name
+    from ida_idaapi import BADADDR
+except ImportError:
+    pass
+
 
-def get_parser():
+def get_parser() -> ArgumentParserNoExit:
     p = ArgumentParserNoExit('set-name',
                              description='set symbol name')
     p.add_argument('name')
     return p
 
 
-def set_name(addresses, name, interpreter):
+def is_address_nameless(addr: int) -> bool:
+    return not ida_bytes.f_has_user_name(ida_bytes.get_flags(addr), None)
+
+
+def set_name(addresses: List[int], name: str, interpreter) -> List[int]:
     for ea in addresses:
+        if IDA_MODULE:
+            current_name = ida_name.get_ea_name(ea)
+            remote_addr = locate_single(current_name)
+            if current_name == name:
+                continue
+
+            # we want to avoid accidental renames from bad sigs, therefore we assert the following:
+            assert remote_addr == BADADDR, f'Rename failed, name already used at {hex(remote_addr)} ({hex(ea)})'
+            assert is_address_nameless(ea), f'Rename failed, address has a different name {current_name} ({hex(ea)})'
+
         interpreter.set_symbol(name, ea)
     return addresses
 
 
-def run(segments, args, addresses, interpreter=None, **kwargs):
+def run(segments, args, addresses: List[int], interpreter=None, **kwargs) -> List[int]:
     return set_name(addresses, args.name, interpreter)

fa/commands/union.py

@@ -0,0 +1,40 @@
+from argparse import RawTextHelpFormatter
+from typing import List
+
+from fa import utils
+
+DESCRIPTION = '''union two or more variables
+
+EXAMPLE:
+    results = [0, 4, 8]
+    store a
+    ...
+    results = [0, 12, 20]
+    store b
+
+    -> union a b
+    results = [0, 4, 8, 12, 20]
+'''
+
+
+def get_parser() -> utils.ArgumentParserNoExit:
+    p = utils.ArgumentParserNoExit('union',
+                                   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: List[int], interpreter=None, **kwargs) -> List[int]:
+    if args.piped:
+        first_var = addresses
+    else:
+        first_var = interpreter.get_variable(args.variables.pop(0))
+
+    results = set(first_var)
+
+    for c in args.variables:
+        results.update(interpreter.get_variable(c))
+
+    return list(results)

fa/fa_types.py

@@ -10,7 +10,7 @@
 
     IDA_MODULE = True
 except ImportError:
-    pass
+    IDA_MODULE = False
 
 
 def del_struct_members(sid: int, offset1: int, offset2: int) -> None:

ide-completions/sublime/sig.sublime-completions

@@ -114,7 +114,7 @@
         {
             "trigger": "locate",
             "kind": "snippet",
-            "contents": "locate ${1:name}"
+            "contents": "locate ${1:name} ${2:name} ..."
         },
         {
             "trigger": "make-code",
@@ -159,7 +159,7 @@
         {
             "trigger": "next-instruction",
             "kind": "snippet",
-            "contents": "next-instruction --limit ${1:LIMIT} --back --op0 ${2:OP0} --op1 ${3:OP1} --op2 ${4:OP2} --op3 ${5:OP3} --op4 ${6:OP4} --op5 ${7:OP5} ${8:mnem} ${9:mnem} ..."
+            "contents": "next-instruction --limit ${1:LIMIT} --back --op0 ${2:OP0} --op1 ${3:OP1} --op2 ${4:OP2} --op3 ${5:OP3} --op4 ${6:OP4} --op5 ${7:OP5}                        ${8:mnem} ${9:mnem} ..."
         },
         {
             "trigger": "offset",
@@ -241,6 +241,11 @@
             "kind": "snippet",
             "contents": "trace "
         },
+        {
+            "trigger": "union",
+            "kind": "snippet",
+            "contents": "union --piped ${1:variables} ${2:variables} ..."
+        },
         {
             "trigger": "unique",
             "kind": "snippet",