From 3a392f2084f4b7e2a8dbd8109076e022f6809ef3 Mon Sep 17 00:00:00 2001
From: Josh Smith <git@breadlysm.com>
Date: Thu, 30 Jul 2020 23:45:23 -0400
Subject: [PATCH] created tag_selection function

---
 main.py | 47 ++++++++++++++++++++++++++++++++++++++++++-----
 1 file changed, 42 insertions(+), 5 deletions(-)

diff --git a/main.py b/main.py
index 9a2e148..228ab69 100755
--- a/main.py
+++ b/main.py
@@ -10,10 +10,13 @@ DB_PORT = int(os.environ.get('INFLUX_DB_PORT'))
 DB_USER = os.environ.get('INFLUX_DB_USER')
 DB_PASSWORD = os.environ.get('INFLUX_DB_PASSWORD')
 DB_DATABASE = os.environ.get('INFLUX_DB_DATABASE')
+DB_TAGS = os.environ.get('INFLUX_DB_TAGS')
 
 # Speedtest Settings
-TEST_INTERVAL = int(os.environ.get('SPEEDTEST_INTERVAL'))  # Time between tests (in seconds).
-TEST_FAIL_INTERVAL = int(os.environ.get('SPEEDTEST_FAIL_INTERVAL'))  # Time before retrying a failed Speedtest (in seconds).
+# Time between tests (in seconds).
+TEST_INTERVAL = int(os.environ.get('SPEEDTEST_INTERVAL'))
+# Time before retrying a failed Speedtest (in seconds).
+TEST_FAIL_INTERVAL = int(os.environ.get('SPEEDTEST_FAIL_INTERVAL'))
 
 influxdb_client = InfluxDBClient(
     DB_ADDRESS, DB_PORT, DB_USER, DB_PASSWORD, None)
@@ -26,13 +29,47 @@ def init_db():
         influxdb_client.create_database(
             DB_DATABASE)  # Create if does not exist.
     else:
-        influxdb_client.switch_database(DB_DATABASE)  # Switch to if does exist.
+        # Switch to if does exist.
+        influxdb_client.switch_database(DB_DATABASE)
+
+
 def pkt_loss(data):
     if 'packetLoss' in data.keys():
         return data['packetLoss']
-    else: 
+    else:
         return 0
 
+
+def tag_selection(data):
+    tags = DB_TAGS
+    tag_switch = {
+        'isp': data['isp'],
+        'interface': data['interface']['name'],
+        'internal_ip': data['interface']['internalIp'],
+        'interface_mac': data['interface']['macAddr'],
+        'vpn_enabled': (False if data['interface']['isVpn'] == 'false' else True),
+        'external_ip': data['interface']['externalIp'],
+        'server_id': data['server']['id'],
+        'server_name': data['server']['name'],
+        'server_location': data['server']['location'],
+        'server_country': data['server']['country'],
+        'server_host': data['server']['host'],
+        'server_port': data['server']['port'],
+        'server_ip': data['server']['ip'],
+        'speedtest_id': data['result']['id'],
+        'speedtest_url': data['result']['url']
+    }
+    options = {}
+    if tags == '':
+        return None
+    else:
+        tags = tags.split(',')
+        for tag in tags:
+            tag = tag.strip()
+            options[tag] = tag_switch[tag]
+        return options
+
+
 def format_for_influx(cliout):
     data = json.loads(cliout)
     # There is additional data in the speedtest-cli output but it is likely not necessary to store.
@@ -99,4 +136,4 @@ def main():
 
 if __name__ == '__main__':
     print('Speedtest CLI Data Logger to InfluxDB')
-    main()
\ No newline at end of file
+    main()