hilda_client: Add module-based retrieval in objc_get_class

hilda/hilda_client.py

@@ -110,8 +110,8 @@ def lsof(self) -> dict:
         Get dictionary of all open FDs
         :return: Mapping between open FDs and their paths
         """
-        with open(os.path.join(Path(__file__).resolve().parent, 'lsof.m'), 'r') as f:
-            result = json.loads(self.po(f.read()))
+        data = (Path(__file__).parent / 'objective_c' / 'lsof.m').read_text()
+        result = json.loads(self.po(data))
         # convert FDs into int
         return {int(k): v for k, v in result.items()}
 
@@ -714,13 +714,18 @@ def lldb_handle_command(self, cmd):
         """
         self.debugger.HandleCommand(cmd)
 
-    def objc_get_class(self, name) -> objective_c_class.Class:
+    def objc_get_class(self, name: str, module_name: Optional[str] = None) -> objective_c_class.Class:
         """
         Get ObjC class object
+        :param module_name:
         :param name:
         :return:
         """
-        return objective_c_class.Class.from_class_name(self, name)
+        if module_name is not None:
+            ret = self.symbol(self._get_module_class_list(module_name)[name]).objc_class
+        else:
+            ret = objective_c_class.Class.from_class_name(self, name)
+        return ret
 
     def CFSTR(self, symbol: int) -> Symbol:
         """ Create CFStringRef object from given string """
@@ -745,7 +750,7 @@ def ns(self, data: CfSerializable) -> Symbol:
         except TypeError as e:
             raise ConvertingToNsObjectError from e
 
-        obj_c_code = (Path(__file__).resolve().parent / 'to_ns_from_json.m').read_text()
+        obj_c_code = (Path(__file__).parent / 'objective_c' / 'to_ns_from_json.m').read_text()
         expression = obj_c_code.replace('__json_object_dump__', json_data.replace('"', r'\"'))
         try:
             return self.evaluate_expression(expression)
@@ -758,8 +763,7 @@ def decode_cf(self, address: Union[int, str]) -> CfSerializable:
         :param address: NS object.
         :return: Python object.
         """
-        with open(os.path.join(Path(__file__).resolve().parent, 'from_ns_to_json.m'), 'r') as f:
-            obj_c_code = f.read()
+        obj_c_code = (Path(__file__).parent / 'objective_c' / 'from_ns_to_json.m').read_text()
         address = f'0x{address:x}' if isinstance(address, int) else address
         expression = obj_c_code.replace('__ns_object_address__', address)
         try:
@@ -1131,3 +1135,15 @@ def _ks(self) -> Optional['Ks']:
                      'arm64e': Ks(KS_ARCH_ARM64, KS_MODE_LITTLE_ENDIAN),
                      'x86_64h': Ks(KS_ARCH_X86, KS_MODE_64)}
         return platforms.get(self.arch)
+
+    def _get_module_class_list(self, module_name: str):
+        for m in self.target.module_iter():
+            if module_name != m.file.basename:
+                continue
+            objc_classlist = m.FindSection('__DATA').FindSubSection('__objc_classlist')
+            objc_classlist_addr = self.symbol(objc_classlist.GetLoadAddress(self.target))
+            obj_c_code = (Path(__file__).parent / 'objective_c' / 'get_objectivec_class_by_module.m').read_text()
+            obj_c_code = obj_c_code.replace('__count_objc_class', f'{objc_classlist.size // 8}').replace(
+                '__objc_class_list',
+                f'{objc_classlist_addr}')
+            return json.loads(self.po(obj_c_code))

hilda/objective_c/get_objectivec_class_by_module.m

@@ -0,0 +1,19 @@
+@import ObjectiveC;
+@import Foundation;
+
+uintptr_t *classList = (uintptr_t *)__objc_class_list;
+size_t count = (int)__count_objc_class;
+NSMutableDictionary *classDictionary = [NSMutableDictionary dictionaryWithCapacity:count];
+
+for (size_t i = 0; i < count; i++) {
+    Class cls = (Class)classList[i];
+    NSString *className = NSStringFromClass(cls);
+
+    if (cls && className) {
+        NSNumber *classAddress = [NSNumber numberWithUnsignedLongLong:(uintptr_t)cls];
+        [classDictionary setObject:classAddress forKey:className];
+    }
+}
+
+NSData *data = [NSJSONSerialization dataWithJSONObject:classDictionary options:0 error:nil];
+[[NSString alloc] initWithData:data encoding:NSUTF8StringEncoding];

hilda/objective_c_class.py

@@ -1,5 +1,4 @@
 import json
-import os
 import time
 from collections import namedtuple
 from dataclasses import dataclass, field
@@ -112,8 +111,7 @@ def from_class_name(client, class_name: str):
         :param hilda.hilda_client.HildaClient client: Hilda client.
         :param class_name: Class name.
         """
-        with open(os.path.join(Path(__file__).resolve().parent, 'get_objectivec_class_description.m'), 'r') as f:
-            obj_c_code = f.read()
+        obj_c_code = (Path(__file__).parent / 'objective_c' / 'get_objectivec_class_description.m').read_text()
         obj_c_code = obj_c_code.replace('__class_address__', '0').replace('__class_name__', class_name)
         class_symbol = Class(client, class_data=json.loads(client.po(obj_c_code)))
         if class_symbol.name != class_name:
@@ -136,8 +134,7 @@ def reload(self):
         Reload class object data.
         Should be used whenever the class layout changes (for example, during method swizzling)
         """
-        with open(os.path.join(Path(__file__).resolve().parent, 'get_objectivec_class_description.m'), 'r') as f:
-            obj_c_code = f.read()
+        obj_c_code = (Path(__file__).parent / 'objective_c' / 'get_objectivec_class_description.m').read_text()
         obj_c_code = obj_c_code.replace('__class_address__', f'{self._class_object:d}')
         obj_c_code = obj_c_code.replace('__class_name__', self.name)
         self._load_class_data(json.loads(self._client.po(obj_c_code)))

hilda/objective_c_symbol.py

@@ -1,5 +1,4 @@
 import json
-import os
 from contextlib import suppress
 from dataclasses import dataclass
 from functools import partial
@@ -61,9 +60,7 @@ def reload(self):
         self.methods.clear()
         self.class_ = None
 
-        with open(os.path.join(Path(__file__).resolve().parent, 'get_objectivec_symbol_data.m'), 'r') as f:
-            obj_c_code = f.read()
-
+        obj_c_code = (Path(__file__).parent / 'objective_c' / 'get_objectivec_symbol_data.m').read_text()
         obj_c_code = obj_c_code.replace('__symbol_address__', f'{self:d}')
         data = json.loads(self._client.po(obj_c_code))