From da659894777447bdf458ae05c17a2d484f359318 Mon Sep 17 00:00:00 2001
From: iRaven4522 <horse@iraven.net>
Date: Mon, 23 Sep 2024 21:42:44 -0500
Subject: [PATCH] Filtering!

---
 conf-template.json |  5 ++++-
 main.py            | 48 +++++++++++++++++++++++++++++++++++++++-------
 2 files changed, 45 insertions(+), 8 deletions(-)

diff --git a/conf-template.json b/conf-template.json
index 70c3880..2348888 100644
--- a/conf-template.json
+++ b/conf-template.json
@@ -6,5 +6,8 @@
     "removedead": "True/False",
     "adddisabled": "True/False",
     "piholeurl": "<Put your pi-hole url here>",
-    "piholepass": "<Put your pi-hole password here>"
+    "piholepass": "<Put your pi-hole password here>",
+    "filterenabled": "True.False",
+    "filteredpiholeurl": "<Put another pi-hole url here>",
+    "filterlist": "<put domains you don't want on other servers here, separated by commas. example: domain,domain2,domain3>"
 }
\ No newline at end of file
diff --git a/main.py b/main.py
index f87c93c..3ec5408 100644
--- a/main.py
+++ b/main.py
@@ -129,14 +129,31 @@ 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):
+def addPiHoleHosts(apiurl, phpassword, targetsvr, list, filter):
     url = apiurl + "/scripts/pi-hole/php/customcname.php"
+    hostname = apiurl.split("/")[2]
     piauth = loginToPihole(apiurl,phpassword)
+    # 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)
+            for host in list[:]:
+                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 + " .")
+                    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...")
     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.")
+            log.debug("Adding host " + i + " to " + hostname + " Pi-hole CNAME list.")
             response = apireq.json()
             log.debug(response)
             try:
@@ -147,7 +164,7 @@ def addPiHoleHosts(apiurl, phpassword, targetsvr, list):
                         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.info("Added " + i + " to " + hostname + " 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)
@@ -157,8 +174,8 @@ 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"]}
+        log.info("Removing host "+ i + " from Pi-hole's CNAME list.")
+        payload = {"action": "delete", "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()
@@ -190,6 +207,23 @@ 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)
-    log.info("Success!")
+    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"])
+                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.")
+
 log.debug("Script ending")
\ No newline at end of file