import configparser
import os
import sys
import argparse
import time
import logging
import re
import socket
from influxdb import InfluxDBClient
from influxdb.exceptions import InfluxDBClientError, InfluxDBServerError
# TODO Move urlopen login in each method call to one central method
# TODO Validate that we get a valid URL from config
# TODO Deal with multiple trackers per torrent
__author__ = 'barry'
class configManager():
#TODO Validate given client url
def __init__(self, silent, config):
self.valid_torrent_clients = ['deluge', 'utorrent', 'rtorrent']
self.valid_log_levels = {
'DEBUG': 0,
'INFO': 1,
'WARNING': 2,
'ERROR': 3,
'CRITICAL': 4
}
self.silent = silent
if not self.silent:
print('Loading Configuration File {}'.format(config))
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()
self._validate_logging_level()
self._validate_torrent_client()
if not self.silent:
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)
self.hostname = self.config['GENERAL'].get('Hostname')
if not self.hostname:
self.hostname = socket.gethostname()
# 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 = self.config['LOGGING'].getboolean('Enable', fallback=False)
self.logging_level = self.config['LOGGING']['Level'].upper()
self.logging_file = self.config['LOGGING']['LogFile']
self.logging_censor = self.config['LOGGING'].getboolean('CensorLogs', fallback=True)
self.logging_print_threshold = self.config['LOGGING'].getint('PrintThreshold', fallback=2)
# TorrentClient
self.tor_client = self.config['TORRENTCLIENT'].get('Client', fallback=None).lower()
self.tor_client_user = self.config['TORRENTCLIENT'].get('Username', fallback=None)
self.tor_client_password = self.config['TORRENTCLIENT'].get('Password', fallback=None)
self.tor_client_url = self.config['TORRENTCLIENT'].get('Url', fallback=None)
def _validate_torrent_client(self):
if self.tor_client not in self.valid_torrent_clients:
print('ERROR: {} Is Not a Valid or Support Torrent Client. Aborting'.format(self.tor_client))
sys.exit(1)
def _validate_logging_level(self):
"""
Make sure we get a valid logging level
:return:
"""
if self.logging_level in self.valid_log_levels:
self.logging_level = self.logging_level.upper()
return
else:
if not self.silent:
print('Invalid logging level provided. {}'.format(self.logging_level))
print('Logging will be disabled')
self.logging = None
class influxdbSeedbox():
def __init__(self, config=None, silent=None):
self.config = configManager(silent, config=config)
self.output = self.config.output
self.logger = None
self.delay = self.config.delay
self.influx_client = InfluxDBClient(
self.config.influx_address,
self.config.influx_port,
database=self.config.influx_database,
ssl=self.config.influx_ssl,
verify_ssl=self.config.influx_verify_ssl
)
self._set_logging()
if self.config.tor_client == 'deluge':
from clients.deluge import DelugeClient
if self.output:
print('Generating Deluge Client')
self.tor_client = DelugeClient(self.send_log,
username=self.config.tor_client_user,
password=self.config.tor_client_password,
url=self.config.tor_client_url,
hostname=self.config.hostname)
elif self.config.tor_client == 'utorrent':
from clients.utorrent import UTorrentClient
if self.output:
print('Generating uTorrent Client')
self.tor_client = UTorrentClient(self.send_log,
username=self.config.tor_client_user,
password=self.config.tor_client_password,
url=self.config.tor_client_url,
hostname=self.config.hostname)
elif self.config.tor_client == 'rtorrent':
from clients.rtorrent import rTorrentClient
if self.output:
print('Generating rTorrent Client')
self.tor_client = rTorrentClient(self.send_log,
username=None,
password=None,
url=self.config.tor_client_url,
hostname=self.config.hostname)
def _set_logging(self):
"""
Create the logger object if enabled in the config
:return: None
"""
if self.config.logging:
if self.output:
print('Logging is enabled. Log output will be sent to {}'.format(self.config.logging_file))
self.logger = logging.getLogger(__name__)
self.logger.setLevel(self.config.logging_level)
formatter = logging.Formatter('%(asctime)s %(levelname)s: %(message)s')
fhandle = logging.FileHandler(self.config.logging_file)
fhandle.setFormatter(formatter)
self.logger.addHandler(fhandle)
def send_log(self, msg, level):
"""
Used as a shim to write log messages. Allows us to sanitize input before logging
:param msg: Message to log
:param level: Level to log message at
:return: None
"""
if not self.logger:
return
if self.output and self.config.valid_log_levels[level.upper()] >= self.config.logging_print_threshold:
print(msg)
# Make sure a good level was given
if not hasattr(self.logger, level):
self.logger.error('Invalid log level provided to send_log')
return
output = self._sanitize_log_message(msg)
log_method = getattr(self.logger, level)
log_method(output)
def _sanitize_log_message(self, msg):
"""
Take the incoming log message and clean and sensitive data out
:param msg: incoming message string
:return: cleaned message string
"""
if not self.config.logging_censor:
return msg
# Remove server addresses
msg = msg.replace(self.config.tor_client_url, 'http://*******:8112/json')
# Remove IP addresses
for match in re.findall(r"\b\d{1,3}\.\d{1,3}\.\d{1,3}\.\d{1,3}\b", msg):
msg = msg.replace(match, '***.***.***.***')
return msg
def write_influx_data(self, json_data):
"""
Writes the provided JSON to the database
:param json_data:
:return:
"""
self.send_log(json_data, 'info')
# TODO This bit of fuckery may turn out to not be a good idea.
"""
The idea is we only write 1 series at a time but we can get passed a bunch of series. If the incoming list
has more than 1 thing, we know it's a bunch of points to write. We loop through them and recursively call this
method which each series. The recursive call will only have 1 series so it passes on to be written to Influx.
I know, brilliant right? Probably not
"""
if len(json_data) > 1:
for series in json_data:
self.write_influx_data(series)
return
try:
self.influx_client.write_points(json_data)
except (InfluxDBClientError, ConnectionError, InfluxDBServerError) as e:
if hasattr(e, 'code') and e.code == 404:
msg = 'Database {} Does Not Exist. Attempting To Create'.format(self.config.influx_database)
self.send_log(msg, 'error')
# TODO Grab exception here
self.influx_client.create_database(self.config.influx_database)
self.influx_client.write_points(json_data)
return
self.send_log('Failed to write data to InfluxDB', 'error')
print('ERROR: Failed To Write To InfluxDB')
print(e)
self.send_log('Written To Influx: {}'.format(json_data), 'debug')
def run(self):
while True:
self.tor_client.get_all_torrents()
torrent_json = self.tor_client.process_torrents()
if torrent_json:
self.write_influx_data(torrent_json)
#self.tor_client.get_active_plugins()
tracker_json = self.tor_client.process_tracker_list()
if tracker_json:
self.write_influx_data(tracker_json)
time.sleep(self.delay)
def main():
parser = argparse.ArgumentParser(description="A tool to send Torrent Client statistics to InfluxDB")
parser.add_argument('--config', default='config.ini', dest='config', help='Specify a custom location for the config file')
# Silent flag allows output prior to the config being loaded to also be suppressed
parser.add_argument('--silent', action='store_true', help='Surpress All Output, regardless of config settings')
args = parser.parse_args()
monitor = influxdbSeedbox(silent=args.silent, config=args.config)
monitor.run()
if __name__ == '__main__':
main()