1
0

main.py 3.2 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394
  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. pihole = ph.PiHole(PIHOLE_HOSTNAME)
  16. influxdb_client = InfluxDBClient(DB_ADDRESS, DB_PORT, DB_USER, DB_PASSWORD, None)
  17. def init_db():
  18. databases = influxdb_client.get_list_database()
  19. if len(list(filter(lambda x: x['name'] == DB_DATABASE, databases))) == 0:
  20. influxdb_client.create_database(DB_DATABASE) # Create if does not exist.
  21. print('{} - Created database {}'.format(datetime.datetime.now(), DB_DATABASE))
  22. else:
  23. # Switch to if does exist.
  24. influxdb_client.switch_database(DB_DATABASE)
  25. print('{} - Switched to database {}'.format(datetime.datetime.now(), DB_DATABASE))
  26. def get_data_for_influxdb():
  27. influx_data = [
  28. {
  29. 'measurement': 'domains',
  30. 'time': datetime.datetime.now(),
  31. 'fields': {
  32. 'domain_count': int(pihole.domain_count.replace(',','')),
  33. 'unique_domains': int(pihole.unique_domains.replace(',','')),
  34. 'forwarded': int(pihole.forwarded.replace(',','')),
  35. 'cached': int(pihole.cached.replace(',',''))
  36. }
  37. },
  38. {
  39. 'measurement': 'queries',
  40. 'time': datetime.datetime.now(),
  41. 'fields': {
  42. 'queries': int(pihole.queries.replace(',','')),
  43. 'blocked': int(pihole.blocked.replace(',','')),
  44. 'ads_percentage': float(pihole.ads_percentage)
  45. }
  46. },
  47. {
  48. 'measurement': 'clients',
  49. 'time': datetime.datetime.now(),
  50. 'fields': {
  51. 'total_clients': int(pihole.total_clients.replace(',','')),
  52. 'unique_clients': int(pihole.unique_clients.replace(',','')),
  53. 'total_queries': int(pihole.total_queries.replace(',',''))
  54. }
  55. },
  56. {
  57. 'measurement': 'other',
  58. 'time': datetime.datetime.now(),
  59. 'fields': {
  60. 'status': True if pihole.status == 'enabled' else False,
  61. 'gravity_last_update': pihole.gravity_last_updated['absolute']
  62. }
  63. }
  64. ]
  65. return influx_data
  66. def main():
  67. init_db()
  68. while(1):
  69. pihole.refresh()
  70. data = get_data_for_influxdb()
  71. if influxdb_client.write_points(data) == True:
  72. print("{} - Data written to DB successfully".format(datetime.datetime.now()))
  73. print("{} - Now sleeping for {}s".format(datetime.datetime.now(), TEST_INTERVAL))
  74. time.sleep(TEST_INTERVAL)
  75. else:
  76. print('{} - Failed to write points to the database'.format(datetime.datetime.now()))
  77. time.sleep(120) # Sleep for two seconds.
  78. if __name__ == '__main__':
  79. print('{} - PiHole Data Logger to InfluxDB'.format(datetime.datetime.now()))
  80. main()