1
0

main.py 5.2 KB

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