From 5cbb03f7742e3683560ec8e06fd98bdfe7f663a9 Mon Sep 17 00:00:00 2001
From: iRaven4522 <horse@iraven.net>
Date: Wed, 30 Apr 2025 20:15:41 -0500
Subject: [PATCH] Updated for Pi-Hole v6!

---
 main.py | 136 +++++++++++++++++++++++++++++++++++++++++++++-----------
 1 file changed, 111 insertions(+), 25 deletions(-)

diff --git a/main.py b/main.py
index f20e577..6fa4cb4 100644
--- a/main.py
+++ b/main.py
@@ -49,12 +49,34 @@ def loginToNPM(apiurl, user, pw): # To retrieve NPM's API session token
         return None
 
 
-    
-def loginToPihole(apiurl, phpassword): 
+def loginToPihole(serveraddr, phpassword, session):
+    # NEW Method as of Pi-Hole version 6.0 and newer.
+    # I truly appreciate them fixing but this broke so many things.
+    apiurl = "http://" + serveraddr + "/api/auth"
+    passwordjson = {"password": phpassword}
+    loginrq = session.request("POST", apiurl, json=passwordjson, verify=False)
+    log.debug("Sent POST for getting Pi-hole API authorization.")
+    if loginrq.status_code == 200: # yay
+        log.debug("API returned a 200, proceeding.") # We're good c:
+        log.info("Logged into Pi-hole API on "+serveraddr+"!")
+        fullresponse = loginrq.json()
+        phpsessid = fullresponse["session"]["sid"]
+        token = fullresponse["session"]["csrf"]
+        log.debug("phpsessid: " + phpsessid + " Token: "+ token)
+        return {'phpsessid': phpsessid, "csrftoken": token}
+    elif loginrq.status_code == 401: # Unauthorized
+        log.debug("Pi-hole returned a 401, halting!")
+        log.error("The Pi-hole retured an unauthorized error!")
+        log.error("Please make sure your Pi-hole password in conf.json is correct and accurate.")
+        return None
+    else:
+        log.error("Pi-hole returned this " + str(apireq.status_code) + " in error: "+ apireq.text)
+
+def loginToPiholeOld(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"
+    url = "http://" + apiurl + "/admin/login.php"
     loginrq = session.post(url,data={"pw": phpassword})
     phpsessid = (str(loginrq.cookies.get_dict()["PHPSESSID"]))
     regex = r'(<div id="token" hidden>)(\S+)(<\/div>)'
@@ -133,13 +155,9 @@ def getNPMHosts(apiurl,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, phpassword, targetsvr, list, filter):
-    url = apiurl + "/scripts/pi-hole/php/customcname.php"
-    hostname = apiurl.split("/")[2]
-    piauth = loginToPihole(apiurl,phpassword)
+def filterPiHoleHosts(targetsvr, list, filter):
     # Adding filter functionality
     if filter is not None:
-        log.info("A domain filter was applied for " + hostname + ". Filtering out domains in this run.")
         filterlist = filter.replace(" ","").split(",")
         for domain in filterlist:
             log.debug("Filtering domain " + domain)
@@ -147,12 +165,38 @@ def addPiHoleHosts(apiurl, phpassword, targetsvr, list, filter):
                 log.debug("Host " + host + " Filter " + domain)
                 if domain in host:
                     index = list.index(host)
-                    log.info("Removing host " + host + " from the hostlist on " + hostname + " due to the filtered domain " + domain + " .")
+                    log.info("Removing host " + host + " from the filtered host list due to the filtered domain " + domain + ".")
                     list.remove(host)
                 else:
                     log.debug("The domain " + host + "does not contain the domain " + domain)
             log.debug(domain + " domain was filtered, moving on to next one")
         log.info("Domains were filtered. Adding the list to the server...")
+    return list
+
+
+def addPiHoleHosts(apiurl, phpassword, targetsvr, list, session):
+    baseurl = "http://" + apiurl + "/api/config/dns/cnameRecords/"
+    piauth = loginToPihole(apiurl,phpassword,session)
+    hostname = apiurl
+    # Moved filter functionality to filterPiHoleHosts().
+    for i in list:
+        log.debug("Adding host " + i + " to " + hostname + " Pi-hole CNAME list.")
+        url = baseurl + i + "," + targetsvr
+        apireq = requests.put(url,cookies={"sid": piauth["phpsessid"]},headers={"X-CSRF-TOKEN": piauth["csrftoken"]})
+        response = apireq.json()
+        if apireq.status_code == 201: # Check if the API returned a "created object" message
+            log.info("Added " + i + " to " + hostname + " CNAME list.")
+            log.debug("PiHole API returned true. Message returned: "+ apireq.text)
+        elif apireq.status_code == 400 and ("Item already present" in response['error']['message']):
+            log.info("There is already a CNAME record for " + i)
+        else:
+            log.error("Pi-hole returned this " + str(apireq.status_code) + " in error: "+ apireq.text)
+
+def addPiHoleHostsOld(apiurl, phpassword, targetsvr, list):
+    url = "http://" + apiurl + "/scripts/pi-hole/php/customcname.php"
+    hostname = apiurl
+    piauth = loginToPiholeOld(apiurl,phpassword)
+    # Moved filter functionality to filterPiHoleHosts().
     for i in list:
         payload = {"action": "add", "domain": i, "target": targetsvr, "token": piauth["csrftoken"]}
         apireq = requests.post(url,data=payload,cookies={"PHPSESSID": piauth["phpsessid"]})
@@ -173,8 +217,25 @@ def addPiHoleHosts(apiurl, phpassword, targetsvr, list, filter):
             except:
                 log.error("Pi-hole returned this message and was not JSON: "+ response)
                 
+def removePiHoleHosts(apiurl, phpassword, targetsvr, list, session):
+    baseurl = "http://" + apiurl + "/api/config/dns/cnameRecords/"
+    piauth = loginToPihole(apiurl,phpassword,session)
+    hostname = apiurl
+    # Moved filter functionality to filterPiHoleHosts().
+    for i in list:
+        log.debug("Adding host " + i + " to " + hostname + " Pi-hole CNAME list.")
+        url = baseurl + i + "," + targetsvr
+        apireq = requests.put(url,cookies={"sid": piauth["phpsessid"]},headers={"X-CSRF-TOKEN": piauth["csrftoken"]})
+        response = apireq.json()
+        if apireq.status_code == 201: # Check if the API returned a "created object" message
+            log.info("Added " + i + " to " + hostname + " CNAME list.")
+            log.debug("PiHole API returned true. Message returned: "+ apireq.text)
+        elif apireq.status_code == 400 and ("Item already present" in response['error']['message']):
+            log.info("There is already a CNAME record for " + i)
+        else:
+            log.error("Pi-hole returned this " + str(apireq.status_code) + " in error: "+ apireq.text)
 
-def removePiHoleHosts(apiurl, phpassword, targetsvr, list):
+def removePiHoleHostsOld(apiurl, phpassword, targetsvr, list):
     url = apiurl + "/scripts/pi-hole/php/customcname.php"
     piauth = loginToPihole(apiurl,phpassword)
     for i in list:
@@ -199,6 +260,7 @@ 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
+session = requests.Session()
 
 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.")
@@ -211,23 +273,47 @@ else:
     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,None)
