#!/usr/bin/python
#
# tcptester.py
#
# (c) 2013,2018 Todd Shadburn <tshadburn@opencomputingsolutions.com>
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 59 Temple Place - Suite 330, Boston, MA 02111-1307, USA.
#


import sys
import getpass
import pickle
import socket


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

# Globals
timeout = 5
infiles = None
outfile = None
old_python = 0


# Parse and validate arguments
argmod = ''
try:
	import argparse
	argmod = 'argparse'
except:
	try:
		from optparse import OptionParser
		argmod = 'optparse'
	except:
		print 'No argparse or optparse modules!'
		sys.exit(2)
if argmod == 'argparse':
	parser = argparse.ArgumentParser(description='Utility for bulk testing of TCP connectivity.')
	parser.add_argument('--infile', nargs=1, help='(Mandatory)Input file containing hosts and ports to test.')
	parser.add_argument('--outfile', nargs=1, help='(Optional)File to output results to.')
	parser.add_argument('--timeout', nargs=1, help='(Optional)Connection timeout in seconds (default=5).')
	args = parser.parse_args()
	infiles = args.infile
	if not args.outfile is None:
		outfile = args.outfile[0]	
	if not args.timeout is None:
		timeout = args.timeout[0]	
else:
	parser = OptionParser()
	parser.add_option('--infile', dest='infiles', help='(Mandatory)Input file containing hosts and ports to test.', metavar='FILE')
	parser.add_option('--outfile', dest='outfile', help='(Optional)File to output results to.')
	parser.add_option('--timeout', dest='timeout', help='(Optional)Connection timeout in seconds (default=5).')
	(options, args) = parser.parse_args()
	if not options.infiles is None:
		infiles = [options.infiles]
	if not options.outfile is None:
		outfile = options.outfile
	if not options.timeout is None:
		timeout = options.timeout

if infiles is None:
	print "You must supply at least one --infile argument."
	sys.exit(2)


# Check for old versions of Python which lack socket.create_connection()
if sys.version_info[0] == 2 and sys.version_info[1] < 6:
	old_python = 1

output_to_file = 0
if not outfile is None:
	ofd = open(outfile, 'a')
	output_to_file = 1


# Process any files provided on the command line
for file in infiles:
	fd = open(file, 'r')
	for ln in fd:
		(host,ports) = ln.split(',', 1)
		portlist = ports.rstrip('\n').split(',')

		goodports = []
		badports = []
		for port in portlist:
			rc = 0
			try:
				if old_python == 0:
					sock = socket.create_connection((host,port),
						timeout)
				else:
					sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
					sock.connect((host,int(port)))
				sock.close()
				goodports.append(port)
			except:
				badports.append(port)
		oln = host + " Success: " + ','.join(goodports) + " Failed: " + ','.join(badports)
		if output_to_file == 1:
			ofd.write(oln + "\n")
		else:
			print oln
	fd.close()

if output_to_file == 1:
	ofd.close()
sys.exit(0)