add app mode selection

main.py

@@ -3,6 +3,7 @@ import datetime
 import os
 import sys
 import logging
+from enum import Enum
 
 from influxdb_client import InfluxDBClient, Point
 from influxdb_client.client.write_api import SYNCHRONOUS
@@ -11,9 +12,14 @@ from pihole import PiHole
 logger = logging.Logger('pihole-to-influxdb')
 logger.addHandler(logging.StreamHandler(sys.stdout))
 
+class AppMode(Enum):
+  Live = 1
+  Totals = 2
+  Raw = 3
+
 try:
   # optional Logger Settings
-  logging.basicConfig(level=os.getenv("LOGLEVEL", "DEBUG"))
+  logging.basicConfig(level=os.getenv("LOG_LEVEL", "DEBUG"))
 
   # InfluxDB Settings
   DB_URL = os.environ['INFLUX_DB_URL']
@@ -22,37 +28,52 @@ try:
   DB_BUCKET = os.environ['INFLUX_DB_BUCKET']
 
   # PiHole Settings
-  PIHOLE_HOSTNAME = str(os.environ['PIHOLE_HOSTNAME'])
-  QUERY_LIVE = bool(os.getenv('PIHOLE_QUERY_LIVE', True))
+  PIHOLE_URL = str(os.environ['PIHOLE_URL'])
   QUERY_INTERVAL = int(os.environ['PIHOLE_INTERVAL'])
 
   # optional Pi-hole authentication
   AUTHENTICATION_TOKEN = os.getenv('PIHOLE_AUTHENTICATION', None)
 
+  # optional App Mode
+  APP_MODE = AppMode(os.getenv('APP_MODE', 'Totals'))
+
 except KeyError as e:
   logger.fatal('Missing environment variable: {}'.format(e))
   sys.exit(1)
+except ValueError as e:
+  logger.fatal('Invalid environment variable: {}'.format(e))
+  sys.exit(1)
 
-influxdb_client = InfluxDBClient(DB_URL, DB_TOKEN, org=DB_ORG)
-pihole = PiHole(PIHOLE_HOSTNAME, AUTHENTICATION_TOKEN)
+if APP_MODE != AppMode.Totals and not AUTHENTICATION_TOKEN:
+  logger.fatal('Pi-hole authentication token is required for live data')
+  sys.exit(1)
 
-class Auth(object):
-  def __init__(self, token):
-    # PiHole's web token is just a double sha256 hash of the utf8 encoded password
-    self.token = token
-    self.auth_timestamp = time.time()
+influxdb_client = InfluxDBClient(DB_URL, DB_TOKEN, org=DB_ORG)
+pihole = PiHole(PIHOLE_URL, AUTHENTICATION_TOKEN)
 
 def main():
   write_api = influxdb_client.write_api(write_options=SYNCHRONOUS)
   next_update = time.monotonic()
 
+  logger.info('Starting PiHole Data Logger to InfluxDB')
+  logger.info('AppMode: {}'.format(APP_MODE))
+
+  # Test Pi-hole connection
+  try:
+    pihole.request_summary()
+  except Exception as e:
+    logger.fatal('Unable to connect to Pi-hole: {}'.format(e))
+    sys.exit(1)
+
   while True:
     try:
-      timestamp = datetime.datetime.now()
-      if QUERY_LIVE:
+      if APP_MODE == AppMode.Live:
+        timestamp = datetime.datetime.now()
         data = list(pihole.get_queries_for_influxdb(timestamp, QUERY_INTERVAL))
-      else:
+      elif APP_MODE == AppMode.Totals:
         data = list(pihole.get_totals_for_influxdb())
+      elif APP_MODE == AppMode.Raw:
+        data = list(pihole.get_query_logs_for_influxdb())
 
       logger.debug('Writing {} points to InfluxDB'.format(len(data)))
       write_api.write(bucket=DB_BUCKET, record=data)
@@ -68,5 +89,4 @@ def main():
     time.sleep(max(0, next_update - time.monotonic()))
 
 if __name__ == '__main__':
-  logger.info('PiHole Data Logger to InfluxDB')
   main()

pihole.py

@@ -2,10 +2,11 @@
 # -*- coding: utf-8 -*-
 
 import requests
-from datetime import datetime, timedelta
+from datetime import datetime
 from enum import Enum
 from influxdb_client import Point
 from pandas import DataFrame
+from urllib.parse import urlparse
 
 class QueryStati(Enum):
   Blocked = 1
@@ -15,12 +16,13 @@ class QueryStati(Enum):
   Unknown = 5
 
 class PiHole:
-  def __init__(self, host, token):
-    self.host = host
+  def __init__(self, url, token):
+    url = urlparse(url)
+    self.url = "{}://{}".format(url.scheme or "http", url.netloc)
     self.token = token
 
   def query(self, endpoint, params={}):
-    url = f"http://{self.host}/admin/{endpoint}.php"
+    url = f"{self.url}/admin/{endpoint}.php"
     return requests.get(url, params=params)
   
   def request_all_queries(self, start: float, end: float):
@@ -209,7 +211,7 @@ class PiHole:
         .tag("destination", key.split('|')[0]) \
         .field("queries", len(group_df))
 
-  def get_logs_for_influxdb(self, query_date: datetime, sample_period: int):
+  def get_query_logs_for_influxdb(self, query_date: datetime, sample_period: int):
     end_time = query_date.timestamp()
     start_time = end_time - sample_period + 1