amfi_detect.py: Adjust to be more Object-oriented

resources/amfi_detect.py

@@ -1,27 +1,54 @@
 # Determine AppleMobileFileIntegrity's OS configuration
+# Copyright (C) 2022-2023, Mykola Grymalyuk
 
+import enum
 from resources import utilities
 
-class amfi_configuration_detection:
+
+class AmfiConfigDetectLevel(enum.IntEnum):
+    """
+    Configuration levels used by AmfiConfigurationDetection
+    """
+
+    NO_CHECK:                   int = 0
+    LIBRARY_VALIDATION:         int = 1  # For Ventura, use LIBRARY_VALIDATION_AND_SIG
+    LIBRARY_VALIDATION_AND_SIG: int = 2
+    ALLOW_ALL:                  int = 3
+
+
+class AmfiConfigurationDetection:
+    """
+    Detect AppleMobileFileIntegrity's OS configuration
+
+    Usage:
+
+    >>> from resources.amfi_detect import AmfiConfigurationDetection
+    >>> can_patch = AmfiConfigurationDetection().check_config(AmfiConfigDetectLevel.AMFI_LEVEL_ALLOW_ALL)
+
+    """
 
     def __init__(self):
-        self.AMFI_ALLOW_TASK_FOR_PID =      False
-        self.AMFI_ALLOW_INVALID_SIGNATURE = False
-        self.AMFI_LV_ENFORCE_THIRD_PARTY =  False
-        self.AMFI_ALLOW_EVERYTHING =        False
-        self.SKIP_LIBRARY_VALIDATION =      False
+        self.AMFI_ALLOW_TASK_FOR_PID:      bool = False
+        self.AMFI_ALLOW_INVALID_SIGNATURE: bool = False
+        self.AMFI_LV_ENFORCE_THIRD_PARTY:  bool = False
+        self.AMFI_ALLOW_EVERYTHING:        bool = False
+        self.SKIP_LIBRARY_VALIDATION:      bool = False
 
-        self.boot_args = []
-        self.oclp_args = []
+        self.boot_args: list = []
+        self.oclp_args: list = []
 
-        self.init_nvram_dicts()
+        self._init_nvram_dicts()
 
-        self.parse_amfi_bitmask()
-        self.parse_amfi_boot_args()
-        self.parse_oclp_configuration()
+        self._parse_amfi_bitmask()
+        self._parse_amfi_boot_args()
+        self._parse_oclp_configuration()
 
 
-    def init_nvram_dicts(self):
+    def _init_nvram_dicts(self):
+        """
+        Initialize the boot-args and OCLP-Settings NVRAM dictionaries
+        """
+
         boot_args = utilities.get_nvram("boot-args", decode=True)
         oclp_args = utilities.get_nvram("OCLP-Settings", "4D1FDA02-38C7-4A6A-9CC6-4BCCA8B30102", decode=True)
 
@@ -32,8 +59,12 @@ class amfi_configuration_detection:
             self.oclp_args = oclp_args.split(" ")
 
 
-    def parse_amfi_bitmask(self):
-        # See data/amfi_data.py for more information
+    def _parse_amfi_bitmask(self):
+        """
+        Parse the AMFI bitmask from boot-args
+        See data/amfi_data.py for more information
+        """
+
         amfi_value = 0
         for arg in self.boot_args:
             if arg.startswith("amfi="):
@@ -65,7 +96,11 @@ class amfi_configuration_detection:
             self.AMFI_ALLOW_INVALID_SIGNATURE = True
 
 
-    def parse_amfi_boot_args(self):
+    def _parse_amfi_boot_args(self):
+        """
+        Parse the AMFI boot-args
+        """
+
         for arg in self.boot_args:
             if arg.startswith("amfi_unrestrict_task_for_pid"):
                 value = arg.split("=")
@@ -86,26 +121,34 @@ class amfi_configuration_detection:
                         self.AMFI_ALLOW_INVALID_SIGNATURE = True
 
 
-    def parse_oclp_configuration(self):
+    def _parse_oclp_configuration(self):
+        """
+        Parse the OCLP configuration
+        """
+
         if "-allow_amfi" in self.oclp_args:
             self.SKIP_LIBRARY_VALIDATION = True
 
 
-    def check_config(self, level):
-        # Levels:
-        # - 0: No checks
-        # - 1. Library Validation (Monterey and Older)
-        # - 2. Library Validation and Signature Checks (Ventura and Newer)
-        # - 3. Disable all AMFI checks
+    def check_config(self, level: int):
+        """
+        Check the AMFI configuration based on provided AMFI level
+        See AmfiConfigLevel enum for valid levels
 
-        if level == 0:
+        Parameters:
+           level (int): The level of AMFI checks to check for
+
+        Returns:
+              bool: True if the AMFI configuration matches the level, False otherwise
+        """
+
+        if level == AmfiConfigDetectLevel.NO_CHECK:
             return True
-
-        if level == 1:
+        if level == AmfiConfigDetectLevel.LIBRARY_VALIDATION:
             return self.SKIP_LIBRARY_VALIDATION
-        if level == 2:
+        if level == AmfiConfigDetectLevel.LIBRARY_VALIDATION_AND_SIG:
             return bool(self.SKIP_LIBRARY_VALIDATION and self.AMFI_ALLOW_INVALID_SIGNATURE)
-        if level == 3:
+        if level == AmfiConfigDetectLevel.ALLOW_ALL:
             return self.AMFI_ALLOW_EVERYTHING
 
         return False

resources/sys_patch/sys_patch_detect.py

@@ -490,9 +490,9 @@ class detect_root_patch:
                     if self.amfi_shim_bins is True:
                         # Currently we require AMFI outright disabled
                         # in Ventura to work with shim'd binaries
-                        return 3
-                return 1
-        return 0
+                        return amfi_detect.AmfiConfigDetectLevel.ALLOW_ALL
+                return amfi_detect.AmfiConfigDetectLevel.LIBRARY_VALIDATION
+        return amfi_detect.AmfiConfigDetectLevel.NO_CHECK
 
     def verify_patch_allowed(self, print_errors=False):
         sip_dict = self.check_sip()
@@ -500,7 +500,7 @@ class detect_root_patch:
         sip_value = sip_dict[1]
 
         self.sip_enabled, self.sbm_enabled, self.fv_enabled, self.dosdude_patched = utilities.patching_status(sip, self.constants.detected_os)
-        self.amfi_enabled = not amfi_detect.amfi_configuration_detection().check_config(self.get_amfi_level_needed())
+        self.amfi_enabled = not amfi_detect.AmfiConfigurationDetection().check_config(self.get_amfi_level_needed())
 
         if self.nvidia_web is True:
             self.missing_nv_web_nvram   = not self.check_nv_web_nvram()