1
0

main.py 5.2 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 = 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. for key, value in pihole.query_types.items():
  57. yield Point("query_types") \
  58. .time(timestamp) \
  59. .tag("hostname", PIHOLE_HOSTNAME) \
  60. .tag("query_type", key) \
  61. .field("value", float(value))
  62. for key, value in pihole.forward_destinations['forward_destinations'].items():
  63. yield Point("forward_destinations") \
  64. .time(timestamp) \
  65. .tag("hostname", PIHOLE_HOSTNAME) \
  66. .tag("destination", key.split('|')[0]) \
  67. .field("value", value)
  68. class Auth(object):
  69. def __init__(self, token):
  70. # PiHole's web token is just a double sha256 hash of the utf8 encoded password
  71. self.token = token
  72. self.auth_timestamp = time.time()
  73. def main():
  74. # pihole ctor has side effects, so we create it locally
  75. pihole = PiHole(PIHOLE_HOSTNAME)
  76. write_api = influxdb_client.write_api(write_options=SYNCHRONOUS)
  77. USE_AUTHENTICATION = False if AUTHENTICATION_TOKEN == None else True
  78. if USE_AUTHENTICATION:
  79. try:
  80. pihole.auth_data = Auth(AUTHENTICATION_TOKEN)
  81. pihole.refresh()
  82. logger.info('Pi-Hole authentication successful')
  83. except Exception as e:
  84. logger.exception('Pi-Hole authentication failed')
  85. USE_AUTHENTICATION = False
  86. raise
  87. next_update = time.monotonic()
  88. while True:
  89. try:
  90. pihole.refresh()
  91. timestamp = datetime.datetime.now()
  92. data = get_data_for_influxdb(pihole, timestamp)
  93. if USE_AUTHENTICATION:
  94. authenticated_data = list(get_authenticated_data_for_influxdb(pihole, timestamp))
  95. try:
  96. write_api.write(bucket=DB_BUCKET, record=authenticated_data)
  97. except Exception as e:
  98. logger.exception('Failed to write authenticated data to InfluxDB')
  99. try:
  100. write_api.write(bucket=DB_BUCKET, record=data)
  101. logger.debug('Wrote data to InfluxDB')
  102. sleep_duration = TEST_INTERVAL
  103. except Exception as e:
  104. logger.exception('Failed to write data to InfluxDB')
  105. # Sleep at most two minutes
  106. sleep_duration = min(TEST_INTERVAL, 120)
  107. except Exception as e:
  108. logger.exception('Failed to get data from Pi-Hole')
  109. # Sleep at most two minutes
  110. sleep_duration = min(TEST_INTERVAL, 120)
  111. next_update = next_update + sleep_duration
  112. logger.debug("Now sleeping for {}".format(datetime.timedelta(seconds=sleep_duration)))
  113. time.sleep(max(0, next_update - time.monotonic()))
  114. if __name__ == '__main__':
  115. logger.info('PiHole Data Logger to InfluxDB')
  116. main()