#!/usr/bin/python
#
# tcplatencyc2d.py
#
# (c) 2013,2015,2018 Todd Shadburn <tshadburn@opencomputingsolutions.com>
#
# Licensed under the GNU GPL v2
#


import sys
import time
import datetime
import getpass
import pickle
import socket
import select
from random import choice
from string import lowercase
import netinfo



def submit_events(ev_list):
	nc = None
	for ev in ev_list:
		#print 'DEBUG sending event'
		try:
			nc.send_event(ev)
		except:
			try:
				nc.close()
			except:
				pass
			nc = None
			nc = netinfo.NetinfoClient(
				host='mon.home.opencomputingsolutions.com',
				#cert='/etc/netinfo/netinfod.cert',
				#key='/etc/netinfo/netinfod.key')
				cert='server.cert',
				key='server.key')
			try:
				# Try to resend event
				nc.send_event(ev)
			except:
				pass
	return



##########################################################################
# main
##########################################################################

# Globals
interval = 300
timeout = 5
host = None
port = 5012
packet_data = 512
warntime = 0.010
crittime = 0.015
conffile = '/etc/tcplatencycd.conf'
service_list = []
console = 0
count = 1
stayalive = 0

# Parse and validate arguments
import argparse
parser = argparse.ArgumentParser(description='Utility for bulk testing of TCP connectivity.')
parser.add_argument('--host', nargs=1, help='(Required)Host to connect to.')
parser.add_argument('--port', nargs=1, help='(Required)Port to connect to.')
parser.add_argument('--len', nargs=1, help='(Optional)Packet data bytes to send.')
parser.add_argument('--interval', nargs=1, help='(Optional)Interval between checks in seconds (default=300).')
parser.add_argument('--timeout', nargs=1, help='(Optional)Connection timeout in seconds (default=5).')
parser.add_argument('--warntime', nargs=1, help='(Optional)Warning threshold in seconds (default=0.010).')
parser.add_argument('--crittime', nargs=1, help='(Optional)Critical threshold in seconds (default=0.030).')
parser.add_argument('--console', nargs=1, help='(Optional)Output results to console rather than netinfo.')
parser.add_argument('--count', nargs=1, help='(Optional)Number of packets to send in each interval.')
parser.add_argument('--conffile', nargs=1, help='(Optional)Alternate configuration file to use (default=/etc/tcplatencycd.conf).')
parser.add_argument('--stayalive', nargs=1, help='(Optional)Keep the connections open rather than closing them each interval.')
args = parser.parse_args()
if not args.port is None:
	port = int(args.port[0])
if not args.len is None:
	packet_data = int(args.len[0])
if not args.interval is None:
	interval = int(args.interval[0])
if not args.timeout is None:
	timeout = args.timeout[0]	
if not args.warntime is None:
	warntime = float(args.warntime[0])
if not args.crittime is None:
	crittime = float(args.crittime[0])
if not args.conffile is None:
	conffile = args.conffile[0]
if not args.host is None:
	service_list.append({'host': args.host[0], 'port': port, 'lastrc': 4, 'socket': None})
if not args.console is None:
	console = float(args.console[0])
if not args.count is None:
	count = int(args.count[0])
if not args.stayalive is None:
	stayalive = 1

# parse the config file (if available)
try:
	fd = open(conffile, 'r')
	for ln in fd:
		h = None
		p = None
		if len(ln) < 2 or ln[0] == '#':
			# blank line or comment line
			continue
		if ':' in ln:
			h,p = ln.strip(' \r\n').split(':', 1)
			p = int(p)
		else:
			h = ln.strip(' \r\n')
			p = port
		service_list.append({'host': h, 'port': p, 'lastrc': 4, 'socket': None})
except:
	# No config file or error reading it
	pass
if not len(service_list):
	print 'Error: host argument or configuration file is required.'
	sys.exit(2)

nc = None
data = "".join(choice(lowercase) for i in range(packet_data))
# Echo server on each port
while 1:
	st = 0
	loop_start = datetime.datetime.now()
	ne_list = []

	# Test each service
	for svc in service_list:
		host = svc['host']
		port = svc['port']
		objectname = 'TCP-%s:%d' % (host,port)
		rc = 0
		monmsg = ''
		ndata = ''

		if not svc['socket']:
			try:
				svc['socket'] = socket.create_connection((host, port))
				monmsg = 'Connection to %s:%d created.' % (host,port)
			except:
				monmsg = 'Connection to %s:%d failed. Service may be down.' % (host,port)
				rc = 2
				if console > 0:
					print monmsg

		st = datetime.datetime.now()

		if svc['socket']:
			try:
				for i in range(count):
					svc['socket'].send(data)
					ndata = svc['socket'].recv(1024)
			except:
				try:
					svc['socket'].close()
				except:
					pass
				svc['socket'] = None
				monmsg = 'I/O failed to %s:%d. Connection interrupted.' % (host,port)
				rc = 2
				if console > 0:
					print monmsg

		if stayalive == 0:
			if svc['socket']:
				svc['socket'].close()
				svc['socket'] = None
			
		et = datetime.datetime.now()


		if rc != 2:
			diff = (et - st)
			dt = diff.seconds + (diff.microseconds/1000000.0)
			dt /= float(count)
			msg = 'Roundtrip to %s:%d elapsed: %0.6f' % (host,port,dt)
			perf = 'gauge=%0.6f;0.0;0.0;%0.6f;%0.6f' % (dt,warntime,crittime)
			monmsg = str('|'.join([msg,perf]))
			if console > 0:
				print msg
			if dt > warntime:
				rc = 1
			if dt > crittime:
				rc = 2
			# performance event
			ne_list.append(netinfo.NetinfoEvent(
				object=str(objectname),
				flags=0,
				rc=0,
				data=str(monmsg)))
		if rc != svc['lastrc']:
			# state change event
			ne_list.append(netinfo.NetinfoEvent(
				object=str(objectname),
				flags=1,
				rc=rc,
				data=str(monmsg)))
			svc['lastrc'] = rc

	if console == 0:
		submit_events(ne_list)

	# Sleep for interval
	loop_end = datetime.datetime.now()
	loop_diff = loop_end-loop_start
	loop_sec = loop_diff.seconds + (loop_diff.microseconds/1000000.0)
	#print 'DEBUG sleep(%0.6f)' % (float(interval)-loop_sec)
	if console > 0:
		time.sleep(float(console)-loop_sec)
	else:
		time.sleep(float(interval)-loop_sec)
sys.exit(0)