feature: select if add missing IAT entries

app/templates/project.html

@@ -137,7 +137,18 @@
                 </select>
             </div>
 
+            <!-- Row 5 -->
+            <div class="col-2">
+                <div class="form-check">
+                    <input class="form-check-input" type="checkbox" value="YES" id="flexCheckDefault"  
+                        name="fix_missing_iat" onchange="this.form.submit()" {{ 'checked' if fix_missing_iat }}>
+                    <label class="form-check-label" for="flexCheckDefault">
+                        Add missing IAT entries
+                    </label>
+                </div>
+            </div>
         </div>
+
         </form>
 
         <div class="row">

app/views_project.py

@@ -123,6 +123,7 @@ def project(name):
         has_rodata_section=has_rodata_section,
 
         has_remote=has_remote,
+        fix_missing_iat=project.settings.fix_missing_iat,
     )
 
 
@@ -170,6 +171,8 @@ def add_project():
             settings.inject_exe_in = request.form['exe']
             settings.inject_exe_out = request.form['exe'].replace(".exe", ".infected.exe")
 
+            settings.fix_missing_iat = True if request.form.get('fix_missing_iat') != None else False
+
             source_style = request.form['source_style']
             settings.source_style = FunctionInvokeStyle[source_style]
 

pe/superpe.py

@@ -2,6 +2,7 @@ import pefile
 import capstone
 import logging
 from typing import List, Dict
+import random
 
 from model.defs import *
 from model.rangemanager import RangeManager
@@ -350,10 +351,16 @@ class SuperPe():
         if not dll_name in iat:
             raise Exception("DLL not found in IAT")
 
+        possible = []
         for entry in iat[dll_name]:
             if len(entry.func_name) >= len(func_name):
-                return entry.func_name
-        return None
+                possible.append(entry.func_name)
+
+        if len(possible) == 0:
+            return None
+        else:
+            # Hope there wont be many collisions
+            return random.choice(possible)
 
 
     def get_iat_offset_by_name(self, dll_name: str, func_name: str) -> int:

supermega.py

@@ -160,9 +160,10 @@ def start_real(settings: Settings):
     if settings.source_style == FunctionInvokeStyle.iat_reuse:
         functions = project.carrier.get_unresolved_iat()
         if len(functions) != 0:
-            #raise Exception("IAT entry not found: {}".format(", ".join(functions)))
-            logger.warn("IAT entry not found: {}".format(", ".join(functions)))
-            pass
+            if settings.fix_missing_iat:
+                logger.info("Fixing missing IAT entries: {}".format(", ".join(functions)))
+            else:
+                raise Exception("IAT entry not found: {}".format(", ".join(functions)))
 
     # Assemble: Assemble .asm to .shc (ASM -> SHC)
     if settings.generate_shc_from_asm: