From a868e03c92c88b29a4e69bde931d6ae3b40911da Mon Sep 17 00:00:00 2001
From: doronz88 <doron88@gmail.com>
Date: Wed, 11 Sep 2024 15:07:48 +0300
Subject: [PATCH] commands: add union

---
 commands.md                                   | 34 ++++++++++++++--
 fa/commands/union.py                          | 40 +++++++++++++++++++
 fa/fa_types.py                                |  2 +-
 .../sublime/sig.sublime-completions           |  9 ++++-
 4 files changed, 78 insertions(+), 7 deletions(-)
 create mode 100644 fa/commands/union.py

diff --git a/commands.md b/commands.md
index 54b6d5c..0e61a5c 100644
--- a/commands.md
+++ b/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]
diff --git a/fa/commands/union.py b/fa/commands/union.py
new file mode 100644
index 0000000..c984031
--- /dev/null
+++ b/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)
diff --git a/fa/fa_types.py b/fa/fa_types.py
index 1051195..747f3cf 100644
--- a/fa/fa_types.py
+++ b/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:
diff --git a/ide-completions/sublime/sig.sublime-completions b/ide-completions/sublime/sig.sublime-completions
index 99e4b3e..63274bf 100644
--- a/ide-completions/sublime/sig.sublime-completions
+++ b/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",