#!/usr/bin/env python3

## Simple TCP port scanner

import argparse
import re
import ast
import socket
import io
import concurrent.futures


def _ports(port):
	"""
	Validates argparse port argument.
	Accepts a single port, a port range or a list (comma separated).
	"""

	valid_range = range(1, 65535 + 1)
	range_re = re.compile("\d+:\d+")
	list_re = re.compile("(\d+,\d+)+")

	try:
		port = int(port)
		if port not in valid_range:
			raise argparse.ArgumentTypeError("Port must be 1-65535")
		return [port]
	except ValueError:
		if range_re.match(port):
			start, stop = (int(i) for i in port.split(":"))
			if (start not in valid_range
				or stop not in valid_range
				or start > stop):
				raise argparse.ArgumentTypeError("Invalid range: " + port)
			return range(start, stop + 1)
		elif list_re.match(port):
			try:
				return [i for i in ast.literal_eval(id) if i in valid_range]
			except:
				raise argparse.ArgumentTypeError("Invalid list: " + port)


def _hostname(name):
	"""
	Validates argparse hostname argument.
	Accepts an ip address or a domain name.
	"""

	try:
		socket.gethostbyname(name)
		return name
	except socket.gaierror:
		raise argparse.ArgumentTypeError("Invalid hostname: " + name)
	try:
		socket.inet_aton(name)
		return name
	except socket.error:
		raise argparse.ArgumentTypeError("Invalid ip address: " + name)


def connect(hostname, port, verbose, banner):
	"""
	Connects to a given hostname and port.
	Set banner to True to get the application banner.
	Set verbose level (1,2,3) to get more informations.
	"""

	out = io.StringIO()
	try:
		connection = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
		connection.connect((hostname, port))
		print("[+] %d open" % port, file=out)
		if banner:
			connection.send("hi\r\n".encode("utf-8"))
			banner = connection.recv(100).decode("utf-8").strip("\n")
			print(" ┖──[b] " + banner, file=out)
		connection.close()
	except ConnectionRefusedError as e:
		if verbose > 0:
			print("[-] %d closed" % port, file=out)
		if verbose > 1:
			print(" ┖──[v] %s" % e, file=out)
	except socket.timeout:
		if banner and verbose > 1:
			print(" ┖──[b] No response", file=out)
	return out


def scan(hostname, ports, verbose=0, banner=False):
	"""Perform a scan on the given hostname and ports."""
	socket.setdefaulttimeout(2)
	try:
		name = socket.gethostbyaddr(hostname)
		print("[*] Scanning " + name[0])
	except socket.error as e:
		print("[*] Scanning " + hostname)
	with concurrent.futures.ThreadPoolExecutor(100) as executor:
		threads = []
		for port in ports:
			threads.append(executor.submit(connect, hostname, port, verbose, banner))
		for i in threads:
			print(i.result().getvalue(), end="")


def main():
	parser = argparse.ArgumentParser(description="Simple TCP port scanner.")
	parser.add_argument(
		"hostname", type=_hostname,
		help="Hostname, ip address.")
	parser.add_argument(
		"ports", type=_ports,
		help="Single port, range a:b or list (comma separated)")
	parser.add_argument(
		"-v", "--verbose",
		action="count", default=0,
		help="Show more information")
	parser.add_argument(
		"-b", "--banner",
		action="store_true", default=False,
		help="Get application banner.")
	args = parser.parse_args()

	scan(**vars(args))

if __name__ == "__main__":
	main()