From 012d1253af790d07859720a3cb9b599089a94d2e Mon Sep 17 00:00:00 2001
From: Dobin <dobin.rutishauser@gmail.com>
Date: Wed, 8 May 2024 12:19:57 +0100
Subject: [PATCH] feature: select if add missing IAT entries

---
 app/templates/project.html | 11 +++++++++++
 app/views_project.py       |  3 +++
 pe/superpe.py              | 11 +++++++++--
 supermega.py               |  7 ++++---
 4 files changed, 27 insertions(+), 5 deletions(-)

diff --git a/app/templates/project.html b/app/templates/project.html
index 94fdccd..aeb4aef 100644
--- a/app/templates/project.html
+++ b/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">
diff --git a/app/views_project.py b/app/views_project.py
index a77fc97..24c56fe 100644
--- a/app/views_project.py
+++ b/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]
 
diff --git a/pe/superpe.py b/pe/superpe.py
index e8ca603..226ca0e 100644
--- a/pe/superpe.py
+++ b/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:
diff --git a/supermega.py b/supermega.py
index 5c1a141..92fb2c1 100644
--- a/supermega.py
+++ b/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: