fix the include counts

modules/spf.py

@@ -18,15 +18,15 @@ def __init__(self, domain, dns_server=None):
             self.too_many_includes = self.num_includes > 10
 
     def get_spf_record(self):
-        """Returns the SPF record for the domain."""
+        """Returns the SPF record for a given domain."""
         try:
             resolver = dns.resolver.Resolver()
-            if self.dns_server:
-                resolver.nameservers = [self.dns_server]
+            resolver.nameservers = [self.dns_server, '1.1.1.1', '8.8.8.8']
             query_result = resolver.resolve(self.domain, 'TXT')
             for record in query_result:
-                if 'v=spf1' in str(record):
-                    return str(record).replace('"', '')
+                if 'spf1' in str(record):
+                    spf_record = str(record).replace('"', '')
+                    return spf_record
             return None
         except Exception:
             return None
@@ -40,25 +40,36 @@ def get_spf_all_string(self):
             elif len(all_matches) > 1:
                 return '2many'
         return None
-
-    def get_spf_includes(self, count=0):
-        """Returns the number of includes in the SPF record for the domain."""
-        if count > 10:  # Assuming a maximum of 10 includes as a threshold
-            return count
-        try:
-            if self.spf_record:
-                count += self.spf_record.count("include:")
-                # Recursively check includes
-                for item in self.spf_record.split(' '):
-                    if "include:" in item:
-                        included_domain = item.replace('include:', '')
-                        # Instantiate SPF class for the included domain to get its includes
-                        include_spf = SPF(included_domain, self.dns_server)
-                        count += include_spf.get_spf_includes(count)
-            return count
-        except Exception:
+
+    def get_spf_includes(self):
+        """Returns the number of includes and other mechanisms in the SPF record for a given domain."""
+        def count_includes(spf_record):
+            count = 0
+            for item in spf_record.split():
+                if item.startswith("include:"):
+                    url = item.replace('include:', '')
+                    count += 1
+                    try:
+                        # Recursively fetch and count includes in the SPF record of the included domain
+                        answers = dns.resolver.resolve(url, 'TXT')
+                        for rdata in answers:
+                            for txt_string in rdata.strings:
+                                txt_record = txt_string.decode('utf-8')
+                                if txt_record.startswith('v=spf1'):
+                                    count += count_includes(txt_record)
+                    except Exception as e:
+                        pass
+
+            # Count occurrences of 'a', 'mx', 'ptr', and 'exists' mechanisms
+            count += len(re.findall(r"[ ,+]a[ ,:]", spf_record))
+            count += len(re.findall(r"[ ,+]mx[ ,:]", spf_record))
+            count += len(re.findall(r"[ ]ptr[ ]", spf_record))
+            count += len(re.findall(r"exists[:]", spf_record))
+
             return count
 
+        return count_includes(self.spf_record)
+
     def __str__(self):
         return (f"SPF Record: {self.spf_record}\n"
                 f"All Mechanism: {self.all_mechanism}\n"