#!/usr/bin/env python
# -*- coding: iso-8859-1 -*-

# Dinko Korunic 'kreator', 2005.
# parallel_ping.py
# script for parallelized ping, not necessarily internal

"""This program is a paralellized ping implementation. It can do standard
ICMP or TCP-style ping to a specified port, using one or more IPs as
targets. It can ping hosts given from a command line or in a file.

Usage:
  program_name [-h|--help] [[-i|--icmp] | [-t|--tcp port]] [-f|--file filename] [-z|--stamp] target1 [target2] [...]

Example 1:
python parallel_ping.py 127.0.0.1 www.google.com
('127.0.0.1', 1)
('www.google.com', 1)

Example 2:
python parallel_ping.py -t 80 -z 127.0.0.1 www.google.com nowhere.wtfwtfx.com
('127.0.0.1', 1, 1141163901)
('www.google.com', 1, 1141163901)
('nowhere.wtfwtfx.com', 0, 1141163901)
"""

__copyright__ = """Copyright (C) 2005  Dinko Korunic, InfoMAR d.o.o.

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
"""

__version__ = '$Id: parallel_ping.py 202 2006-02-28 22:52:11Z kreator $'

import sys, getopt, os, re, socket, threading, time

class ping:
	"""Simple ping implementation with TCP and ICMP style ping routines"""

	def __init__(self, cmd = None, cmdout = None, timeout = 1, \
			timestamp = False):
		"""Initialise local vars, command to be used, etc."""

		# if no cmd/cmdout is specified, use fping and it standard output	
		if cmd is None:
			cmd = 'fping -t %s' % (timeout * 1000)
		if cmdout is None:
			cmdout = r'^.* is alive'
    
		# set local vars
		self.cmd = cmd
		self.cmdout = cmdout
		self.timeout = timeout
		self.timestamp = timestamp

	def icmp_ping(self, host, port):
		"""ICMP style ping implementation

		>>> a = ping()
		>>> a.icmp_ping('127.0.0.1', None)
		('127.0.0.1', 1)

		>>> a = ping()
		>>> a.icmp_ping('xx127.0.0.1', None)
		('xx127.0.0.1', 0)
		"""

		# compile the output parser, initialise the result to false
		restats = re.compile(self.cmdout)
		result = 0

		# fork and exec the cmd
		pipeout, pipe = os.popen4('%s %s' % (self.cmd, host))

		try:
			# grep the output for the positive ping answer - break
			# immediately if found
			for line in pipe:
				statsmatch = restats.match(line)
				if statsmatch:
					result = 1
					break
		finally:
			# close both ends of the pipe
			pipe.close()
			pipeout.close()

		try:
			# collect the running process
			pid, sts = os.wait()
		except OSError:
			pass

		# get current time
		ts = int(time.time())

		# return timestamps or not, depending on flags
		if self.timestamp:
			return (host, result, ts)
		else:
			return (host, result)

	def tcp_ping(self, host, port):
		"""TCP style ping implementation

		>>> a = ping()
		>>> a.tcp_ping('127.0.0.1', 80)
		('127.0.0.1', 1)

		>>> a = ping()
		>>> a.tcp_ping('xx127.0.0.1', 80)
		('xx127.0.0.1', 0)
		"""

		# set the result to false
		result = 0

		try:
			# allocate IPv4 socket
			sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
			# set the timeout
			sock.settimeout(self.timeout);
			# connect
			sock.connect((host, port))
			result = 1
		except socket.error:
			pass

		sock.close()

		# get current time
		ts = int(time.time())

		if self.timestamp:
			return (host, result, ts)
		else:
			return (host, result)

	def ping(self, host, ptype = 'icmp', port = None):
		"""Ping interface with ICMP as default"""

		# get the dispatcher
		dispatchername = '%s_ping' % ptype

		# call proper dispatcher at all times
		dispatcher = getattr(self, dispatchername, self.icmp_ping)
		return dispatcher(host, port)

class pingthread(threading.Thread):
	"""Threading ping class for parallelized scan"""

	def __init__(self, func, arg, ptype, port):
		"""Initialise the ping thread object"""

		self.func = func
		self.arg = arg
		self.ptype = ptype
		self.port = port
		self.result = ''
		threading.Thread.__init__(self)

	def run(self):
		"""Self-run thread function"""

		self.result = self.func(self.arg, self.ptype, self.port)

def usage():
	"""Define the usage for this program"""

	usage = """Usage:
  %s [-h|--help] [[-i|--icmp] | [-t|--tcp port]] [-f|--file filename] [-z|--stamp] target1 [target2] [...]

Summary:
  This program is a paralellized ping implementation. It can do
standard ICMP or TCP-style ping to a specified port, using one or more IPs
as targets. It can ping hosts given from a command line or in a file."""
	print usage % sys.argv[0]
	sys.exit(1)

def _test():
	"""Self testing..."""

	import doctest, parallel_ping
	return doctest.testmod(parallel_ping)

def main(argv):
	"""Main function, mainly to parse the arguments

	>>> main(['-t', 80, '127.0.0.1', 'xx127.0.0.1'])
	('127.0.0.1', 1)
	('xx127.0.0.1', 0)

	>>> main(['127.0.0.1', 'xx127.0.0.1'])
	('127.0.0.1', 1)
	('xx127.0.0.1', 0)
	"""

	try:
		opts, args = getopt.gnu_getopt(argv, 'hit:f:z', \
			['help', 'icmp', 'tcp=', 'file=', 'stamp'])
	except getopt.GetoptError:
		usage()

	# set the default values
	ptype = 'icmp'
	port = None
	needarg = 1
	tstamp = False
	filename = None
	pinglist = []
  
	# ugly argument parse
	for opt, optarg in opts:
		if opt in ('-h', '--help'):
			usage()
		elif opt in ['-i', '--icmp', '-t', '--tcp']:
			params = {'-i': 'icmp', '-t': 'tcp', \
				'--tcp': 'tcp', '--icmp': 'icmp'}
			ptype = params[opt]
			try:
				port = int(optarg)
			except ValueError:
				pass
		elif opt in ('-f', '--file'):
			filename = optarg
			needarg = 0
		elif opt in ('-z', '--stamp'):
			tstamp = True
		else:
			usage()

	# instantiate the generic ping object
	pingobj = ping(timeout = 2, timestamp = tstamp)

	# sanity check: need another arguments and there is none
	if (needarg == 1) and (len(args) < 1):
		usage()

	# do some pinging for hosts in arguments
	for arg in args:
		current = pingthread(pingobj.ping, arg, ptype, port)
		pinglist.append(current)
		current.start()
    
	# and for hosts in file argument, if there are some
	if filename:
		try:
			ftargets = file(filename, 'r')
		except IOError:
			usage()

		for line in ftargets:
			current = pingthread(pingobj.ping, line[:-1], ptype, port)
			pinglist.append(current)
			current.start()

		ftargets.close()

	# wait for all the threads to finish and display results
	for current in pinglist:
		current.join()
		print current.result

if __name__ == '__main__':
	# selftest
	#_test()
	# start the program...
	sys.exit(main(sys.argv[1:]))