From 7ab6752d67a919c2b05f7f8c39fb65cfe31e8a1e Mon Sep 17 00:00:00 2001
From: DoronZ <doron88@gmail.com>
Date: Sat, 25 Jul 2020 18:25:36 +0300
Subject: [PATCH] command: function_start: and optional cmd arg

---
 commands.md                                    | 16 ++++++++++++++--
 fa/commands/function_start.py                  | 18 +++++++++++++++++-
 .../test-project-ida/test-ida-context.sig      |  6 ++++++
 tests/test_commands/test_idalink.py            |  1 +
 4 files changed, 38 insertions(+), 3 deletions(-)

diff --git a/commands.md b/commands.md
index 8ccd053..64b6ec3 100644
--- a/commands.md
+++ b/commands.md
@@ -372,18 +372,30 @@ optional arguments:
 ```
 ## function-start
 ```
-usage: function-start [-h]
+usage: function-start [-h] [cmd [cmd ...]]
 
 goto function's start
 
 EXAMPLE:
     0x00000000: push {r4-r7, lr} -> function's prolog
     ...
-    0x000000f0: push {r4-r7, pc} -> function's epilog
+    0x000000f0: pop {r4-r7, pc} -> function's epilog
 
     results = [0xf0]
     -> function-start
     result = [0]
+    
+EXAMPLE 2:
+    0x00000000: push {r4-r7, lr} -> function's prolog
+    ...
+    0x000000f0: pop {r4-r7, pc} -> function's epilog
+
+    results = []
+    -> function-start arm-find-all 'pop {r4-r7, pc}'
+    result = [0]
+
+positional arguments:
+  cmd         command
 
 optional arguments:
   -h, --help  show this help message and exit
diff --git a/fa/commands/function_start.py b/fa/commands/function_start.py
index a170a00..61519dc 100644
--- a/fa/commands/function_start.py
+++ b/fa/commands/function_start.py
@@ -11,11 +11,20 @@
 EXAMPLE:
     0x00000000: push {r4-r7, lr} -> function's prolog
     ...
-    0x000000f0: push {r4-r7, pc} -> function's epilog
+    0x000000f0: pop {r4-r7, pc} -> function's epilog
 
     results = [0xf0]
     -> function-start
     result = [0]
+    
+EXAMPLE 2:
+    0x00000000: push {r4-r7, lr} -> function's prolog
+    ...
+    0x000000f0: pop {r4-r7, pc} -> function's epilog
+
+    results = []
+    -> function-start arm-find-all 'pop {r4-r7, pc}'
+    result = [0]
 '''
 
 
@@ -30,6 +39,7 @@ def get_parser():
     p = utils.ArgumentParserNoExit('function-start',
                                    description=DESCRIPTION,
                                    formatter_class=RawTextHelpFormatter)
+    p.add_argument('cmd', nargs='*', default='', help='command')
     return p
 
 
@@ -43,4 +53,10 @@ def function_start(addresses):
 
 
 def run(segments, args, addresses, interpreter=None, **kwargs):
+    if len(args.cmd) > 0:
+        cmd = args.cmd[0] + ' ' + \
+              ''.join('"{}"'.format(c) for c in args.cmd[1:])
+        addresses = interpreter.find_from_instructions_list(
+            [cmd],
+            addresses=addresses)
     return list(function_start(addresses))
diff --git a/fa/signatures/test-project-ida/test-ida-context.sig b/fa/signatures/test-project-ida/test-ida-context.sig
index 84334cd..85c0dd5 100644
--- a/fa/signatures/test-project-ida/test-ida-context.sig
+++ b/fa/signatures/test-project-ida/test-ida-context.sig
@@ -68,5 +68,11 @@
 
         offset 4
         set-struct-member test_struct_t test_member_offset_4 'unsigned int'
+
+        clear
+
+        function-start arm-find-all 'pop {r4-r7, pc}'
+        single 0
+        set-name test_function_start_cmd
 	]
 }
diff --git a/tests/test_commands/test_idalink.py b/tests/test_commands/test_idalink.py
index 84bbc8a..18b3ae5 100644
--- a/tests/test_commands/test_idalink.py
+++ b/tests/test_commands/test_idalink.py
@@ -109,3 +109,4 @@ def test_ida_symbols(ida, sample_elf):
         assert symbols['test_find_immediate'] == 0x1240
         assert symbols['test_find_immediate'] == 0x1240
         assert symbols['test_operand'] == 1
+        assert symbols['test_function_start_cmd'] == 0x1248