1
0

main.py 3.7 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107
  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 main():
  70. init_db()
  71. if USE_AUTHENTICATION:
  72. try:
  73. pihole.authenticate(AUTHENTICATION_TOKEN)
  74. pihole.refresh()
  75. except:
  76. print("{} - Authentication failed using token: {}, disabling authentication.".format(datetime.datetime.now(), AUTHENTICATION_TOKEN))
  77. USE_AUTHENTICATION = False
  78. raise
  79. while(1):
  80. pihole.refresh()
  81. data = get_data_for_influxdb()
  82. if influxdb_client.write_points(data) == True:
  83. print("{} - Data written to DB successfully".format(datetime.datetime.now()))
  84. print("{} - Now sleeping for {}s".format(datetime.datetime.now(), TEST_INTERVAL))
  85. time.sleep(TEST_INTERVAL)
  86. else:
  87. print('{} - Failed to write points to the database'.format(datetime.datetime.now()))
  88. time.sleep(120) # Sleep for two seconds.
  89. if __name__ == '__main__':
  90. print('{} - PiHole Data Logger to InfluxDB'.format(datetime.datetime.now()))
  91. main()