main.py 4.8 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137
  1. import time
  2. import datetime
  3. import os
  4. from influxdb import InfluxDBClient
  5. import pihole as ph # PiHole API Wrapper
  6. # InfluxDB Settings
  7. DB_ADDRESS = os.environ.get('INFLUX_DB_ADDRESS')
  8. DB_PORT = int(os.environ.get('INFLUX_DB_PORT'))
  9. DB_USER = os.environ.get('INFLUX_DB_USER')
  10. DB_PASSWORD = os.environ.get('INFLUX_DB_PASSWORD')
  11. DB_DATABASE = os.environ.get('INFLUX_DB_DATABASE')
  12. # PiHole Settings
  13. PIHOLE_HOSTNAME = os.environ.get('PIHOLE_HOSTNAME')
  14. TEST_INTERVAL = int(os.environ.get('PIHOLE_INTERVAL'))
  15. # Authentication
  16. AUTHENTICATION_TOKEN = os.environ.get('PIHOLE_AUTHENTICATION')
  17. USE_AUTHENTICATION = False if AUTHENTICATION_TOKEN == None else True
  18. pihole = ph.PiHole('192.168.113.250')
  19. influxdb_client = InfluxDBClient(DB_ADDRESS, DB_PORT, DB_USER, DB_PASSWORD, None)
  20. def init_db():
  21. databases = influxdb_client.get_list_database()
  22. if len(list(filter(lambda x: x['name'] == DB_DATABASE, databases))) == 0:
  23. influxdb_client.create_database(DB_DATABASE) # Create if does not exist.
  24. print('{} - Created database {}'.format(datetime.datetime.now(), DB_DATABASE))
  25. else:
  26. # Switch to if does exist.
  27. influxdb_client.switch_database(DB_DATABASE)
  28. print('{} - Switched to database {}'.format(datetime.datetime.now(), DB_DATABASE))
  29. def get_data_for_influxdb():
  30. influx_data = [
  31. {
  32. 'measurement': 'domains',
  33. 'time': datetime.datetime.now(),
  34. 'fields': {
  35. 'domain_count': int(pihole.domain_count.replace(',','')),
  36. 'unique_domains': int(pihole.unique_domains.replace(',','')),
  37. 'forwarded': int(pihole.forwarded.replace(',','')),
  38. 'cached': int(pihole.cached.replace(',',''))
  39. }
  40. },
  41. {
  42. 'measurement': 'queries',
  43. 'time': datetime.datetime.now(),
  44. 'fields': {
  45. 'queries': int(pihole.queries.replace(',','')),
  46. 'blocked': int(pihole.blocked.replace(',','')),
  47. 'ads_percentage': float(pihole.ads_percentage)
  48. }
  49. },
  50. {
  51. 'measurement': 'clients',
  52. 'time': datetime.datetime.now(),
  53. 'fields': {
  54. 'total_clients': int(pihole.total_clients.replace(',','')),
  55. 'unique_clients': int(pihole.unique_clients.replace(',','')),
  56. 'total_queries': int(pihole.total_queries.replace(',',''))
  57. }
  58. },
  59. {
  60. 'measurement': 'other',
  61. 'time': datetime.datetime.now(),
  62. 'fields': {
  63. 'status': True if pihole.status == 'enabled' else False,
  64. 'gravity_last_update': pihole.gravity_last_updated['absolute']
  65. }
  66. }
  67. ]
  68. return influx_data
  69. def get_formatted_authenticated_forward_destinations():
  70. formatted_dict = {}
  71. for key in pihole.forward_destinations['forwarded_destinations']:
  72. formatted_dict[key.split('|')[0]] = pihole.forward_destinations['forwarded_destinations'][key]
  73. return formatted_dict
  74. def get_authenticated_data_for_influxdb():
  75. influx_data = [
  76. {
  77. 'measurement': 'authenticated_query_types',
  78. 'time': datetime.datetime.now(),
  79. 'fields': pihole.query_types
  80. },
  81. {
  82. 'measurement': 'authenticated_forward_destinations',
  83. 'time': datetime.datetime.now(),
  84. 'fields': get_formatted_authenticated_forward_destinations()
  85. }
  86. ]
  87. return influx_data
  88. def main():
  89. init_db()
  90. if USE_AUTHENTICATION:
  91. try:
  92. pihole.authenticate(AUTHENTICATION_TOKEN)
  93. pihole.refresh()
  94. except:
  95. print("{} - Authentication failed using token: {}, disabling authentication.".format(datetime.datetime.now(), AUTHENTICATION_TOKEN))
  96. USE_AUTHENTICATION = False
  97. raise
  98. while(1):
  99. pihole.refresh()
  100. data = get_data_for_influxdb()
  101. if USE_AUTHENTICATION:
  102. authenticated_data = get_authenticated_data_for_influxdb()
  103. if influxdb_client.write_points(authenticated_data) == True:
  104. print("{} - Authenticated data written to DB successfully".format(datetime.datetime.now()))
  105. else:
  106. print('{} - Failed to write authenticated points to the database'.format(datetime.datetime.now()))
  107. if influxdb_client.write_points(data) == True:
  108. print("{} - Data written to DB successfully".format(datetime.datetime.now()))
  109. print("{} - Now sleeping for {}s".format(datetime.datetime.now(), TEST_INTERVAL))
  110. time.sleep(TEST_INTERVAL)
  111. else:
  112. print('{} - Failed to write points to the database'.format(datetime.datetime.now()))
  113. time.sleep(120) # Sleep for two seconds.
  114. if __name__ == '__main__':
  115. print('{} - PiHole Data Logger to InfluxDB'.format(datetime.datetime.now()))
  116. main()