Improve code readability

Author: NeffIsBack

nxc/modules/sccm.py

@@ -22,8 +22,8 @@ class NXCModule:
     multiple_hosts = True
 
     def __init__(self):
-        self.sccm_site_servers = []    # List of dns host names of the SCCM site servers
-        self.sccm_sites = {}
+        self.sccm_site_servers = []     # List of dns host names of the SCCM site servers
+        self.sccm_sites = {}            # List of SCCM sites with their management points (Sorted by site code)
         self.base_dn = ""
 
     def options(self, context, module_options):
@@ -39,6 +39,7 @@ def on_login(self, context, connection):
         self.sc = ldap.SimplePagedResultsControl()
 
         try:
+            # Search for SCCM root object
             search_filter = f"(distinguishedName=CN=System Management,CN=System,{self.base_dn})"
             controls = security_descriptor_control(sdflags=0x04)
             context.log.display(f"Looking for the SCCM container with filter: '{search_filter}'")
@@ -49,18 +50,21 @@ def on_login(self, context, connection):
                 searchControls=controls,
                 searchBase=self.base_dn,
             )
+
+            # There should be only one result
             for item in result:
                 if isinstance(item, ldapasn1_impacket.SearchResultEntry):
                     self.context.log.success(f"Found SCCM object: {item[0]}")
                     self.get_site_servers(item)
                     self.get_sites()
                     self.get_management_points()
 
-                    self.context.log.success("Site Servers:")
+                    # Print results
+                    self.context.log.success(f"Found {len(self.sccm_site_servers)} Site Servers:")
                     for site in self.sccm_site_servers:
                         ip = self.connection.resolver(site)
                         self.context.log.highlight(f"{site} - {ip['host'] if ip else 'unknown'}")
-                    self.context.log.success("SCCM Sites:")
+                    self.context.log.success(f"Found {len(self.sccm_sites)} SCCM Sites:")
                     for site in self.sccm_sites:
                         self.context.log.highlight(f"{self.sccm_sites[site]['cn']}")
                         self.context.log.highlight(f"  Site Code: {site.rjust(14)}")
@@ -72,19 +76,20 @@ def on_login(self, context, connection):
                             self.context.log.highlight(f"\t  IP Address:{' ':<4}{mp['IPAddress']}")
                             self.context.log.highlight(f"\t  Default MP:{' ':<4}{mp['mSSMSDefaultMP']}")
                     self.context.log.highlight("")
-
         except LDAPSearchError as e:
             context.log.fail(f"Got unexpected exception: {e}")
 
     def get_management_points(self):
-        """Searches for all SCCM management points in the Active Directory and maps them to their SCCM site."""
+        """Searches for all SCCM management points in the Active Directory and maps them to their SCCM site via the site code."""
         try:
             response = self.connection.ldapConnection.search(
+                searchBase=self.base_dn,
                 searchFilter="(objectClass=mSSMSManagementPoint)",
-                attributes="*",
+                attributes=["cn", "dNSHostName", "mSSMSDefaultMP", "mSSMSSiteCode"],
             )
+
             response_parsed = parse_result_attributes(response)
-            self.context.log.success("SCCM Management Points:")
+
             for mp in response_parsed:
                 ip = self.connection.resolver(mp["dNSHostName"])
                 self.sccm_sites[mp["mSSMSSiteCode"]]["ManagementPoints"].append({
@@ -98,13 +103,16 @@ def get_management_points(self):
             self.context.log.error(f"Error searching for management points: {e}")
 
     def get_sites(self):
-        """Searches for all SCCM sites in the Active Directory."""
+        """Searches for all SCCM sites in the Active Directory, sorted by site code."""
         try:
             response = self.connection.ldapConnection.search(
+                searchBase=self.base_dn,
                 searchFilter="(objectClass=mSSMSSite)",
                 attributes=["cn", "mSSMSSiteCode", "mSSMSAssignmentSiteCode"],
             )
+
             response_parsed = parse_result_attributes(response)
+
             for site in response_parsed:
                 self.sccm_sites[site["mSSMSSiteCode"]] = {
                     "cn": site["cn"],
@@ -116,7 +124,7 @@ def get_sites(self):
             self.context.log.error(f"Error searching for sites: {e}")
 
     def get_site_servers(self, item):
-        """Extracts the site servers from the SCCM object."""
+        """Extracts the site servers from the root SCCM object."""
         raw_sec_descriptor = str(item[1][0][1][0]).encode("latin-1")
         principal_security_descriptor = ldaptypes.SR_SECURITY_DESCRIPTOR(data=raw_sec_descriptor)
         self.parse_dacl(principal_security_descriptor["Dacl"])
@@ -147,7 +155,9 @@ def resolve_SID(self, sid):
                 searchFilter=f"(objectSid={sid})",
                 attributes=["sAMAccountName", "sAMAccountType", "member", "dNSHostName"],
             )
+
             parsed_result = parse_result_attributes(result)
+
             if not parsed_result:
                 return None
             else:
@@ -179,5 +189,6 @@ def dn_to_sid(self, dn) -> str:
             attributes=["sAMAccountName", "objectSid"],
         )
 
+        # Extract the SID of the object
         sid_raw = bytes(result[0][1][0][1].components[0])
         return ldaptypes.LDAP_SID(data=sid_raw).formatCanonical()