Refactor and fix config parsing and utility functions

Improved type checks in rat_config_parser.py by using isinstance, fixed a bug in config length validation, and enhanced key reporting logic. Corrected parameter names and error messages in data_utils.py utility functions. Fixed logging and key extraction logic in config_decryptor_rijndael.py. Improved error handling and logging in config_decryptor_aes_with_iv_pbkdf.py and config_decryptor_random_hardcoded.py. Optimized CustomAttribute lookup in dotnetpe_payload.py for better performance and accuracy.

Author: doomedraven

src/rat_king_parser/config_parser/rat_config_parser.py

@@ -81,22 +81,21 @@ def __init__(
             if data is None and not isfile(file_path):
                 raise ConfigParserException("File not found")
             # Filled in _decrypt_and_decode_config()
-            self._incompatible_decryptors: list[int] = []
-            try:
-                self._dnpp = DotNetPEPayload(file_path, yara_rule, data)
-            except Exception as e:
-                raise e
+            self._incompatible_decryptors: list[Any] = []
+            self._dnpp = DotNetPEPayload(file_path, yara_rule, data)
             self.report["sha256"] = self._dnpp.sha256
             self.report["yara_possible_family"] = self._dnpp.yara_match
 
             # Assigned in _decrypt_and_decode_config()
             self._decryptor: ConfigDecryptor = None
             self.report["config"] = self._get_config()
-            self.report["key"] = (
-                self._decryptor.key.hex()
-                if self._decryptor is not None and self._decryptor.key is not None
-                else "None"
-            )
+            key_hex = "None"
+            if self._decryptor is not None and self._decryptor.key is not None:
+                if isinstance(self._decryptor.key, bytes):
+                    key_hex = self._decryptor.key.hex()
+                else:
+                    key_hex = str(self._decryptor.key)
+            self.report["key"] = key_hex
             self.report["salt"] = (
                 self._decryptor.salt.hex()
                 if self._decryptor is not None and self._decryptor.salt is not None
@@ -122,7 +121,7 @@ def _decrypt_and_decode_config(
                 config_fields_map[k] = field_name
                 item_data[field_name] = v
             if len(item_data) > 0:
-                if type(item) is config_item.EncryptedStringConfigItem:
+                if isinstance(item, config_item.EncryptedStringConfigItem):
                     # Translate config value RVAs to string values
                     for k in item_data:
                         item_data[k] = self._dnpp.user_string_from_rva(item_data[k])
@@ -159,7 +158,7 @@ def _decrypt_and_decode_config(
                     if self._decryptor is None:
                         raise ConfigParserException("All decryptors failed")
 
-                elif type(item) is config_item.ByteArrayConfigItem:
+                elif isinstance(item, config_item.ByteArrayConfigItem):
                     for k in item_data:
                         arr_size, arr_rva = item_data[k]
                         item_data[k] = self._dnpp.byte_array_from_size_and_rva(
@@ -168,7 +167,7 @@ def _decrypt_and_decode_config(
 
                 decoded_config.update(item_data)
         # UrlHost is a marker of a special case until this can be standardized
-        if len(decoded_config) < min_config_len and "UrlHost" not in item_data:
+        if len(decoded_config) < min_config_len and "UrlHost" not in decoded_config:
             raise ConfigParserException(
                 f"Minimum threshold of config items not met: {len(decoded_config)}/{min_config_len}"
             )

src/rat_king_parser/config_parser/utils/data_utils.py

@@ -31,11 +31,11 @@
 
 
 # Converts a bytes object to an int object using the specified byte order
-def bytes_to_int(bytes: bytes, order: str = "little") -> int:
+def bytes_to_int(data: bytes, order: str = "little") -> int:
     try:
-        return int.from_bytes(bytes, byteorder=order)
+        return int.from_bytes(data, byteorder=order)
     except Exception:
-        raise ConfigParserException(f"Error parsing int from value: {bytes}")
+        raise ConfigParserException(f"Error parsing int from value: {data}")
 
 
 # Decodes a bytes object to a Unicode string, using UTF-16LE for byte values
@@ -57,8 +57,8 @@ def decode_bytes(byte_str: bytes | str) -> str:
 
 
 # Converts an int to a bytes object, with the specified length and order
-def int_to_bytes(int: int, length: int = 4, order: str = "little") -> bytes:
+def int_to_bytes(value: int, length: int = 4, order: str = "little") -> bytes:
     try:
-        return int.to_bytes(length, order)
+        return value.to_bytes(length, order)
     except Exception:
-        raise ConfigParserException(f"Error parsing bytes from value: {int}")
+        raise ConfigParserException(f"Error parsing bytes from value: {value}")

src/rat_king_parser/config_parser/utils/decryptors/config_decryptor_aes_with_iv.py

@@ -173,7 +173,13 @@ def decrypt_encrypted_strings(
         return decrypted_config_strings
 
     # Extracts AES key candidates from the payload
+<<<<<<< Updated upstream
     def _get_aes_key_candidates(self, encrypted_strings: dict[str, str]) -> list[bytes]:
+=======
+    def _get_aes_key_candidates(
+        self, encrypted_strings: dict[str, str]
+    ) -> list[bytes]:
+>>>>>>> Stashed changes
         logger.debug("Extracting AES key candidates...")
         keys = []
 

src/rat_king_parser/config_parser/utils/decryptors/config_decryptor_aes_with_iv_pbkdf.py

@@ -67,7 +67,6 @@ def _decrypt(self, iv: bytes, ciphertext: bytes) -> bytes:
         )
 
         cipher = AES.new(self.key, mode=CBC, iv=iv)
-        unpadded_text = ""
 
         try:
             unpadded_text = cipher.decrypt(ciphertext)
@@ -113,7 +112,7 @@ def decrypt_encrypted_strings(
                 result = decode_bytes(self._decrypt(self.iv, decoded_val))
             except ConfigParserException as e:
                 last_exc = e
-                print("error", e)
+                logger.debug(f"Error decrypting {k}: {e}")
 
             if result is None:
                 logger.debug(

src/rat_king_parser/config_parser/utils/decryptors/config_decryptor_random_hardcoded.py

@@ -79,9 +79,9 @@ def _get_hardcoded_hosts(self) -> list[str]:
         hardcoded_hosts = []
         for rva in hardcoded_host_rvas:
             try:
-                harcoded_host = self._payload.user_string_from_rva(bytes_to_int(rva))
-                if harcoded_host != ".":
-                    hardcoded_hosts.append(harcoded_host)
+                hardcoded_host = self._payload.user_string_from_rva(bytes_to_int(rva))
+                if hardcoded_host != ".":
+                    hardcoded_hosts.append(hardcoded_host)
             except Exception as e:
                 logger.error(f"Error translating hardcoded host at {hex(rva)}: {e}")
                 continue

src/rat_king_parser/config_parser/utils/decryptors/config_decryptor_rijndael.py

@@ -29,7 +29,7 @@
 # LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
 # SOFTWARE.
-import logging
+from logging import getLogger
 from base64 import b64decode
 from hashlib import md5
 from re import DOTALL, compile, search
@@ -42,7 +42,7 @@
 from ..dotnetpe_payload import DotNetPEPayload
 from .config_decryptor import ConfigDecryptor, IncompatibleDecryptorException
 
-logger = logging.getLogger(__name__)
+logger = getLogger(__name__)
 
 
 # Is old AES - specifically Rijndael in CBC mode with MD5 hashing for key derivation
@@ -110,11 +110,14 @@ def decrypt_encrypted_strings(self, encrypted_strings: dict[str, str]) -> dict[s
                 else:
                     for key_pattern in self._KEY_PATTERNS:
                         key_hit = search(key_pattern, self._payload.data)
+                        if not key_hit:
+                            continue
                         key_rva = bytes_to_int(key_hit.groups()[0])
                         raw_key_field = self._payload.field_name_from_rva(key_rva)
-                        key = encrypted_strings[raw_key_field]
-                        self.key = self._derive_key(key)
-                        break
+                        if raw_key_field in encrypted_strings:
+                            key = encrypted_strings[raw_key_field]
+                            self.key = self._derive_key(key)
+                            break
             except Exception as e:
                 raise ConfigParserException(f"Failed to derive AES key: {e}")
 

src/rat_king_parser/config_parser/utils/dotnetpe_payload.py

@@ -234,6 +234,13 @@ def custom_attribute_from_type(self, typespacename: str, typename: str) -> dict:
         """
         config = {}
         try:
+            ca_map = {}
+            for ca in self.dotnetpe.net.mdtables.CustomAttribute.rows:
+                idx = ca.Parent.row_index
+                if idx not in ca_map:
+                    ca_map[idx] = []
+                ca_map[idx].append(ca)
+
             for td in self.dotnetpe.net.mdtables.TypeDef.rows:
                 if (
                     td.TypeNamespace.value != typespacename
@@ -244,19 +251,21 @@ def custom_attribute_from_type(self, typespacename: str, typename: str) -> dict:
                     self.dotnetpe.net.mdtables.Property.rows
                 ):
                     if pd.Name.value.startswith(
-                        "Boolean_",
-                        "BorderStyle_",
-                        "Color_",
-                        "Byte",
-                        "Int32_",
-                        "SizeF_",
-                        "String_",
+                        (
+                            "Boolean_",
+                            "BorderStyle_",
+                            "Color_",
+                            "Byte",
+                            "Int32_",
+                            "SizeF_",
+                            "String_",
+                        )
                     ):
                         continue
-                    for ca in self.dotnetpe.net.mdtables.CustomAttribute.rows:
-                        if (
-                            ca.Parent.row_index == pd_row_index + 1
-                        ):  # CustomAttribute Parent index is 1-based
+                    # CustomAttribute Parent index is 1-based
+                    target_index = pd_row_index + 1
+                    if target_index in ca_map:
+                        for ca in ca_map[target_index]:
                             # Extract the value from the CustomAttribute blob
                             blob_data = ca.Value.value
                             if blob_data and blob_data != b"\x01\x00\x00\x00":