Merge pull request #101 from DomainTools/id-1388-find-averages

ID-1388 Python Wrapper: Find averages

Author: wesleya

domaintools/utils.py

@@ -10,7 +10,8 @@ def get_domain_age(create_date):
 
     Returns: Number of days between now and the create_date.
     """
-    time_diff = datetime.now() - dateparser.parse(create_date)
+    time_diff = datetime.now() - dateparser.parse(create_date, date_formats=['%Y-%m-%d', '%Y%m%d'])
+
     return time_diff.days
 
 
@@ -27,4 +28,48 @@ def get_threat_component(components, threat_type):
         if component.get("name") == threat_type:
             return component
     else:
-        return None
+        return None
+
+
+def get_average_risk_score(domains):
+    """
+    Gets average domain risk score for Investigate and Detect result sets
+    Args:
+        domains: Investigate or Detect result set
+
+    Returns: average risk score
+    """
+    count = 0
+    total = 0
+    for d in domains:
+        # investigate result set
+        if "risk_score" in d.get("domain_risk", {}):
+            count += 1
+            total += d.get("domain_risk").get("risk_score")
+        # detect result set
+        elif d.get("risk_score"):
+            count += 1
+            total += d.get("risk_score")
+
+    return total // count if count else None
+
+
+def get_average_age(domains):
+    """
+    Gets average domain age for Investigate and Detect result sets
+    Args:
+        domains: Investigate or Detect result set
+
+    Returns: average age
+    """
+    count = 0
+    total = 0
+    for d in domains:
+        if isinstance(d.get("create_date"), dict) and d.get("create_date").get("value"):
+            count += 1
+            total += get_domain_age(d.get("create_date").get("value"))
+        elif isinstance(d.get("create_date"), int):
+            count += 1
+            total += get_domain_age(str(d.get("create_date")))
+
+    return total // count if count else None

tests/test_utils.py

@@ -13,4 +13,93 @@ def test_get_threat_component():
         {"name": "threat_profile_spam", "risk_score": 73, "threats": ["spam"]}
     ]
     result = utils.get_threat_component(threat_components, "threat_profile_malware")
-    assert result.get("risk_score") == 31
+    assert result.get("risk_score") == 31
+
+
+def test_investigate_average_risk_score():
+    domains = [
+        {"domain_risk": {"risk_score": 25}},
+        {"domain_risk": {"risk_score": 27}}
+    ]
+    result = utils.get_average_risk_score(domains)
+    assert result == 26
+
+    domains = [
+        {"domain_risk": {"risk_score": 25}},
+        {}
+    ]
+    result = utils.get_average_risk_score(domains)
+    assert result == 25
+
+    domains = []
+    result = utils.get_average_risk_score(domains)
+    assert result == None
+
+def test_detect_average_risk_score():
+    domains = [
+        {"risk_score": 25},
+        {"risk_score": 27}
+    ]
+    result = utils.get_average_risk_score(domains)
+    assert result == 26
+
+    domains = [
+        {"risk_score": 25},
+        {"risk_score": None}
+    ]
+    result = utils.get_average_risk_score(domains)
+    assert result == 25
+
+    domains = []
+    result = utils.get_average_risk_score(domains)
+    assert result == None
+
+def test_investigate_average_age():
+    two_days_ago = (datetime.now() - timedelta(days=2)).strftime("%Y-%m-%d")
+    five_days_ago = (datetime.now() - timedelta(days=5)).strftime("%Y-%m-%d")
+
+    domains = [
+        {"create_date": {"value": two_days_ago}},
+        {"create_date": {"value": five_days_ago}}
+    ]
+    result = utils.get_average_age(domains)
+    assert result == 3
+
+    domains = [
+        {"create_date": {"value": two_days_ago}},
+        {}
+    ]
+    result = utils.get_average_age(domains)
+    assert result == 2
+
+    domains = [
+        {"create_date": {"value": two_days_ago}},
+        {"create_date": {"value": ""}}
+    ]
+    result = utils.get_average_age(domains)
+    assert result == 2
+
+    domains = []
+    result = utils.get_average_age(domains)
+    assert result == None
+
+def test_detect_average_age():
+    two_days_ago = int((datetime.now() - timedelta(days=2)).strftime("%Y%m%d"))
+    five_days_ago = int((datetime.now() - timedelta(days=5)).strftime("%Y%m%d"))
+    domains = [
+        {"create_date": two_days_ago},
+        {"create_date": five_days_ago}
+    ]
+    result = utils.get_average_age(domains)
+    assert result == 3
+
+    domains = [
+        {"create_date": two_days_ago},
+        {"create_date": None}
+    ]
+    result = utils.get_average_age(domains)
+    assert result == 2
+
+    domains = []
+    result = utils.get_average_risk_score(domains)
+    assert result == None