-    log.info("Succeeded in adding all hosts in the list to specified Pi-Hole server at "+ cfg["piholeurl"])
-    if cfg['filterenabled'] == "True":
-        log.debug("The filter server option was enabled.")
-        try:
-            if cfg["filterlist"] is not str(''):
-                log.info("Adding all hosts in the list to filtered Pi-Hole server at "+ cfg["filteredpiholeurl"] + "...")
-                if cfg["filteredpiholepass"] is not str(''):
-                    addPiHoleHosts(cfg["filteredpiholeurl"],cfg["filteredpiholepass"],cfg["npmdnshostname"],allhosts,cfg["filterlist"])
+    ## I'll clean up the following later lmfao
+    if cfg["phOldVersion"] == True:
+        addPiHoleHostsOld(cfg["piholeurl"],cfg["piholepass"],cfg["npmdnshostname"],allhosts)
+        log.info("Succeeded in adding all hosts in the list to specified Pi-Hole server at "+ cfg["piholeurl"])
+        if cfg['filterenabled'] == "True":
+            log.debug("The filter server option was enabled.")
+            try:
+                if cfg["filterlist"] is not str(''):
+                    filteredhosts = filterPiHoleHosts(cfg["npmdnshostname"],allhosts,cfg["filterlist"])
+                    log.info("Adding all hosts in the list to filtered Pi-Hole server at "+ cfg["filteredpiholeurl"] + "...")
+                    if cfg["filteredpiholepass"] is not str(''):
+                        addPiHoleHostsOld(cfg["filteredpiholeurl"],cfg["filteredpiholepass"],cfg["npmdnshostname"],filteredhosts)
+                    else:
+                        addPiHoleHostsOld(cfg["filteredpiholeurl"],cfg["piholepass"],cfg["npmdnshostname"],filteredhosts)
+                        log.info("Succeeded in adding all hosts in the list to filtered Pi-Hole server at "+ cfg["filteredpiholeurl"])
                 else:
