diff --git a/Dockerfile b/Dockerfile index 96b4be4f4c5eee4b523e641b2582f1cf086873e8..b3fa16e2d1487b8865305f44afbc630c4ed76e8e 100644 --- a/Dockerfile +++ b/Dockerfile @@ -2,7 +2,7 @@ FROM python:3.3-slim MAINTAINER Allan Tribe <atribe13@gmail.com> VOLUME /src/ -COPY InfluxdbSpeedtest.py requirements.txt /src/ +COPY influxspeedtest/InfluxdbSpeedtest.py requirements.txt /src/ WORKDIR /src RUN pip install -r requirements.txt diff --git a/InfluxdbSpeedtest.py b/InfluxdbSpeedtest.py deleted file mode 100644 index 08121ec1584cf4d6cbee97ea7ab5cea252bbccef..0000000000000000000000000000000000000000 --- a/InfluxdbSpeedtest.py +++ /dev/null @@ -1,172 +0,0 @@ -import configparser -import os -import sys -import argparse -from influxdb import InfluxDBClient -from influxdb.exceptions import InfluxDBClientError, InfluxDBServerError -import speedtest -import time - -class configManager(): - - def __init__(self, config): - print('Loading Configuration File {}'.format(config)) - self.test_server = [] - config_file = os.path.join(os.getcwd(), config) - if os.path.isfile(config_file): - self.config = configparser.ConfigParser() - self.config.read(config_file) - else: - print('ERROR: Unable To Load Config File: {}'.format(config_file)) - sys.exit(1) - - self._load_config_values() - print('Configuration Successfully Loaded') - - def _load_config_values(self): - - # General - self.delay = self.config['GENERAL'].getint('Delay', fallback=2) - self.output = self.config['GENERAL'].getboolean('Output', fallback=True) - - # InfluxDB - self.influx_address = self.config['INFLUXDB']['Address'] - self.influx_port = self.config['INFLUXDB'].getint('Port', fallback=8086) - self.influx_database = self.config['INFLUXDB'].get('Database', fallback='speedtests') - self.influx_user = self.config['INFLUXDB'].get('Username', fallback='') - self.influx_password = self.config['INFLUXDB'].get('Password', fallback='') - self.influx_ssl = self.config['INFLUXDB'].getboolean('SSL', fallback=False) - self.influx_verify_ssl = self.config['INFLUXDB'].getboolean('Verify_SSL', fallback=True) - - # Speedtest - test_server = self.config['SPEEDTEST'].get('Server', fallback=None) - if test_server: - self.test_server.append(test_server) - - -class InfluxdbSpeedtest(): - - def __init__(self, config=None, single_run=False): - - self.config = configManager(config=config) - self.output = self.config.output - self.single_run = single_run - self.influx_client = InfluxDBClient( - self.config.influx_address, - self.config.influx_port, - username=self.config.influx_user, - password=self.config.influx_password, - database=self.config.influx_database, - ssl=self.config.influx_ssl, - verify_ssl=self.config.influx_verify_ssl - ) - - self.speedtest = None - self.results = None - - def setup_speedtest(self): - - speedtest.build_user_agent() - - print('Getting speedtest.net Configuration') - try: - self.speedtest = speedtest.Speedtest() - except speedtest.ConfigRetrievalError: - print('ERROR: Failed to get speedtest.net configuration. Aborting') - sys.exit(1) - - try: - self.speedtest.get_servers(self.config.test_server) - except speedtest.NoMatchedServers: - print('ERROR: No matched servers: {}'.format(self.config.test_server[0])) - sys.exit(1) - except speedtest.ServersRetrievalError: - print('ERROR: Cannot retrieve speedtest.net server list') - sys.exit(1) - except speedtest.InvalidServerIDType: - print('{} is an invalid server type, must be int'.format(self.config.test_server[0])) - sys.exit(1) - - print('Picking the closest server') - self.speedtest.get_best_server() - - self.results = self.speedtest.results - - def send_results(self): - - result_dict = self.results.dict() - - input_points = [ - { - 'measurement': 'speed_test_results', - 'fields': { - 'download': result_dict['download'], - 'upload': result_dict['upload'], - 'ping': result_dict['server']['latency'] - }, - 'tags': { - 'server': result_dict['server']['sponsor'] - } - } - ] - - if self.output: - print('Download: {}'.format(str(result_dict['download']))) - print('Upload: {}'.format(str(result_dict['upload']))) - - self.write_influx_data(input_points) - - def run(self): - - while True: - self.setup_speedtest() - self.speedtest.download() - self.speedtest.upload() - - self.send_results() - if self.single_run: - return - - time.sleep(self.config.delay) - - def write_influx_data(self, json_data): - """ - Writes the provided JSON to the database - :param json_data: - :return: - """ - if self.output: - print(json_data) - - try: - self.influx_client.write_points(json_data) - except (InfluxDBClientError, ConnectionError, InfluxDBServerError) as e: - if hasattr(e, 'code') and e.code == 404: - - print('Database {} Does Not Exist. Attempting To Create'.format(self.config.influx_database)) - - # TODO Grab exception here - self.influx_client.create_database(self.config.influx_database) - self.influx_client.write_points(json_data) - - return - - print('ERROR: Failed To Write To InfluxDB') - print(e) - - if self.output: - print('Written To Influx: {}'.format(json_data)) - - -def main(): - - parser = argparse.ArgumentParser(description="A tool to send Plex statistics to InfluxDB") - parser.add_argument('--config', default='config.ini', dest='config', help='Specify a custom location for the config file') - parser.add_argument('--singlerun', action='store_true', help='Only runs through once, does not keep monitoring') - args = parser.parse_args() - collector = InfluxdbSpeedtest(config=args.config, single_run=args.singlerun) - collector.run() - - -if __name__ == '__main__': - main() diff --git a/README.md b/README.md index 52e976554cf3f1d6cce20e582acf79cf50695ac4..479557e125c832a0b297bf670c813736870a9c77 100644 --- a/README.md +++ b/README.md @@ -5,18 +5,12 @@ This tool is a wrapper for speedtest-cli which allows you to run periodic speedtets and save the results to Influxdb -## Command Line Args ---config whatever.ini - Use a different ini file - ---singlerun - Run once and exit - ## Configuration within config.ini #### GENERAL |Key |Description | |:--------------|:-------------------------------------------------------------------------------------------------------------------| |Delay |Delay between runs | -|Output |Write console output while tool is running | #### INFLUXDB |Key |Description | |:--------------|:-------------------------------------------------------------------------------------------------------------------| @@ -28,7 +22,7 @@ This tool is a wrapper for speedtest-cli which allows you to run periodic speedt #### SPEEDTEST |Key |Description | |:--------------|:-------------------------------------------------------------------------------------------------------------------| -|Server |Server ID of speedtest.net server. Leave blank for auto | +|Server |Comma sperated list of servers. Leave blank for auto | @@ -37,8 +31,7 @@ This tool is a wrapper for speedtest-cli which allows you to run periodic speedt Before the first use run pip3 install -r requirements.txt Enter your desired information in config.ini and run InfluxdbSpeedtest.py - -Optionally, you can specify the --config argument to load the config file from a different location. + ***Requirements*** diff --git a/config.ini b/config.ini index 34c4c9f3a0555750b21038200356b670f8865071..d96c9ee650612fad521aac901895bc4881df574f 100644 --- a/config.ini +++ b/config.ini @@ -1,6 +1,5 @@ [GENERAL] Delay = 300 -Output = False [INFLUXDB] Address = @@ -12,4 +11,8 @@ Verify_SSL = False [SPEEDTEST] # Leave blank to auto pick server -Server = \ No newline at end of file +Server = + +[LOGGING] +# Valid Options: critical, error, warning, info, debug +Level = info \ No newline at end of file diff --git a/influxspeedtest.py b/influxspeedtest.py new file mode 100644 index 0000000000000000000000000000000000000000..4f1d6b8b78436c722218109cc0f84dbd4327418a --- /dev/null +++ b/influxspeedtest.py @@ -0,0 +1,8 @@ +import argparse + +from influxspeedtest.InfluxdbSpeedtest import InfluxdbSpeedtest + +parser = argparse.ArgumentParser(description="A tool to take network speed test and send the results to InfluxDB") +args = parser.parse_args() +collector = InfluxdbSpeedtest() +collector.run() diff --git a/influxspeedtest/InfluxdbSpeedtest.py b/influxspeedtest/InfluxdbSpeedtest.py new file mode 100644 index 0000000000000000000000000000000000000000..fb16179d1382371883eff4904d93749862e0f766 --- /dev/null +++ b/influxspeedtest/InfluxdbSpeedtest.py @@ -0,0 +1,170 @@ +import sys +import time + +import speedtest +from influxdb import InfluxDBClient +from influxdb.exceptions import InfluxDBClientError, InfluxDBServerError +from requests import ConnectTimeout + +from influxspeedtest.common import log +from influxspeedtest.config import config + + +class InfluxdbSpeedtest(): + + def __init__(self): + + self.influx_client = self._get_influx_connection() + self.speedtest = None + self.results = None + + def _get_influx_connection(self): + """ + Create an InfluxDB connection and test to make sure it works. + We test with the get all users command. If the address is bad it fails + with a 404. If the user doesn't have permission it fails with 401 + :return: + """ + + influx = InfluxDBClient( + config.influx_address, + config.influx_port, + database=config.influx_database, + ssl=config.influx_ssl, + verify_ssl=config.influx_verify_ssl, + username=config.influx_user, + password=config.influx_password, + timeout=5 + ) + try: + log.debug('Testing connection to InfluxDb using provided credentials') + influx.get_list_users() # TODO - Find better way to test connection and permissions + log.debug('Successful connection to InfluxDb') + except (ConnectTimeout, InfluxDBClientError) as e: + if isinstance(e, ConnectTimeout): + log.critical('Unable to connect to InfluxDB at the provided address (%s)', config.influx_address) + elif e.code == 401: + log.critical('Unable to connect to InfluxDB with provided credentials') + + sys.exit(1) + + return influx + + def setup_speedtest(self, server=None): + """ + Initializes the Speed Test client with the provided server + :param server: Int + :return: None + """ + speedtest.build_user_agent() + + log.debug('Setting up SpeedTest.net client') + + if server is None: + server = [] + else: + server = server.split() # Single server to list + + try: + self.speedtest = speedtest.Speedtest() + except speedtest.ConfigRetrievalError: + log.critical('Failed to get speedtest.net configuration. Aborting') + sys.exit(1) + + try: + self.speedtest.get_servers(server) + except speedtest.NoMatchedServers: + log.error('No matched servers: %s', server) + return + except speedtest.ServersRetrievalError: + log.critical('Cannot retrieve speedtest.net server list. Aborting') + sys.exit(1) + except speedtest.InvalidServerIDType: + log.error('%s is an invalid server type, must be int', server) + return + + log.debug('Picking the closest server') + + self.speedtest.get_best_server() + + log.info('Selected Server %s in %s', self.speedtest.best['id'], self.speedtest.best['name']) + + self.results = self.speedtest.results + + def send_results(self): + """ + Formats the payload to send to InfluxDB + :rtype: None + """ + result_dict = self.results.dict() + + input_points = [ + { + 'measurement': 'speed_test_results', + 'fields': { + 'download': result_dict['download'], + 'upload': result_dict['upload'], + 'ping': result_dict['server']['latency'] + }, + 'tags': { + 'server': result_dict['server']['id'] + } + } + ] + + self.write_influx_data(input_points) + + def run_speed_test(self, server=None): + """ + Performs the speed test with the provided server + :param server: Server to test against + """ + log.info('Starting Speed Test For Server %s', server) + self.setup_speedtest(server) + log.info('Starting download test') + self.speedtest.download() + log.info('Starting upload test') + self.speedtest.upload() + self.send_results() + + results = self.results.dict() + log.info('Download: %sMbps - Upload: %sMbps - Latency: %sms', + round(results['download'] / 1000000, 2), + round(results['upload'] / 1000000, 2), + results['server']['latency'] + ) + + + + def write_influx_data(self, json_data): + """ + Writes the provided JSON to the database + :param json_data: + :return: None + """ + log.debug(json_data) + + try: + self.influx_client.write_points(json_data) + except (InfluxDBClientError, ConnectionError, InfluxDBServerError) as e: + if hasattr(e, 'code') and e.code == 404: + log.error('Database %s Does Not Exist. Attempting To Create', config.influx_database) + self.influx_client.create_database(config.influx_database) + self.influx_client.write_points(json_data) + return + + log.error('Failed To Write To InfluxDB') + print(e) + + log.debug('Data written to InfluxDB') + + def run(self): + + while True: + if not config.servers: + self.run_speed_test() + else: + for server in config.servers: + self.run_speed_test(server) + log.info('Waiting %s seconds until next test', config.delay) + time.sleep(config.delay) diff --git a/influxspeedtest/__init__.py b/influxspeedtest/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..e69de29bb2d1d6434b8b29ae775ad8c2e48c5391 diff --git a/influxspeedtest/common/__init__.py b/influxspeedtest/common/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..cbee26db5481c05d3fd4ac717e45b46ef0e5ac05 --- /dev/null +++ b/influxspeedtest/common/__init__.py @@ -0,0 +1 @@ +from .utils import log \ No newline at end of file diff --git a/influxspeedtest/common/logfilters.py b/influxspeedtest/common/logfilters.py new file mode 100644 index 0000000000000000000000000000000000000000..bd015b3b2c9066da87dcb7ff9d167188cf5c69d8 --- /dev/null +++ b/influxspeedtest/common/logfilters.py @@ -0,0 +1,13 @@ +import logging + + +class SingleLevelFilter(logging.Filter): + def __init__(self, passlevel, above=True): + self.passlevel = passlevel + self.above = above + + def filter(self, record): + if self.above: + return record.levelno >= self.passlevel + else: + return record.levelno <= self.passlevel \ No newline at end of file diff --git a/influxspeedtest/common/utils.py b/influxspeedtest/common/utils.py new file mode 100644 index 0000000000000000000000000000000000000000..a5cc12031a6c1f15df5f10d9854124da73c57e7d --- /dev/null +++ b/influxspeedtest/common/utils.py @@ -0,0 +1,23 @@ +import logging +import sys + +from influxspeedtest.common.logfilters import SingleLevelFilter +from influxspeedtest.config import config + +log = logging.getLogger(__name__) +log.setLevel(config.logging_level) +formatter = logging.Formatter('%(asctime)s - %(levelname)s: %(message)s') + +general_handler = logging.StreamHandler(sys.stdout) +general_filter = SingleLevelFilter(logging.INFO, False) +general_handler.setFormatter(formatter) +general_handler.addFilter(general_filter) +log.addHandler(general_handler) + +error_handler = logging.StreamHandler(sys.stderr) +error_filter = SingleLevelFilter(logging.WARNING) +error_handler.setFormatter(formatter) +error_handler.addFilter(error_filter) +log.addHandler(error_handler) + +log.propagate = False \ No newline at end of file diff --git a/influxspeedtest/config/__init__.py b/influxspeedtest/config/__init__.py new file mode 100644 index 0000000000000000000000000000000000000000..7f57bfafecd08b073b5cfbddc04af2d59f94f85c --- /dev/null +++ b/influxspeedtest/config/__init__.py @@ -0,0 +1,10 @@ +import os + +from influxspeedtest.config.configmanager import ConfigManager + +if os.getenv('influxspeedtest'): + config = os.getenv('influxspeedtest') +else: + config = 'config.ini' + +config = ConfigManager(config) \ No newline at end of file diff --git a/influxspeedtest/config/configmanager.py b/influxspeedtest/config/configmanager.py new file mode 100644 index 0000000000000000000000000000000000000000..cc448062b733b4900c3e91794caa9637e0e510e5 --- /dev/null +++ b/influxspeedtest/config/configmanager.py @@ -0,0 +1,43 @@ +import configparser +import os +import sys + + +class ConfigManager(): + + def __init__(self, config): + print('Loading Configuration File {}'.format(config)) + self.servers = [] + config_file = os.path.join(os.getcwd(), config) + if os.path.isfile(config_file): + self.config = configparser.ConfigParser() + self.config.read(config_file) + else: + print('ERROR: Unable To Load Config File: {}'.format(config_file)) + sys.exit(1) + + self._load_config_values() + print('Configuration Successfully Loaded') + + def _load_config_values(self): + + # General + self.delay = self.config['GENERAL'].getint('Delay', fallback=2) + + # InfluxDB + self.influx_address = self.config['INFLUXDB']['Address'] + self.influx_port = self.config['INFLUXDB'].getint('Port', fallback=8086) + self.influx_database = self.config['INFLUXDB'].get('Database', fallback='speedtests') + self.influx_user = self.config['INFLUXDB'].get('Username', fallback='') + self.influx_password = self.config['INFLUXDB'].get('Password', fallback='') + self.influx_ssl = self.config['INFLUXDB'].getboolean('SSL', fallback=False) + self.influx_verify_ssl = self.config['INFLUXDB'].getboolean('Verify_SSL', fallback=True) + + # Logging + self.logging_level = self.config['LOGGING'].get('Level', fallback='debug') + self.logging_level = self.logging_level.upper() + + # Speedtest + test_server = self.config['SPEEDTEST'].get('Server', fallback=None) + if test_server: + self.servers = test_server.split(',') diff --git a/requirements.txt b/requirements.txt index 0cce7ddc53fc2b6b9d8ffd1ad86d09354d7a9fdf..b61200b82e6e2c509f4898fd1af56b78d1ba2e7e 100644 --- a/requirements.txt +++ b/requirements.txt @@ -1,2 +1,3 @@ influxdb -speedtest-cli \ No newline at end of file +speedtest-cli +requests \ No newline at end of file