From 1633b975c3876e47880cca3ce4c55d21d7df2798 Mon Sep 17 00:00:00 2001
From: Joost Agterhoek <mail@joostagterhoek.nl>
Date: Tue, 10 Jun 2025 21:04:43 +0200
Subject: [PATCH] minor changes

---
 flask_soc_site/src/host_data.py      |  9 ++++++---
 flask_soc_site/src/host_lookup.py    | 21 +++++++++++++++++++--
 flask_soc_site/src/virustotal_api.py |  3 ++-
 3 files changed, 27 insertions(+), 6 deletions(-)

diff --git a/flask_soc_site/src/host_data.py b/flask_soc_site/src/host_data.py
index 044c693..57ef389 100644
--- a/flask_soc_site/src/host_data.py
+++ b/flask_soc_site/src/host_data.py
@@ -5,6 +5,7 @@ import validators
 from ipwhois import IPWhois
 from whois import whois
 from .constants import DOMAIN, EMAIL, IPV4, IPV6, URL
+from datetime import datetime
 
 
 def determine(host):
@@ -25,9 +26,11 @@ def determine(host):
 
 
 def domain(host):
-    result = dict(whois(host))
-    if type(result["domain_name"]) is list:
-        result["domain_name"] = result["domain_name"][0]
+    info = dict(whois(host))
+    if type(info["domain_name"]) is list:
+        info["domain_name"] = info["domain_name"][0]
+    selection = {"registrar_country", "registrar"}
+    result = {key: info.get(key) for key in selection}
     return result
 
 
diff --git a/flask_soc_site/src/host_lookup.py b/flask_soc_site/src/host_lookup.py
index f78f316..28505fb 100644
--- a/flask_soc_site/src/host_lookup.py
+++ b/flask_soc_site/src/host_lookup.py
@@ -17,12 +17,21 @@ from . import host_data
 import tldextract
 
 
+class Test(object):
+    def __init__(self):
+        self.host = "joostagterhoek.nl"
+        self.host_type = "domain"
+        self.vt = {"whatever": "whatever", "second_key": 4}
+        self.email_security = ["1", 1, "2", "3"]
+
+
 class Lookedup(object):
     def __init__(self, host):
         self.host = host
         self.host_type = determine(self.host)
         self = self.specific()
         # TODO: consolidate all below functions if possible
+        # Implement an __eq__ method based on host name: https://dnmtechs.com/comparing-object-instances-by-attributes-in-python-3-a-comprehensive-guide/
 
     def url_lookup(self):
         self.domain = urlparse(self.host).netloc
@@ -46,7 +55,8 @@ class Lookedup(object):
         self.metadata = domain(self.host)
         self.domain = self.metadata["domain_name"]
         self.email_security = spf_dmarc(self.domain)
-        self.vt, self.vt_dict = virustotal_api.analyse2(self.host, self.host_type)
+        # self.vt, self.vt_dict = virustotal_api.analyse2(self.host, self.host_type)
+        self.vt = virustotal_api.analyse2(self.host, self.host_type)
         self.abuseipdb = abuseipdb_api.analyse(self.ip_address)
         return self
 
@@ -127,18 +137,25 @@ def extract(user_input):
             hosts.append(item)
         else:
             errors.append(item)
+    print(hosts)
+    print(errors)
     return hosts, errors
 
 
 def domain(host):
     result = dict(whois(host))
+    print("RESULT IS: ", result)
     if type(result["creation_date"]) is list:
         result["creation_date"] = result["creation_date"][0].strftime("%d-%m-%Y")
     else:
         result["creation_date"] = result["creation_date"].strftime("%d-%m-%Y")
     if type(result["domain_name"]) is list:
         result["domain_name"] = result["domain_name"][0]
-    return result
+    # result["creation_date"] = result["creation_date"].isoformat()
+    included = {"domain_name", "creation_date", "registrar", "registrar_country"}
+    filtered = {key: value for key, value in result.items() if key in included}
+    print("FILTERED IS: ", filtered)
+    return filtered
 
 
 def spf_dmarc(domain):
diff --git a/flask_soc_site/src/virustotal_api.py b/flask_soc_site/src/virustotal_api.py
index 3283593..57405ef 100644
--- a/flask_soc_site/src/virustotal_api.py
+++ b/flask_soc_site/src/virustotal_api.py
@@ -108,4 +108,5 @@ def analyse2(host, host_type):
         "%d-%m-%Y",
         time.gmtime(last_update),
     )
-    return summary, analysis_json
+    # return summary, analysis_json
+    return summary