From 05aa90a1cb7a9d357a5c742164f852c8e9757299 Mon Sep 17 00:00:00 2001
From: iRaven4522 <horse@iraven.net>
Date: Thu, 19 Sep 2024 02:12:15 -0500
Subject: [PATCH] finally got this working!!!!

---
 conf-template.json |  15 ++--
 main.py            | 171 +++++++++++++++++++++++++++++++++------------
 2 files changed, 136 insertions(+), 50 deletions(-)

diff --git a/conf-template.json b/conf-template.json
index 16fbbf2..70c3880 100644
--- a/conf-template.json
+++ b/conf-template.json
@@ -1,7 +1,10 @@
 {
-    "loglevel": "debug",
-    "npmapi": "https://npm.corp.iraven.net",
-    "npmkey": "Bearer eyJhbGciOiJSUzI1NiIsInR5cCI6IkpXVCJ9.eyJpc3MiOiJhcGkiLCJzY29wZSI6WyJ1c2VyIl0sImF0dHJzIjp7ImlkIjoxfSwiZXhwaXJlc0luIjoiMWQiLCJqdGkiOiJaT0IrNmV6WkRVVHR0RW9SIiwiaWF0IjoxNzIyODMyMjU2LCJleHAiOjE3MjI5MTg2NTZ9.utpeUVqOWSq1N2qXTsjjMjmvIenTFkLvQTsYPSFMAwK7v2U8XJgDhuQaqaas0Tv1tVpAuVK8LpJC7WQUPE1TkOzQormp5tb7EqGvQy4jRQRZ5q6xkUuE_KW2ME5M4rboy1cjANNiWeRMhCpkWBg4JreFYd0-vWZIQf-3LnkXux-rcaWyRvDsWoB3xkNXvQLnuTIREAAH3CuGkFpbsJD8AjMN8ciYofWP8ZaE9zzbPDMFQ84mXRhpmmQVgorWC9JlGEzz_-AYUrq-bB5v03tHeqVfWvFtmVlfUpObUlDIDQr2MBGXHErIPa8d9amqqZLHSbo7D96QsFMrfavLn9TBRg",
-    "piholeapi": "https://pihole.corp.iraven.net",
-    "piholekey": "balls"
-}
+    "loglevel": "info",
+    "npmdnshostname": "<put the DNS hostname of your Nginx Proxy Manager server here>",
+    "npmemail": "<put your Nginx Proxy Manager email here>",
+    "npmpassword": "<put your Nginx Proxy Manager password here>",
+    "removedead": "True/False",
+    "adddisabled": "True/False",
+    "piholeurl": "<Put your pi-hole url here>",
+    "piholepass": "<Put your pi-hole password here>"
+}
\ No newline at end of file
diff --git a/main.py b/main.py
index c33add0..f87c93c 100644
--- a/main.py
+++ b/main.py
@@ -2,6 +2,7 @@ import os
 import json
 import logging
 import requests
+import re
 from ipaddress import ip_address
 
 # Open the config file and make it accessible via "cfg"
@@ -25,88 +26,170 @@ log.debug("Logger was initialized")
 def checkIfIP(host):
     try:
         ip_address(host)
-        logging.debug(str(host) + " is an IP address")
+        log.debug(str(host) + " is an IP address")
         return True
     except:
-        logging.debug(host + " is not an IP address.")
+        log.debug(host + " is not an IP address.")
         return False
 
-def getNPMHosts(apiurl,apikey,type):
+def loginToNPM(apiurl, user, pw): # To retrieve NPM's API session token
+    # Returns the token used for API requests.
+    session = requests.Session()
+    url = apiurl + "/api/tokens"
+    loginreq = session.post(url,data={"identity": user, "secret": pw})
+    if loginreq.status_code == 200: # Check if the API accepted our creds.
+        log.info("Logged into Nginx Proxy Manager")
+        response = loginreq.json()
+        return response['token']
+    elif loginreq.status_code == 401:
+        log.error("Your Nginx Proxy Manager credentials are incorrect. Please verify you put in the right ones!")
+        exit()
+    else:
+        log.error("Nginx Proxy Manager: " + response['error']['message'])
+        return None
+
+
+    
+def loginToPihole(apiurl, phpassword): 
+    # For whatever reason, Pi-hole's web CSRF token is in a hidden <div> element returned by the login.php in the homepage right after you login.
+    # This took *WAY* too long to figure out. This should be in the header for ease-of-use.
+    session = requests.Session()
+    url = apiurl + "/login.php"
+    loginrq = session.post(url,data={"pw": phpassword})
+    phpsessid = (str(loginrq.cookies.get_dict()["PHPSESSID"]))
+    regex = r'(<div id="token" hidden>)(\S+)(<\/div>)'
+    token = re.findall(regex, loginrq.text, re.MULTILINE)[0][1]
+    log.debug("phpsessid: " + phpsessid + " Token: "+ token)
+    return {'phpsessid': phpsessid, "csrftoken": token}
+
+def getNPMHosts(apiurl,type):
     log.debug("Type is set to "+str(type)+".")
