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 @@
+
+
+
+
+
+
+
+
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: