Merge pull request #23 from doronz88/features/upgrade

Features/upgrade

README.md

@@ -98,9 +98,9 @@ Confused? That's alright :grinning:. [Just look at the examples below](#examples
 
 User may also use his own python script files to perform 
 the search. Just create a new `.py` file in your project 
-directory and implement the `run(**kwargs)` method. Also, the project's
-path is appended to python's `sys.path` so you may import
-your scripts from one another.
+directory and implement the `run(interpreter)` method. 
+Also, the project's path is appended to python's `sys.path` 
+so you may import your scripts from one another.
 
 To view the list of available commands, [view the list below](#available-commands)
 
@@ -258,18 +258,16 @@ To view the list of available commands, [view the list below](#available-command
 
 ```python
 from fa.commands.find_str import find_str
+from fa.commands.set_name import set_name
 from fa import context
 
-def run(**kwargs):
+def run(interpreter):
     # throw an exception if not running within ida context
     context.verify_ida('script-name')
-    interp = kwargs['interpreter']
 
     # locate the global string
-    resultset = list(find_str('hello world', null_terminated=True)
-
-    if len(resultset) == 1:
-        interp.set_symbol('g_hello_world', resultset[0])
+    set_name(find_str('hello world', null_terminated=True),
+             'g_hello_world', interpreter)
 ```
 
 #### Python script to automate SIG files interpreter
@@ -283,17 +281,12 @@ unique
 set-name '{function_name}'
 '''
 
-def run(**kwargs):
-    results = {}
-    interp = kwargs['interpreter']
-
+def run(interpreter):
     for function_name in ['func1', 'func2', 'func3']:
         instructions = TEMPLATE.format(unique_string=function_name, 
                                        function_name=function_name).split('\n')
 
-        results[function_name] = interp.find_from_instructions_list(instructions)
-
-    return results
+        interpreter.find_from_instructions_list(instructions)
 ```
 
 #### Python script to dynamically add structs
@@ -332,8 +325,6 @@ def run(**kwargs):
             # the set_type can receive either a string, FaStruct
             # or FaEnum :-)
             set_type(ea, special_struct_t)
-
-    return {}
 ```
 
 ### Aliases

commands.md

@@ -9,6 +9,7 @@ Below is the list of available commands:
 - [back-to-checkpoint](#back-to-checkpoint)
 - [checkpoint](#checkpoint)
 - [clear](#clear)
+- [find](#find)
 - [find-bytes](#find-bytes)
 - [find-bytes-ida](#find-bytes-ida)
 - [find-immediate](#find-immediate)
@@ -195,6 +196,18 @@ EXAMPLE:
     -> clear
     results = []
 
+optional arguments:
+  -h, --help  show this help message and exit
+```
+## find
+```
+usage: find [-h] name
+
+find another symbol defined in other SIG files
+
+positional arguments:
+  name        symbol name
+
 optional arguments:
   -h, --help  show this help message and exit
 ```

fa/commands/find.py

@@ -0,0 +1,16 @@
+from fa import utils
+
+
+def get_parser():
+    p = utils.ArgumentParserNoExit('find',
+                                   description='find another symbol defined '
+                                               'in other SIG files')
+    p.add_argument('name', help='symbol name')
+    return p
+
+
+def run(segments, args, addresses, interpreter=None, **kwargs):
+    interpreter.find(args.name)
+
+    # return an empty result-set
+    return []

fa/commands/set_name.py

@@ -13,7 +13,11 @@ def get_parser():
     return p
 
 
-def run(segments, args, addresses, interpreter=None, **kwargs):
+def set_name(addresses, name, interpreter):
     for ea in addresses:
-        interpreter.set_symbol(args.name, ea)
+        interpreter.set_symbol(name, ea)
     return addresses
+
+
+def run(segments, args, addresses, interpreter=None, **kwargs):
+    return set_name(addresses, args.name, interpreter)

fa/fa_types.py

@@ -50,33 +50,27 @@ def update_idb(self):
 
 
 class FaStruct(FaType):
-    Field = namedtuple('Field', ['name', 'type', 'size'])
+    Field = namedtuple('Field', ['name', 'type', 'offset'])
 
     def __init__(self, name):
         super(FaStruct, self).__init__(name)
         self._fields = []
-        self._size = 0
 
-    def add_field(self, name, type_, size=0, offset=0):
-        if (offset != 0) and (offset != self._size):
-            self.add_field('padd_{:x}'.format(self._size),
-                           'unsigned char[{}]'.format(offset - self._size),
-                           offset - self._size)
-
-        self._fields.append(self.Field(name, type_, size))
-        self._size += size
+    def add_field(self, name, type_, offset=0xffffffff):
+        self._fields.append(self.Field(name, type_, offset))
 
     def update_idb(self):
         sid = ida_struct.get_struc_id(self._name)
-        if sid != -1:
-            sptr = ida_struct.get_struc(sid)
-            ida_struct.del_struc(sptr)
-
-        sid = ida_struct.add_struc(idc.BADADDR, self._name, 0)
         sptr = ida_struct.get_struc(sid)
 
+        if sid == idc.BADADDR:
+            sid = ida_struct.add_struc(idc.BADADDR, self._name, 0)
+            sptr = ida_struct.get_struc(sid)
+        else:
+            ida_struct.del_struc_members(sptr, 0, 0xffffffff)
+
         for f in self._fields:
-            ida_struct.add_struc_member(sptr, f.name, idc.BADADDR,
+            ida_struct.add_struc_member(sptr, f.name, f.offset,
                                         (idc.FF_BYTE | idc.FF_DATA)
                                         & 0xFFFFFFFF,
                                         None, 1)

fa/fainterp.py

@@ -198,7 +198,7 @@ def log(message):
         :param message:
         :return:
         """
-        for line in message.splitlines():
+        for line in str(message).splitlines():
             print('FA> {}'.format(line))
 
     @abstractmethod
@@ -410,9 +410,12 @@ def get_python_symbols(self, file_name=None):
 
                 if not file_name or file_name == filename:
                     name = os.path.splitext(filename)[0]
-                    filename = os.path.join(project_root, filename)
+                    filename = os.path.join(root, filename)
                     m = FaInterp.get_module(name, filename)
-                    m.run(interpreter=self)
+                    if not hasattr(m, 'run'):
+                        self.log('skipping: {}'.format(filename))
+                    else:
+                        m.run(self)
 
     def get_json_signatures(self, symbol_name=None):
         """

fa/signatures/test-project-elf/test-basic.sig

@@ -87,5 +87,9 @@
 
         find-str '3DUfw'
         set-name test_find_str
+
+        clear
+
+        find test_find
 	]
 }

fa/signatures/test-project-elf/test_find.sig

@@ -0,0 +1,8 @@
+{
+    "type": "function",
+	"name": "test_find",
+	"instructions": [
+	    add 76
+	    set-name test_find
+	]
+}

fa/signatures/test-project-ida/add_structs.py

@@ -1,7 +1,7 @@
 from fa import fa_types
 
 
-def run(**kwargs):
+def run(interpreter):
     fa_types.add_const('CONST7', 7)
     fa_types.add_const('CONST8', 8)
 
@@ -11,6 +11,9 @@ def run(**kwargs):
     foo_e.update_idb()
 
     special_struct_t = fa_types.FaStruct('special_struct_t')
-    special_struct_t.add_field('member1', 'const char *', size=4)
-    special_struct_t.add_field('member2', 'const char *', size=4, offset=0x20)
+    special_struct_t.add_field('member1', 'const char *')
+    special_struct_t.add_field('member2', 'const char *', offset=0x20)
+    special_struct_t.add_field('member3', 'char', offset=0x60)
+    special_struct_t.add_field('member4', 'char', offset=0x61)
+    special_struct_t.add_field('member5', 'const char *', offset=0x80)
     special_struct_t.update_idb()

fa/signatures/test-project-ida/explore.py

@@ -4,6 +4,5 @@
 '''
 
 
-def run(**kwargs):
-    interp = kwargs['interpreter']
-    interp.find_from_instructions_list(TEMPLATE.splitlines())
+def run(interpreter):
+    interpreter.find_from_instructions_list(TEMPLATE.splitlines())

fa/signatures/test-project-ida/test-basic.sig

@@ -87,5 +87,9 @@
 
         find-str '3DUfw'
         set-name test_find_str
+
+        clear
+
+        find test_find
 	]
 }

fa/signatures/test-project-ida/test_find.sig

@@ -0,0 +1,8 @@
+{
+    "type": "function",
+	"name": "test_find",
+	"instructions": [
+	    add 76
+	    set-name test_find
+	]
+}

setup.py

@@ -2,7 +2,7 @@
 
 setup(
     name='fa',
-    version='0.1.3',
+    version='0.1.4',
     description='FA Plugin',
     author='DoronZ',
     author_email='[email protected]',

tests/test_commands/test_elf.py

@@ -30,3 +30,4 @@ def test_elf_symbols(sample_elf):
     assert symbols['test_append'] == 2
     assert symbols['test_find_bytes'] == 0x1240
     assert symbols['test_find_str'] == 0x1242
+    assert symbols['test_find'] == 76

tests/test_commands/test_idalink.py

@@ -87,6 +87,7 @@ def test_ida_symbols(ida, sample_elf):
         assert symbols['test_append'] == 2
         assert symbols['test_find_bytes'] == 0x1240
         assert symbols['test_find_str'] == 0x1242
+        assert symbols['test_find'] == 76
 
         # from test-ida-context
         assert symbols['test_find_bytes_ida'] == 0x1240