From 80f8e87b56675df3f28563849e19d929aa13f4d5 Mon Sep 17 00:00:00 2001
From: NotoriousRebel <matt12299@yahoo.com>
Date: Fri, 5 Jun 2020 21:17:48 -0400
Subject: [PATCH] Fixed how results were displayed when using API for html
 reporting as well as new endpoint for user to query for available sources.

---
 theHarvester/__main__.py    | 14 ++++++++++++--
 theHarvester/lib/stash.py   | 35 ++++++++++++++++++++++++-----------
 theHarvester/lib/web/api.py |  9 +++++++++
 3 files changed, 45 insertions(+), 13 deletions(-)

diff --git a/theHarvester/__main__.py b/theHarvester/__main__.py
index 13feafa4..a7399f7b 100644
--- a/theHarvester/__main__.py
+++ b/theHarvester/__main__.py
@@ -50,6 +50,8 @@ async def start(rest_args=None):
     rest_filename = ""
     # indicates this from the rest API
     if rest_args:
+        if rest_args.source and rest_args.source == "getsources":
+            return list(sorted(Core.get_supportedengines()))
         args = rest_args
         # We need to make sure the filename is random as to not overwrite other files
         filename: str = args.filename
@@ -653,12 +655,20 @@ async def handler(lst):
         try:
             print('\n[*] Reporting started.')
             db = stash.StashManager()
-            scanboarddata = await db.getscanboarddata()
+            if rest_args and rest_args.domain is not None and len(rest_args.domain) > 1:
+                # If using rest API filter by domain
+                scanboarddata = await db.getscanboarddata(domain=rest_args.domain)
+            else:
+                scanboarddata = await db.getscanboarddata()
             latestscanresults = await db.getlatestscanresults(word)
             previousscanresults = await db.getlatestscanresults(word, previousday=True)
             latestscanchartdata = await db.latestscanchartdata(word)
             scanhistorydomain = await db.getscanhistorydomain(word)
-            pluginscanstatistics = await db.getpluginscanstatistics()
+            if rest_args and rest_args.domain is not None and len(rest_args.domain) > 1:
+                # If using rest API filter by domain
+                pluginscanstatistics = await db.getpluginscanstatistics(domain=rest_args.domain)
+            else:
+                pluginscanstatistics = await db.getpluginscanstatistics()
             generator = statichtmlgenerator.HtmlGenerator(word)
             HTMLcode = await generator.beginhtml()
             HTMLcode += await generator.generatedashboardcode(scanboarddata)
diff --git a/theHarvester/lib/stash.py b/theHarvester/lib/stash.py
index c5ae21a4..d5412359 100644
--- a/theHarvester/lib/stash.py
+++ b/theHarvester/lib/stash.py
@@ -146,11 +146,14 @@ async def getlatestscanresults(self, domain, previousday=False):
         except Exception as e:
             print(f'Error connecting to theHarvester database: {e}')
 
-    async def getscanboarddata(self):
+    async def getscanboarddata(self, domain=""):
         try:
             async with aiosqlite.connect(self.db, timeout=30) as conn:
-
-                cursor = await conn.execute('''SELECT COUNT(*) from results WHERE type="host"''')
+                if len(domain) != 0:
+                    cursor = await conn.execute('''SELECT COUNT(*) from results WHERE type="host" and domain=?''',
+                                                (domain,))
+                else:
+                    cursor = await conn.execute('''SELECT COUNT(*) from results WHERE type="host"''')
                 data = await cursor.fetchone()
                 self.scanboarddata["host"] = data[0]
                 cursor = await conn.execute('''SELECT COUNT(*) from results WHERE type="email"''')
@@ -211,16 +214,26 @@ async def getscanhistorydomain(self, domain):
         except Exception as e:
             print(e)
 
-    async def getpluginscanstatistics(self):
+    async def getpluginscanstatistics(self, domain=""):
         try:
             async with aiosqlite.connect(self.db, timeout=30) as conn:
-                cursor = await conn.execute('''
-                SELECT domain,find_date, type, source, count(*)
-                FROM results
-                GROUP BY domain, find_date, type, source
-                ''')
-                results = await cursor.fetchall()
-                self.scanstats = results
+                if len(domain) == 0:
+                    cursor = await conn.execute('''
+                    SELECT domain,find_date, type, source, count(*)
+                    FROM results
+                    GROUP BY domain, find_date, type, source
+                    ''')
+                    results = await cursor.fetchall()
+                    self.scanstats = results
+                else:
+                    cursor = await conn.execute('''
+                                       SELECT domain,find_date, type, source, count(*) 
+                                       FROM results WHERE domain=?
+                                       GROUP BY domain, find_date, type, source
+                                       ''', (domain,))
+                    results = await cursor.fetchall()
+                    self.scanstats = results
+
             return self.scanstats
         except Exception as e:
             print(e)
diff --git a/theHarvester/lib/web/api.py b/theHarvester/lib/web/api.py
index 0ccf86cb..a76eabce 100644
--- a/theHarvester/lib/web/api.py
+++ b/theHarvester/lib/web/api.py
@@ -42,6 +42,15 @@ async def picture():
     return StreamingResponse(io.BytesIO(base64.b64decode(string)))
 
 
+@app.get("/sources", response_class=ORJSONResponse)
+@limiter.limit("5/minute")
+async def getsources(request: Request):
+    # Endpoint for user to query for available sources theHarvester supports
+    # Rate limit of 5 requests per minute
+    sources = await __main__.start(Namespace(source="getsources"))
+    return {'sources': sources}
+
+
 @app.get("/query", response_class=ORJSONResponse)
 @limiter.limit("2/minute")
 async def query(request: Request, dns_server: str = Query(""), user_agent: str = Header(None),