From 7cdad9a5b7c40617b027ed98491b4b05d369c4db Mon Sep 17 00:00:00 2001
From: Fish <ltfish@users.noreply.github.com>
Date: Sat, 27 Jul 2024 09:43:52 -0700
Subject: [PATCH] Backend: Save load args for future uses. (#499)

* Backend: Save load args for future uses.

* Call set_load_args() in more backends.
---
 cle/backends/backend.py  | 14 ++++++++++++++
 cle/backends/blob.py     |  1 +
 cle/backends/cartfile.py |  1 +
 cle/backends/elf/elf.py  |  6 ++++++
 cle/backends/ihex.py     |  1 +
 cle/backends/pe/pe.py    |  1 +
 cle/backends/srec.py     |  1 +
 7 files changed, 25 insertions(+)

diff --git a/cle/backends/backend.py b/cle/backends/backend.py
index 5930b315..949a2a79 100644
--- a/cle/backends/backend.py
+++ b/cle/backends/backend.py
@@ -144,6 +144,17 @@ def __init__(
         :param binary_stream:   The open stream to this binary. The reference to this will be held until you call close.
         :param is_main_bin:     Whether this binary should be loaded as the main executable
         """
+
+        self.load_args: dict[str, Any] = {} | kwargs
+        self.set_load_args(
+            loader=loader,
+            is_main_bin=is_main_bin,
+            entry_point=entry_point,
+            arch=arch,
+            base_addr=base_addr,
+            force_rebase=force_rebase,
+            has_memory=has_memory,
+        )
         self.binary = binary
         self._binary_stream: BufferedReader = binary_stream
         if self.binary is not None:
@@ -292,6 +303,9 @@ def set_arch(self, arch):
         self._arch = arch
         self.memory = Clemory(arch)  # Private virtual address space, without relocations
 
+    def set_load_args(self, **kwargs) -> None:
+        self.load_args |= kwargs
+
     @property
     def image_base_delta(self):
         return self.mapped_base - self.linked_base
diff --git a/cle/backends/blob.py b/cle/backends/blob.py
index 778a7a0b..8a2103ab 100644
--- a/cle/backends/blob.py
+++ b/cle/backends/blob.py
@@ -30,6 +30,7 @@ def __init__(self, *args, offset=None, segments=None, **kwargs):
             offset = kwargs.pop("custom_offset")
             log.critical("Deprecation warning: the custom_offset parameter has been renamed to offset")
         super().__init__(*args, **kwargs)
+        self.set_load_args(offset=offset, segments=segments)
 
         if self._arch is None:
             raise CLEError("Must specify arch when loading blob!")
diff --git a/cle/backends/cartfile.py b/cle/backends/cartfile.py
index e9f443fe..62070cfb 100644
--- a/cle/backends/cartfile.py
+++ b/cle/backends/cartfile.py
@@ -35,6 +35,7 @@ def __init__(self, binary, binary_stream, *args, arc4_key=None, **kwargs):
                 "Please install the cart Python package before loading a CART file. You may run `pip install cart`."
             )
         super().__init__(binary, binary_stream, *args, **kwargs)
+        self.set_load_args(arc4_key=arc4_key)
 
         # hack: we are using a loader internal method in a non-kosher way which will cause our children to be
         # marked as the main binary if we are also the main binary
diff --git a/cle/backends/elf/elf.py b/cle/backends/elf/elf.py
index 8cf8abbe..15691b5c 100644
--- a/cle/backends/elf/elf.py
+++ b/cle/backends/elf/elf.py
@@ -81,6 +81,12 @@ def __init__(
         **kwargs,
     ):
         super().__init__(*args, **kwargs)
+        self.set_load_args(
+            addend=addend,
+            debug_symbols=debug_symbols,
+            discard_section_headers=discard_section_headers,
+            discard_program_headers=discard_program_headers,
+        )
         patch_undo = []
         try:
             self._reader = elffile.ELFFile(self._binary_stream)
diff --git a/cle/backends/ihex.py b/cle/backends/ihex.py
index 6efa4795..a7d8c840 100644
--- a/cle/backends/ihex.py
+++ b/cle/backends/ihex.py
@@ -80,6 +80,7 @@ def coalesce_regions(regions):
 
     def __init__(self, *args, ignore_missing_arch: bool = False, **kwargs):
         super().__init__(*args, **kwargs)
+        self.set_load_args(ignore_missing_arch=ignore_missing_arch)
 
         if self._arch is None:
             if ignore_missing_arch:
diff --git a/cle/backends/pe/pe.py b/cle/backends/pe/pe.py
index 0c4f97a4..bb71c349 100644
--- a/cle/backends/pe/pe.py
+++ b/cle/backends/pe/pe.py
@@ -36,6 +36,7 @@ class PE(Backend):
 
     def __init__(self, *args, debug_symbols=None, **kwargs):
         super().__init__(*args, **kwargs)
+        self.set_load_args(debug_symbols=debug_symbols)
         self.segments = self.sections  # in a PE, sections and segments have the same meaning
         self.os = "windows"
         if self.binary is None:
diff --git a/cle/backends/srec.py b/cle/backends/srec.py
index 789a275e..0950da21 100644
--- a/cle/backends/srec.py
+++ b/cle/backends/srec.py
@@ -86,6 +86,7 @@ def coalesce_regions(regions):
 
     def __init__(self, *args, ignore_missing_arch: bool = False, **kwargs):
         super().__init__(*args, **kwargs)
+        self.set_load_args(ignore_missing_arch=ignore_missing_arch)
 
         if self._arch is None:
             if ignore_missing_arch: