400 lines
11 KiB
Python
Executable File
400 lines
11 KiB
Python
Executable File
#!/usr/bin/env python3
|
|
'''Client for the routerstats collection'''
|
|
|
|
import socket
|
|
import logging
|
|
import time
|
|
import os
|
|
import sys
|
|
import argparse
|
|
import configparser
|
|
|
|
import rrdtool
|
|
|
|
logging.basicConfig(
|
|
format='%(asctime)s %(funcName)20s %(levelname)-8s %(message)s',
|
|
level=logging.INFO,
|
|
datefmt='%Y-%m-%d %H:%M:%S')
|
|
|
|
class NotConnected(Exception):
|
|
'''Raise to make it known that you are not connected'''
|
|
|
|
class RouterstatsClient():
|
|
'''Client itself'''
|
|
def __init__(self, host, port, passwd_file):
|
|
self.host = host
|
|
self.port = port
|
|
self.connected = False
|
|
self.sock = None
|
|
self.passwd_file = passwd_file
|
|
self.received_lines = 0
|
|
|
|
def connect(self):
|
|
'''Do connect'''
|
|
self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
|
|
try:
|
|
self.sock.connect((self.host, self.port))
|
|
self.sock.settimeout(1)
|
|
logging.info('Connected to %s', self.host)
|
|
self.login()
|
|
self.connected = True
|
|
self.received_lines = 0
|
|
except ConnectionRefusedError as error:
|
|
logging.error('Could not connect to %s:%s', self.host, self.port)
|
|
raise ConnectionRefusedError from error
|
|
except ConnectionError as error:
|
|
logging.error('Could not connect to %s:%s: %s',
|
|
self.host,
|
|
self.port,
|
|
error)
|
|
except TimeoutError as error:
|
|
logging.error('Timed out')
|
|
|
|
def login(self):
|
|
'''Do the login routine.
|
|
Wait for server to greet us with "hello"
|
|
reply with the password
|
|
wait for "Welcome"
|
|
'''
|
|
|
|
logging.debug('Logging in')
|
|
try:
|
|
hello = self.sock.recv(5)
|
|
except TimeoutError as exception:
|
|
logging.error('Timeout while waiting for Hello')
|
|
raise ConnectionError('Timed out waiting for Hello when connecting') from exception
|
|
if not hello == b'Hello':
|
|
logging.error('No Hello from server: %s', hello)
|
|
raise ConnectionError('Server did not greet us with Hello during login')
|
|
with open(self.passwd_file, 'r', encoding='utf-8') as passwd_file:
|
|
passwd = passwd_file.readline()
|
|
passwd = passwd.rstrip()
|
|
logging.debug('Sending password: %s', passwd)
|
|
self.sock.send(passwd.encode('utf-8'))
|
|
try:
|
|
response = self.sock.recv(7)
|
|
if not response == b'Welcome':
|
|
logging.error('Not Welcome: %s', response)
|
|
raise ConnectionError('We are not greeted with Welcome after sending password')
|
|
return True
|
|
except TimeoutError as exception:
|
|
raise ConnectionError(
|
|
'Timed out while waiting for server to greet us after sending password') from exception
|
|
|
|
def send(self, tosend):
|
|
'''Send some data'''
|
|
if self.connected:
|
|
logging.debug('Sending %s', tosend)
|
|
send_bytes = bytes(tosend + "\n", 'utf-8')
|
|
try:
|
|
self.sock.send(send_bytes)
|
|
except OSError as error:
|
|
logging.error('Cound not send to server: %s', error)
|
|
self.connected = False
|
|
raise NotConnected from error
|
|
try:
|
|
self.sock.getpeername()
|
|
except OSError as error:
|
|
logging.error('Could not send to server: %s', error)
|
|
self.connected = False
|
|
raise NotConnected from error
|
|
else:
|
|
logging.error('Not connected to server')
|
|
raise NotConnected('Not connected to server')
|
|
|
|
def recv(self):
|
|
'''Receive some data'''
|
|
if self.connected is not True:
|
|
logging.error('Trying to recv when not connected')
|
|
raise NotConnected
|
|
line = b''
|
|
blanks = 0
|
|
while True:
|
|
try:
|
|
toret = self.sock.recv(1)
|
|
if toret:
|
|
line += toret
|
|
if blanks > 0:
|
|
blanks -= 1
|
|
if line.endswith(b'\n'):
|
|
#We're done for now, returning value
|
|
line = line.strip().decode('utf-8')
|
|
logging.debug('Received from server: %s', line)
|
|
self.received_lines += 1
|
|
if not self.received_lines % 1000:
|
|
logging.info('Received %s lines', self.received_lines)
|
|
return line
|
|
else:
|
|
blanks += 1
|
|
if blanks >= 5:
|
|
blanks = 0
|
|
# break
|
|
logging.debug('Too many blank reads, and still no complete line')
|
|
self.connected = False
|
|
raise NotConnected
|
|
except TimeoutError:
|
|
if line == b'':
|
|
break
|
|
logging.error('Timeout while fetching data, got: %s' ,line)
|
|
break
|
|
except OSError as error:
|
|
if str(error) == 'timed out':
|
|
#This is expected, as it is how we loop
|
|
break
|
|
logging.error('OSError: %s', error)
|
|
logging.debug('Got: %s', line)
|
|
self.connected = False
|
|
raise NotConnected(error) from error
|
|
return None
|
|
|
|
def handle_server_msg(received: str):
|
|
'''Do something about something from the server'''
|
|
try:
|
|
timestamp, zone = received.split()
|
|
except ValueError:
|
|
logging.error('Could not parse %s into two distinct values', received)
|
|
return None
|
|
try:
|
|
timestamp = int(timestamp)
|
|
except ValueError:
|
|
logging.error('Could not parse %s as an int', timestamp)
|
|
return None
|
|
toret = {'timestamp': timestamp, 'net_dnat': 0, 'loc-net': 0}
|
|
try:
|
|
toret[zone] += 1
|
|
except KeyError as error:
|
|
logging.debug('Ignoring zone: %s', error)
|
|
logging.debug('Parsed to: %s', toret)
|
|
return toret
|
|
|
|
def handle(received, client):
|
|
'''Do things with whatever came from the server'''
|
|
if not received:
|
|
return None
|
|
if received == 'ping':
|
|
client.send('pong')
|
|
return True
|
|
if received is not None:
|
|
if received:
|
|
return handle_server_msg(received)
|
|
return None
|
|
|
|
class UpdateRRD:
|
|
'''Handle updates to the rrd-file, since we cannot update more than once every second'''
|
|
def __init__(self, rrdfile):
|
|
self.rrdfile = rrdfile
|
|
if os.path.isfile(rrdfile) is not True:
|
|
self.create()
|
|
self.toupdate = {'timestamp': None, 'net_dnat': 0, 'loc-net': 0}
|
|
self.freshdict = self.toupdate.copy()
|
|
|
|
def create(self):
|
|
'''Create self.rrdfile'''
|
|
rrdtool.create(
|
|
self.rrdfile,
|
|
"--start",
|
|
"1000000000",
|
|
"--step",
|
|
"10",
|
|
"DS:net_dnat:ABSOLUTE:10:0:U",
|
|
"DS:loc-net:ABSOLUTE:10:0:U",
|
|
"RRA:AVERAGE:0.5:1:1d",
|
|
"RRA:AVERAGE:0.5:1:1M")
|
|
logging.debug('Created rrdfile %s', self.rrdfile)
|
|
|
|
def push(self):
|
|
'''Send data to rrdtool'''
|
|
#Make sure we're not doing something stupid..
|
|
info = rrdtool.info(self.rrdfile)
|
|
if self.toupdate['timestamp'] is None:
|
|
return False
|
|
if info['last_update'] > self.toupdate['timestamp']:
|
|
logging.error(
|
|
'Trying to update when rrdfile is newer than our timestamp. Ignoring line and resetting.')
|
|
self.toupdate = self.freshdict.copy()
|
|
return False
|
|
if info['last_update'] == self.toupdate['timestamp']:
|
|
logging.error('last update and toupdate timestamp are the same, this should not happen')
|
|
self.toupdate = self.freshdict.copy()
|
|
return False
|
|
try:
|
|
rrdtool.update(
|
|
self.rrdfile,
|
|
str(self.toupdate['timestamp'])
|
|
+ ':'
|
|
+ str(self.toupdate['net_dnat'])
|
|
+ ':'
|
|
+ str(self.toupdate['loc-net']))
|
|
self.toupdate = self.freshdict.copy()
|
|
logging.debug('Updated rrdfile')
|
|
return True
|
|
except rrdtool.OperationalError as error:
|
|
if str(error) == 'could not lock RRD':
|
|
return False
|
|
logging.error(str(error))
|
|
return False
|
|
|
|
def add(self, input_dict):
|
|
'''Add data to be added later'''
|
|
if self.toupdate['timestamp'] is None:
|
|
#Never touched, just overwrite with whatever we got
|
|
self.toupdate = input_dict
|
|
elif input_dict['timestamp'] > self.toupdate['timestamp']:
|
|
#What we get is fresher than what we have, and noone has done a push yet
|
|
self.push()
|
|
self.toupdate = input_dict
|
|
elif input_dict['timestamp'] == self.toupdate['timestamp']:
|
|
#Same timestamp, just up values and be happy
|
|
self.toupdate['net_dnat'] += input_dict['net_dnat']
|
|
self.toupdate['loc-net'] += input_dict['loc-net']
|
|
elif input_dict['timestamp'] < self.toupdate['timestamp']:
|
|
diff = self.toupdate['timestamp'] - input_dict['timestamp']
|
|
if diff <= 5:
|
|
#Might be because something is a bit slow sometimes?
|
|
input_dict['timestamp'] = self.toupdate['timestamp']
|
|
else:
|
|
logging.error('Newly fetched data is older than what we have in the queue already. Passing.')
|
|
else:
|
|
logging.error('Not sure what to do here? %s: %s', input_dict, self.toupdate)
|
|
|
|
def initial_connect(client):
|
|
'''Loop here until first connection is made'''
|
|
while True:
|
|
try:
|
|
client.connect()
|
|
break
|
|
except ConnectionRefusedError:
|
|
time.sleep(1)
|
|
|
|
def loop(client, rrdupdater):
|
|
'''Main loop'''
|
|
initial_connect(client)
|
|
|
|
tries = 0
|
|
loops = 0
|
|
while True:
|
|
try:
|
|
retval = handle(client.recv(), client)
|
|
if retval:
|
|
loops = 0
|
|
if retval is True:
|
|
pass
|
|
else:
|
|
rrdupdater.add(retval)
|
|
else:
|
|
loops += 1
|
|
if loops >= 5:
|
|
#Want to wait until at least 5 secs have passed since last actual data fetch..
|
|
rrdupdater.push()
|
|
if loops >= 60:
|
|
logging.error(
|
|
'No data in 60 seconds. We expect a ping/pong every 30. Lost connection, probably')
|
|
loops = 0
|
|
raise NotConnected
|
|
except NotConnected:
|
|
try:
|
|
client.connect()
|
|
tries = 0
|
|
except ConnectionRefusedError as error:
|
|
logging.debug('%s', error)
|
|
tries += 1
|
|
tries = min(tries, 5)
|
|
time.sleep(tries)
|
|
except ConnectionRefusedError as error:
|
|
logging.debug('%s', error)
|
|
time.sleep(1)
|
|
except KeyboardInterrupt:
|
|
break
|
|
|
|
def main():
|
|
'''Main is main'''
|
|
config_section = 'client'
|
|
|
|
config = configparser.ConfigParser()
|
|
parser = argparse.ArgumentParser(exit_on_error=False, add_help=False)
|
|
|
|
parser.add_argument('-c', '--config', help='config file to load')
|
|
parser.add_argument('-d', '--debug', action='store_true', help='enable debug')
|
|
|
|
args, _ = parser.parse_known_args()
|
|
|
|
if args.debug:
|
|
logging.root.setLevel(logging.DEBUG)
|
|
|
|
logging.debug('Starting as PID %s', os.getpid())
|
|
|
|
found = False
|
|
config_dirs = ('/etc/routerstats/', '/usr/local/etc/routerstats/', '/opt/routerstats/', './')
|
|
if args.config:
|
|
if os.path.isfile(args.config):
|
|
config.read(args.config)
|
|
found = True
|
|
else:
|
|
logging.error('Specified config file does not exist: %s', args.config)
|
|
else:
|
|
logging.debug('Trying to find config')
|
|
#Try to find in "usual" places
|
|
for directory in config_dirs:
|
|
trytoread = directory + 'routerstats.config'
|
|
if os.path.isfile(trytoread):
|
|
logging.debug('Reading config file %s', trytoread)
|
|
config.read(trytoread)
|
|
found = True
|
|
if not found:
|
|
logging.error('routerstats.config not found in %s', config_dirs)
|
|
sys.exit(0)
|
|
|
|
parser.add_argument(
|
|
'-r',
|
|
'--host',
|
|
dest='host',
|
|
help='ip/hostname to connect to',
|
|
default=config[config_section]['host'])
|
|
parser.add_argument(
|
|
'-p',
|
|
'--port',
|
|
dest='port',
|
|
type=int,
|
|
help='port to connect to',
|
|
default=config[config_section]['port'])
|
|
parser.add_argument(
|
|
'-v',
|
|
'--vardir',
|
|
dest='vardir',
|
|
help='directory for storing rrd file',
|
|
default=config[config_section]['var_dir'])
|
|
parser.add_argument(
|
|
'-w',
|
|
'--pwdfile',
|
|
dest='passwd_file',
|
|
help='password file',
|
|
default=config[config_section]['passwd_file'])
|
|
parser.add_argument(
|
|
'-h',
|
|
'--help',
|
|
help='show this help and exit',
|
|
action='store_true',
|
|
default=False)
|
|
|
|
args = parser.parse_args()
|
|
logging.debug(args)
|
|
|
|
if args.help:
|
|
parser.print_help()
|
|
sys.exit()
|
|
|
|
#Make sure the file specified is to be found..
|
|
if not os.path.isfile(args.passwd_file):
|
|
logging.error('Cannot find passwd-file %s', args.passwd_file)
|
|
sys.exit()
|
|
|
|
if not os.path.isdir(args.vardir):
|
|
logging.error('Cannot find var dir %s', args.vardir)
|
|
|
|
rrdupdater = UpdateRRD(args.vardir + '/routerstats.rrd')
|
|
client = RouterstatsClient(args.host, args.port, args.passwd_file)
|
|
loop(client, rrdupdater)
|
|
|
|
if __name__ == '__main__':
|
|
main()
|