-                    addPiHoleHosts(cfg["filteredpiholeurl"],cfg["piholepass"],cfg["npmdnshostname"],allhosts,cfg["filterlist"])
-                log.info("Succeeded in adding all hosts in the list to filtered Pi-Hole server at "+ cfg["filteredpiholeurl"])
-            else:
+                    log.error("There is nothing in the filter!")
+                    log.error("Make sure you put at least 1 (one) domain in the filter list.")
+            except KeyError:
+                log.error("There is nothing in the filter!")
+                log.error("Make sure you put at least 1 (one) domain in the filter list.")
+    else:
+        addPiHoleHosts(cfg["piholeurl"],cfg["piholepass"],cfg["npmdnshostname"],allhosts,session)
+        log.info("Succeeded in adding all hosts in the list to specified Pi-Hole server at "+ cfg["piholeurl"])
+        if cfg['filterenabled'] == "True":
+            log.debug("The filter server option was enabled.")
+            try:
+                if cfg["filterlist"] is not str(''):
+                    log.info("Filtering all hosts for filtered Pi-hole server " + cfg["filteredpiholeurl"])
+                    filteredhosts = filterPiHoleHosts(cfg["npmdnshostname"],allhosts,cfg["filterlist"])
+                    log.info("Adding all hosts in the list to filtered Pi-Hole server at "+ cfg["filteredpiholeurl"] + "...")
+                    if cfg["filteredpiholepass"] is not str(''):
+                        addPiHoleHosts(cfg["filteredpiholeurl"],cfg["filteredpiholepass"],cfg["npmdnshostname"],filteredhosts,session)
+                    else:
+                        addPiHoleHosts(cfg["filteredpiholeurl"],cfg["piholepass"],cfg["npmdnshostname"],filteredhosts,session)
+                        log.info("Succeeded in adding all hosts in the list to filtered Pi-Hole server at "+ cfg["filteredpiholeurl"])
+                else:
+                    log.error("There is nothing in the filter!")
+                    log.error("Make sure you put at least 1 (one) domain in the filter list.")
+            except KeyError:
                 log.error("There is nothing in the filter!")
                 log.error("Make sure you put at least 1 (one) domain in the filter list.")
-        except KeyError:
-            log.error("There is nothing in the filter!")
-            log.error("Make sure you put at least 1 (one) domain in the filter list.")
 
 log.debug("Script ending")
\ No newline at end of file