#!/usr/bin/env python
import socket
import sys
import argparse
import signal
import multiprocessing
import logging
logging.getLogger("scapy.runtime").setLevel(logging.ERROR)
from datetime import datetime
from multiprocessing import Pool
from scapy.all import *

#The function that actually does the stuff
def tcp_scan(arg):
	#This is for ^C when using pool.imap
	signal.signal(signal.SIGINT, signal.SIG_IGN)

	target_ip, port = arg
	sock 			= socket.socket(socket.AF_INET, socket.SOCK_STREAM)
	sock.settimeout(2)

	#Try to connect to the port, and return status
	try:
		sock.connect((target_ip, port))
		sock.close()
		return port, True

	except (socket.timeout, socket.error):
		return port, False

def syn_scan(arg):
	signal.signal(signal.SIGINT, signal.SIG_IGN)

	target_ip, port = arg
	src_port		= port
	syn_packet		= IP(dst=target_ip)/TCP(sport=src_port, dport=port, flags='S')
	rst_packet		= IP(dst=target_ip)/TCP(sport=src_port, dport=port, flags='R')

	resp			= sr1(syn_packet, timeout=2, verbose=False)
	if str(type(resp)) == "<type 'NoneType'>":
		return port, False
	elif resp and resp.getlayer(TCP).flags == 0x12:
		sr(rst_packet, timeout=1, verbose=False)
		return port, True
	return port, False

if __name__ == '__main__':
	#Command-line menu
	parser = argparse.ArgumentParser()
	parser.add_argument('hostname', help="Hostname / IP to scan.")
	parser.add_argument('-v', '--verbose', help="Verbose mode.", action='store_true')
	parser.add_argument('-s', '--syn', help="SYN scanning.", action='store_true')
	parser.add_argument('-t', '--tcp', help="TCP connect scanning, used by default.", action='store_true')
	parser.add_argument('-p', help="Number of processes.", default=multiprocessing.cpu_count() * 10, type=int)
	parser.add_argument('-r', help="Ports range, 1024 by default.", default=1024, type=int)
	args = parser.parse_args()
	
	try:
		#Check if the hostname is valid
		target_ip		= socket.gethostbyname(args.hostname)
		pool 			= Pool(processes=args.p)

	except socket.gaierror:
		sys.exit("Host name doesn't exist.")

	except Exception, e:
		sys.exit("Error: {0}".format(e))

	message = "Currently scanning " + target_ip
	print "-" * len(message)
	print message
	print "-" * len(message)

	t1 					= datetime.now()

	try:
		ports = range(1, args.r + 1)
		spaces = len(str(args.r))
		closed = 0
		#This is the multiprocessing, it calls scan with every port in the range
		for port, status in pool.imap(syn_scan if args.syn else tcp_scan, [(target_ip, port) for port in ports]):
			output_spaces =  " " * (spaces - len(str(port)))
			if args.verbose:
				print "|", port, output_spaces, "| OPEN   |" if status else "| CLOSED |"
			elif status:
				print "|", port, output_spaces, "| OPEN |"
			if not status:
				closed += 1

	except KeyboardInterrupt:
		pool.terminate()
		pool.join()
		sys.exit("\n\nTerminated by user. Bye !")

	except Exception, e:
		sys.exit("Error: {0}".format(e))

	t2 					= datetime.now()

	message = "Scanning complete in " + str(t2 - t1)
	print "-" * len(message)
	print message
	print closed, " ports closed"
	print args.r - closed
	print "-" * len(message), "\n"