From 1abd97f37a00f564cd39c36422352ba8102f55bf Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Rapha=C3=ABl=20Vinot?= <raphael@vinot.info>
Date: Mon, 26 Dec 2022 23:01:44 +0100
Subject: [PATCH] new: switch to return emails only on whois query

---
 lookyloo/modules/uwhois.py | 26 ++++++++++++++++++++++----
 website/web/__init__.py    | 11 +++++++----
 2 files changed, 29 insertions(+), 8 deletions(-)

diff --git a/lookyloo/modules/uwhois.py b/lookyloo/modules/uwhois.py
index f63fa7b2..c5bf0873 100644
--- a/lookyloo/modules/uwhois.py
+++ b/lookyloo/modules/uwhois.py
@@ -1,8 +1,10 @@
 #!/usr/bin/env python3
 
 import logging
+import re
 import socket
-from typing import Any, Dict
+
+from typing import Any, Dict, overload, Literal, List, Union
 
 from har2tree import CrawledTree, Har2TreeError, HostNode
 
@@ -35,9 +37,11 @@ class UniversalWhois():
 
     def query_whois_hostnode(self, hostnode: HostNode) -> None:
         if hasattr(hostnode, 'resolved_ips'):
+            ip: str
             for ip in hostnode.resolved_ips:
                 self.whois(ip)
         if hasattr(hostnode, 'cnames'):
+            cname: str
             for cname in hostnode.cnames:
                 self.whois(cname)
         self.whois(hostnode.name)
@@ -58,7 +62,19 @@ class UniversalWhois():
             for n in hostnode.get_ancestors():
                 self.query_whois_hostnode(n)
 
-    def whois(self, query: str) -> str:
+    @overload
+    def whois(self, query: str, contact_email_only: Literal[True]) -> List[str]:
+        ...
+
+    @overload
+    def whois(self, query: str, contact_email_only: Literal[False]) -> str:
+        ...
+
+    @overload
+    def whois(self, query: str, contact_email_only: bool=False) -> Union[str, List[str]]:
+        ...
+
+    def whois(self, query: str, contact_email_only: bool=False) -> Union[str, List[str]]:
         if not self.available:
             return ''
         bytes_whois = b''
@@ -70,5 +86,7 @@ class UniversalWhois():
                 if not data:
                     break
                 bytes_whois += data
-        to_return = bytes_whois.decode()
-        return to_return
+        if not contact_email_only:
+            return bytes_whois.decode()
+        emails = list(set(re.findall(rb'[\w\.-]+@[\w\.-]+', bytes_whois)))
+        return [e.decode() for e in sorted(emails)]
diff --git a/website/web/__init__.py b/website/web/__init__.py
index caa8fc90..f91aebf8 100644
--- a/website/web/__init__.py
+++ b/website/web/__init__.py
@@ -1062,10 +1062,13 @@ def statsfull():
 
 
 @app.route('/whois/<string:query>', methods=['GET'])
-def whois(query: str):
-    to_return = lookyloo.uwhois.whois(query)
-    return send_file(BytesIO(to_return.encode()),
-                     mimetype='test/plain', as_attachment=True, download_name=f'whois.{query}.txt')
+@app.route('/whois/<string:query>/<int:email_only>', methods=['GET'])
+def whois(query: str, email_only: int=0):
+    to_return = lookyloo.uwhois.whois(query, bool(email_only))
+    if isinstance(to_return, str):
+        return send_file(BytesIO(to_return.encode()),
+                         mimetype='test/plain', as_attachment=True, download_name=f'whois.{query}.txt')
+    return jsonify(to_return)
 
 
 # ##### Methods related to a specific URLNode #####