From 581a2fff727745729e89b52957b1bf50da7f3684 Mon Sep 17 00:00:00 2001
From: SimonGurney <simon.gurney@punksecurity.co.uk>
Date: Wed, 19 Jul 2023 18:01:30 +0100
Subject: [PATCH] feat: added Project Discovery Provider (#158)

* feat: added Project Discovery Provider

* feat: fixed some issues

* chore: ran black

* feat: fixed typo

* feat: removed unecessary comment

* feat: fixed readme

* feat: fixed readme

---------

Co-authored-by: Victoria Kotiwcki <victoria.kotwicki@punksecurity.co.uk>
---
 README.md                     |  9 +++++
 providers/projectdiscovery.py | 70 +++++++++++++++++++++++++++++++++++
 2 files changed, 79 insertions(+)
 create mode 100644 providers/projectdiscovery.py

diff --git a/README.md b/README.md
index 653f374..ec42d61 100644
--- a/README.md
+++ b/README.md
@@ -156,6 +156,15 @@ digitalocean:
   --do-domains DO_DOMAINS
                         Optional
 
+projectdiscovery:
+  Scan multiple domains by fetching them from Project Discovery
+
+  --pd-api-key PD_API_KEY
+                        Required
+  --pd-domains PD_DOMAINS
+                        Required
+
+
 file:
   Read domains from a file (or folder of files), one per line
 
diff --git a/providers/projectdiscovery.py b/providers/projectdiscovery.py
new file mode 100644
index 0000000..b44407f
--- /dev/null
+++ b/providers/projectdiscovery.py
@@ -0,0 +1,70 @@
+import requests, logging, json
+
+from domain import Domain
+
+description = "Scan multiple domains by fetching them from ProjectDiscovery"
+
+
+class DomainNotFoundError(Exception):
+    def __init__(self, domain):
+        self.message = "Domain not found: " + domain
+        super().__init__(self.message)
+
+
+class PDApi:
+    def __init__(self, api_key):
+        self.session = requests.session()
+        self.session.headers.update(
+            {"Content-Type": "application/json", "Authorization": api_key}
+        )
+
+    @staticmethod
+    def check_response(response: requests.Response):
+        if response.status_code == 401:
+            raise ValueError("Invalid API key specified.")
+
+        if response.status_code < 200 or response.status_code >= 300:
+            raise ValueError("Invalid response received from API: " + response.json())
+
+        return response
+
+    def make_request(self, endpoint):
+        return self.session.prepare_request(
+            requests.Request("GET", "https://dns.projectdiscovery.io/dns/" + endpoint)
+        )
+
+    def list_domains(self):
+        req = self.make_request("domains")
+
+        return self.check_response(self.session.send(req))
+
+    def get_subdomains(self, domain):
+        req = self.make_request(f"{domain}/subdomains")
+        res = self.session.send(req)
+
+        if 404 == res.status_code:
+            raise DomainNotFoundError(domain)
+
+        return self.check_response(res)
+
+
+def fetch_domains(pd_api_key: str, pd_domains: str = None, **args):
+    root_domains = []
+    domains = []
+    api = PDApi(pd_api_key)
+
+    if pd_domains is not None and len(pd_domains):
+        root_domains = [domain.strip(" ") for domain in pd_domains.split(",")]
+    else:
+        print("Domain required")
+        exit()
+
+    for domain in root_domains:
+        if "" == domain or domain is None:
+            continue
+
+        raw_domains = api.get_subdomains(domain).json()
+        logging.info("Testing", len(raw_domains["subdomains"]), "subdomains")
+        domains.extend([Domain(f"{sb}.{domain}") for sb in raw_domains["subdomains"]])
+
+    return domains