+        # Add the port if it exists in the config
+    if cfg["npmadminport"] is not None:
+        log.debug("Added port "+ cfg["npmadminport"] + " to the URL")
+        apiurl = "http://"+ apiurl + ":" + cfg["npmadminport"]
+
     if type == "proxy":
         log.info("Retrieving proxy hosts from NPM API...")
         url = apiurl + '/api/nginx/proxy-hosts'
-        logging.debug("Set hosts API URL to " + url + ".")
+        log.debug("Set hosts API URL to " + url + ".")
     elif type == "redir":
         log.info("Retrieving redirection hosts from NPM API...")
         url = apiurl + '/api/nginx/redirection-hosts'
-        logging.debug("Set hosts API URL to " + url + ".")
+        log.debug("Set hosts API URL to " + url + ".")
     elif type == "404" or type == "dead":
         log.info("Retrieving 404 (dead) hosts from NPM API...")
         url = apiurl + '/api/nginx/dead-hosts'
-        logging.debug("Set hosts API URL to " + url + ".")
+        log.debug("Set hosts API URL to " + url + ".")
     else:
         log.error("No type of hosts to retrieve from NPM was specified.")
         return None
     hostlist = [] # Make the list.
-    # url = apiurl + '/api/nginx/proxy-hosts'
-    # logging.debug("Set hosts API URL to " + url + ".")
-    apireq = requests.get(url, headers={'Authorization': apikey})
+    apikey = loginToNPM(apiurl,cfg["npmemail"],cfg["npmpassword"])
+    apireq = requests.get(url, headers={'Authorization': "Bearer "+apikey})
     if apireq.headers['content-type'] == "application/json; charset=utf-8": # Check if the API shat us a JSON.
-        logging.debug("Passed content-type = application/json check.")
-        logging.debug("HTTP Status Code: "+ str(apireq.status_code))
+        log.debug("Passed content-type = application/json check.")
         if apireq.status_code == 200: # Check if the API returned a 200 and accepted our token/key.
-            logging.debug("API returned a 200, proceeding.") # We're good c:
+            log.debug("API returned a 200, proceeding.") # We're good c:
             fullresponse = apireq.json()
             for i in range(len(fullresponse)):
-                logging.debug("Adding indice " + str(i) + " to list, containing "+ str(fullresponse[i]['domain_names']))
-                hostlist = hostlist + fullresponse[i]['domain_names']
-            logging.debug("List has been created.")
+                if cfg['addDisabled']:
+                    log.debug("Adding indice " + str(i) + " to list, containing "+ str(fullresponse[i]['domain_names']))
+                    hostlist = hostlist + fullresponse[i]['domain_names']
+                else:
+                    if fullresponse[i]['enabled'] == 1:
+                        log.debug("Adding indice " + str(i) + " to list, containing "+ str(fullresponse[i]['domain_names']))
+                        hostlist = hostlist + fullresponse[i]['domain_names']
+                    else:
+                        log.debug("Not adding indice " + str(i) + "to list.")
+            log.debug(type + " list has been created.")
             for h in range(len(hostlist)):
                 if checkIfIP(hostlist[h]):
-                    logging.debug("Deleting IP address "+ hostlist[h] + " from the list of hosts.")
+                    log.info("Deleting IP address "+ hostlist[h] + " from the list of " + type + " hosts.")
                     hostlist.pop(h)
                     break # Temp fix!!
+            log.info(type + " hosts have been retrieved. Count: " + str(len(hostlist)))
             return hostlist
         elif apireq.status_code == 403: # If the API gave us a 403 Forbidden/Permission Denied
-            logging.debug("API returned a 403, halting!")
-            logging.error("The API retured a permission denied error!")
-            logging.error("Please make sure your Nginx Proxy Manager API key in conf.json is correct and accurate.")
+            log.debug("API returned a 403, halting!")
+            log.error("The API retured a permission denied error!")
+            log.error("Please make sure your Nginx Proxy Manager API key in conf.json is correct and accurate.")
+            return None
+        elif apireq.status_code == 401: # If the API gave us a 401 Unauthorized
+            log.debug("API returned a 401, halting!")
+            log.error("The API retured an unauthorized error!")
+            log.error("Please make sure your Nginx Proxy Manager API key in conf.json is correct and accurate.")
             return None
         else:
-            logging.error("The API returned a " + str(apireq.status_code)+".")
-            logging.error("Please make sure your Nginx Proxy Manager API key or URL in conf.json is correct and accurate.")
+            log.error("The API returned a " + str(apireq.status_code)+".")
+            log.error("Please make sure your Nginx Proxy Manager API key or URL in conf.json is correct and accurate.")
             return None
     else: # If we don't get a JSON
-        logging.error("The API did not return a JSON, and instead a content type of "+str(apireq.headers['content-type'])+".")
-        logging.error("Please make sure your Nginx Proxy Manager API key or URL in conf.json is correct and accurate.")
+        log.error("The API did not return a JSON, and instead a content type of "+str(apireq.headers['content-type'])+".")
+        log.error("Please make sure your Nginx Proxy Manager API key or URL in conf.json is correct and accurate.")
         return None
 
