1
0

main.py 5.1 KB

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