1
0

main.py 5.1 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142
  1. import time
  2. import datetime
  3. import os
  4. import sys
  5. import logging
  6. from influxdb_client import InfluxDBClient, Point
  7. from influxdb_client.client.write_api import SYNCHRONOUS
  8. from pihole import PiHole # PiHole API Wrapper
  9. # Logger Settings
  10. logger = logging.Logger('pihole-to-influxdb')
  11. logger.addHandler(logging.StreamHandler(sys.stdout))
  12. logging.basicConfig(level=os.environ.get("LOGLEVEL", "DEBUG"))
  13. # InfluxDB Settings
  14. DB_URL = os.environ.get('INFLUX_DB_URL')
  15. DB_ORG = os.environ.get('INFLUX_DB_ORG')
  16. DB_TOKEN = os.environ.get('INFLUX_DB_TOKEN')
  17. DB_BUCKET = os.environ.get('INFLUX_DB_BUCKET')
  18. # PiHole Settings
  19. PIHOLE_HOSTNAME = str(os.environ.get('PIHOLE_HOSTNAME'))
  20. TEST_INTERVAL = int(os.environ.get('PIHOLE_INTERVAL'))
  21. # Authentication
  22. AUTHENTICATION_TOKEN = os.environ.get('PIHOLE_AUTHENTICATION')
  23. influxdb_client = InfluxDBClient(DB_URL, DB_TOKEN, org=DB_ORG)
  24. def get_data_for_influxdb(pihole: PiHole, timestamp: datetime.datetime):
  25. return [
  26. Point("domains") \
  27. .time(timestamp) \
  28. .tag("hostname", PIHOLE_HOSTNAME) \
  29. .field("domain_count", int(pihole.domain_count.replace(',',''))) \
  30. .field("unique_domains", int(pihole.unique_domains.replace(',',''))) \
  31. .field("forwarded", int(pihole.forwarded.replace(',',''))) \
  32. .field("cached", int(pihole.cached.replace(',',''))),
  33. Point("queries") \
  34. .time(timestamp) \
  35. .tag("hostname", PIHOLE_HOSTNAME) \
  36. .field("queries", int(pihole.queries.replace(',',''))) \
  37. .field("blocked", int(pihole.blocked.replace(',',''))) \
  38. .field("ads_percentage", float(pihole.ads_percentage)),
  39. Point("clients") \
  40. .time(timestamp) \
  41. .tag("hostname", PIHOLE_HOSTNAME) \
  42. .field("total_clients", int(pihole.total_clients.replace(',',''))) \
  43. .field("unique_clients", int(pihole.unique_clients.replace(',',''))) \
  44. .field("total_queries", int(pihole.total_queries.replace(',',''))),
  45. Point("other") \
  46. .time(timestamp) \
  47. .tag("hostname", PIHOLE_HOSTNAME) \
  48. .field("status", pihole.status == 'enabled') \
  49. .field("gravity_last_update", pihole.gravity_last_updated['absolute'])
  50. ]
  51. def get_authenticated_data_for_influxdb(pihole: PiHole, timestamp: datetime.datetime):
  52. query_type_point = Point("query_types") \
  53. .time(timestamp) \
  54. .tag("hostname", PIHOLE_HOSTNAME)
  55. for key, value in pihole.query_types.items():
  56. query_type_point.field(key, float(value))
  57. forward_destinations_point = Point("forward_destinations") \
  58. .time(timestamp) \
  59. .tag("hostname", PIHOLE_HOSTNAME)
  60. for key, value in pihole.forward_destinations['forward_destinations'].items():
  61. forward_destinations_point.field(key.split('|')[0], value)
  62. return [
  63. query_type_point,
  64. forward_destinations_point
  65. ]
  66. class Auth(object):
  67. def __init__(self, token):
  68. # PiHole's web token is just a double sha256 hash of the utf8 encoded password
  69. self.token = token
  70. self.auth_timestamp = time.time()
  71. def main():
  72. # pihole ctor has side effects, so we create it locally
  73. pihole = PiHole(PIHOLE_HOSTNAME)
  74. write_api = influxdb_client.write_api(write_options=SYNCHRONOUS)
  75. USE_AUTHENTICATION = False if AUTHENTICATION_TOKEN == None else True
  76. if USE_AUTHENTICATION:
  77. try:
  78. pihole.auth_data = Auth(AUTHENTICATION_TOKEN)
  79. pihole.refresh()
  80. logger.info('Pi-Hole authentication successful')
  81. except Exception as e:
  82. logger.exception('Pi-Hole authentication failed')
  83. USE_AUTHENTICATION = False
  84. raise
  85. next_update = time.monotonic()
  86. while True:
  87. try:
  88. pihole.refresh()
  89. timestamp = datetime.datetime.now()
  90. data = get_data_for_influxdb(pihole, timestamp)
  91. if USE_AUTHENTICATION:
  92. authenticated_data = get_authenticated_data_for_influxdb(pihole, timestamp)
  93. try:
  94. write_api.write(bucket=DB_BUCKET, record=authenticated_data)
  95. except Exception as e:
  96. logger.exception('Failed to write authenticated data to InfluxDB')
  97. try:
  98. write_api.write(bucket=DB_BUCKET, record=data)
  99. logger.debug('Wrote data to InfluxDB')
  100. sleep_duration = TEST_INTERVAL
  101. except Exception as e:
  102. logger.exception('Failed to write data to InfluxDB')
  103. # Sleep at most two minutes
  104. sleep_duration = min(TEST_INTERVAL, 120)
  105. except Exception as e:
  106. logger.exception('Failed to get data from Pi-Hole')
  107. # Sleep at most two minutes
  108. sleep_duration = min(TEST_INTERVAL, 120)
  109. next_update = next_update + sleep_duration
  110. logger.debug("Now sleeping for {}".format(datetime.timedelta(seconds=sleep_duration)))
  111. time.sleep(max(0, next_update - time.monotonic()))
  112. if __name__ == '__main__':
  113. logger.info('PiHole Data Logger to InfluxDB')
  114. main()