-def addPiHoleHosts(apiurl, apikey,targetsvr, list):
+def addPiHoleHosts(apiurl, phpassword, targetsvr, list):
     url = apiurl + "/scripts/pi-hole/php/customcname.php"
-    for i in list:    
-        payload = "action=add&domain="+list[i]+"&target="+targetsvr+"&token="+apikey
-        apireq = requests.post(url,data=payload)
+    piauth = loginToPihole(apiurl,phpassword)
+    for i in list:
+        payload = {"action": "add", "domain": i, "target": targetsvr, "token": piauth["csrftoken"]}
+        apireq = requests.post(url,data=payload,cookies={"PHPSESSID": piauth["phpsessid"]})
+        if apireq.status_code == 200: # Check if the API returned a 200 and accepted our token/key.
+            log.debug("Adding host "+ i + " to Pi-hole's CNAME list.")
+            response = apireq.json()
+            log.debug(response)
+            try:
+                if response['success'] == False:
+                    if "There is already" in response['message']:
+                        log.info("There is already a CNAME record for " + i)
+                    else:
+                        log.debug("Pi-Hole API returned false!")
+                        log.warning("The Pi-Hole API gave the following message:"+ response['message'])
+                elif response['success'] == True:
+                    log.info("Added " + i + " to CNAME list.")
+                    log.debug("PiHole API returned true. Message returned: "+ response['message'])
+            except:
+                log.error("Pi-hole returned this message and was not JSON: "+ response)
+                
+
+def removePiHoleHosts(apiurl, phpassword, targetsvr, list):
+    url = apiurl + "/scripts/pi-hole/php/customcname.php"
+    piauth = loginToPihole(apiurl,phpassword)
+    for i in list:
+        log.debug("Removing host "+ i + " from Pi-hole's CNAME list.")
+        payload = {"action": "add", "domain": i, "target": targetsvr, "token": piauth["csrftoken"]}
+        apireq = requests.post(url,data=payload,cookies={"PHPSESSID": piauth["phpsessid"]})
         if apireq.status_code == 200: # Check if the API returned a 200 and accepted our token/key.
             response = apireq.json()
-            if response['success'] == "false":
-                log.debug("Pi-Hole API returned false!")
-                log.warning("The Pi-Hole API gave the following message:"+ response['message'])
-            elif response['success'] == "true":
-                logging.debug("PiHole API returned true. Message returned: "+ response['message'])
-        else:
+            try:
+                if response['success'] == "false":
+                    log.debug("Pi-Hole API returned false!")
+                    log.warning("The Pi-Hole API gave the following message:"+ response['message'])
+                elif response['success'] == "true":
+                    log.debug("PiHole API returned true. Message returned: "+ response['message'])
+            except:
+                log.error("Pi-hole returned this message: "+ response)
+
 
 
 log.debug("Script has started!")
-proxyhosts = getNPMHosts(cfg['npmapi'],cfg['npmkey'],"proxy") # Get Proxy Hosts
-redirhosts = getNPMHosts(cfg['npmapi'],cfg['npmkey'],"redir") # Get Redir Hosts
-deadhosts = getNPMHosts(cfg['npmapi'],cfg['npmkey'],"dead") # Get 404 Hosts
-log.debug("Adding all the hosts together")
-allhosts = proxyhosts + redirhosts
-if cfg['removedead'] == False: # if "removedead" in the config is true, add the dead hosts to the main list, but otherwise no.
-    allhosts = allhosts + deadhosts
+log.info("Getting all NPM hosts.")
+proxyhosts = getNPMHosts(cfg['npmdnshostname'],"proxy") # Get All Proxy Hosts
+redirhosts = getNPMHosts(cfg['npmdnshostname'],"redir") # Get All Redir Hosts
+deadhosts = getNPMHosts(cfg['npmdnshostname'],"dead") # Get All 404 Hosts
+
+if ((proxyhosts == None) or (redirhosts == None) or (deadhosts == None)):
+    log.error("One (or more) of the lists of hosts returned None. Check above for any errors while retriving hosts from NPM.")
+    exit()
 else:
-    removePiHoleHosts(cfg["piholeapi"],cfg["piholekey"],deadhosts)
-print(allhosts)
-import random
-print(allhosts[random.randint()])
\ No newline at end of file
+    log.info("NPM Hosts were retrieved!")
+    log.debug("Adding all the hosts together")
+    if cfg['removeDead'] == False: # if "removeDead" in the config is true, add the dead hosts to the main list, but otherwise no.
+        allhosts = proxyhosts + redirhosts + deadhosts
+    else:
+        allhosts = proxyhosts + redirhosts
+    log.info("Adding all hosts in the list to specified Pi-Hole server at "+ cfg["piholeurl"] + "...")
+    addPiHoleHosts(cfg["piholeurl"],cfg["piholepass"],cfg["npmdnshostname"],allhosts)
+    log.info("Success!")
+log.debug("Script ending")
\ No newline at end of file