Introduce a cache

subdomain_takeover_tools/confirm_takeover.py

@@ -1,3 +1,4 @@
+import re
 import sys
 
 from subdomain_takeover_tools.confirm_agile_crm import is_valid as agile_crm_is_valid
@@ -16,6 +17,9 @@
 from subdomain_takeover_tools.confirm_surge import is_valid as surge_is_valid
 from subdomain_takeover_tools.confirm_tumblr import is_valid as tumblr_is_valid
 from subdomain_takeover_tools.confirm_unclaimed import is_valid as unclaimed_is_valid
+from subdomain_takeover_tools.helper.prepare import get_cached, store_cache
+
+cacheable = ["azure", "elasticbeanstalk"]
 
 
 def main():
@@ -29,8 +33,8 @@ def main():
     for line in lines:
         if not line.strip():
             continue
-        elif ']\t\t' not in line:
-            raise IOError("Unexpected input received, currently only subtake output is supported")
+        elif ']\t\t' not in line and '.]   ' not in line:
+            raise IOError("Unexpected input s received, currently only subtake output is supported")
 
         (service, target, domain) = _process_line(line)
 
@@ -44,7 +48,10 @@ def main():
 
 
 def _process_line(line):
-    (parts, domain) = line.split('\t\t')
+    if '\t\t' in line:
+        (parts, domain) = line.split('\t\t')
+    else:
+        (parts, domain) = re.split(r'  +', line)
     if ': ]' in parts:
         service = parts[1:-3]
         target = ''
@@ -65,6 +72,18 @@ def _process_subtake_output(service, target, domain, inverse):
 
 
 def _perform_check(service, target, domain):
+    if service in cacheable:
+        cached = get_cached(service, target)
+        if cached is not None:
+            return cached
+
+    fresh_result = perform_uncached(domain, service, target)
+    if service in cacheable and fresh_result is not None:
+        store_cache(service, target, fresh_result)
+    return fresh_result
+
+
+def perform_uncached(domain, service, target):
     if service == 'agilecrm':
         return agile_crm_is_valid(domain, target)
     elif service == 'azure':

subdomain_takeover_tools/confirm_unclaimed.py

@@ -6,7 +6,7 @@
 from tldextract import tldextract
 
 from subdomain_takeover_tools.helper.main import bootstrap, settings
-from subdomain_takeover_tools.helper.prepare import prepare_domain_name
+from subdomain_takeover_tools.helper.prepare import prepare_domain_name, get_cached, store_cache
 
 # once we include additional support we might need to lazy initialize
 transip_session = requests.Session()
@@ -49,14 +49,22 @@ def cname_still_valid(domain, cname):
 def confirm_unclaimed(cname):
     domain_name = _extract_single_domain_name(prepare_domain_name(cname))
 
+    cached = get_cached('unclaimed', domain_name)
+    if cached is not None:
+        return cached
+
+    return store_cache('unclaimed', domain_name, confirm_uncached(domain_name))
+
+
+def confirm_uncached(domain_name):
     for blacklisted in tld_blacklist:
         if domain_name.endswith(blacklisted):
             return False
-
     if settings['transip'] is not None:
         return confirm_unclaimed_transip(domain_name)
     else:
         sys.stderr.write("No supported unclaimed configuration is setup, please consult the documentation")
+        return None
 
 
 def confirm_unclaimed_transip(domain_name):

subdomain_takeover_tools/helper/prepare.py

@@ -2,6 +2,8 @@
 
 import dns.resolver
 
+cache = {}
+
 
 def prepare_domain_name(host_in):
     host = re.sub(r'https?://', '', host_in)
@@ -10,6 +12,24 @@ def prepare_domain_name(host_in):
     return host
 
 
+def get_cached(cname, target):
+    key = (cname, target)
+    if key in cache:
+        return cache[key]
+    else:
+        return None
+
+
+def store_cache(cname, target, result):
+    if result is None:
+        return None
+
+    key = (cname, target)
+    cache[key] = result
+
+    return result
+
+
 def process_subtake_output(is_valid, line, check, inverse, strict):
     (parts, domain) = line.split('\t\t')
     if ': ]' in parts: