routerstats/routerstats_client.py

#!/usr/bin/env python3

import socket
import logging
import time
import os

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):
	logging.error('NotConnected error')
	pass

class routerstats_client():
	def __init__(self, host, port):
		self.host = host
		self.port = port
		self.connected = False

	def connect(self):
		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 ' + str(self.host))
			self.connected = True
		except ConnectionRefusedError as error:
			logging.error('Could not connect to ' + str(self.host) + ':' + str(self.port))
			raise ConnectionRefusedError(error)

	def send(self, tosend):
		if self.connected:
			logging.debug('Sending ' + str(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: ' + str(error))
				self.connected = False
				raise NotConnected
			try:
				self.sock.getpeername()
			except OSError as error:
				logging.error('Could not send to server: ' + str(error))
				self.connected = False
				raise NotConnected
		else:
			logging.error('Not connected to server')
			raise NotConnected('Not connected to server')

	def recv(self):
		if self.connected != 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: ' + str(line))
						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 OSError as error:
				if str(error) == 'timed out':
					#This is expected, as it is how we loop
					break
				logging.error('OSError: ' + str(error))
				logging.debug('Got: ' + str(line))
				self.connected = False
				raise NotConnected(error)
			except TimeoutError:
				logging.error('Timeout while fetching data, got: ' + str(line))
				break

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 ' + str(received) + ' into two distinct values')
		return
	try:
		timestamp = int(timestamp)
	except ValueError:
		logging.error('Could not parse ' + str(timestamp) + ' as an int')
		return
	toret = {'timestamp': timestamp, 'net_dnat': 0, 'loc-net': 0}
	try:
		toret[zone] += 1
	except KeyError as error:
		logging.debug('Ignoring zone: ' + str(error))
	logging.debug('Parsed to: ' + str(toret))
	return toret

def handle(received, client):
	'''Do things with whatever came from the server'''
	if not received:
		return
	if received == 'ping':
		client.send('pong')
		return True
	elif received is not None:
		if received:
			return handle_server_msg(received)

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) != True:
			self.create()
		self.toupdate = {'timestamp': None, 'net_dnat': 0, 'loc-net': 0}
		self.freshdict = self.toupdate.copy()

	def create(self):
		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 ' + str(self.rrdfile))

	def push(self):
		#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
		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:
			logging.error(str(error))
			return False

	def add(self, input_dict):
		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']:
			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? ' + str(input_dict) + str(self.toupdate))

def main():
	rrdfile = 'test.rrd'
	rrdupdater = UpdateRRD(rrdfile)
	client = routerstats_client('127.0.0.1', 9999)
	while True:
		try:
			client.connect()
			break
		except ConnectionRefusedError:
			time.sleep(1)

	tries = 0
	loops = 0
	while True:
		try:
			retval = handle(client.recv(), client)
			if retval:
				if retval == True:
					loops = 0
				else:
					rrdupdater.add(retval)
			else:
				rrdupdater.push()
				loops += 1
				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:
				logging.debug('ConnectionRefused')
				tries += 1
				if tries >= 5:
					tries = 5
				time.sleep(tries)
		except ConnectionRefusedError:
			logging.debug('ConnectionRefused')
			time.sleep(1)

if __name__ == '__main__':
	main()