# -*- coding: utf-8 -*-
#
# Enteletaor - https://github.com/cr0hn/enteletaor
#
# Redistribution and use in source and binary forms, with or without modification, are permitted provided that the
# following conditions are met:
#
# 1. Redistributions of source code must retain the above copyright notice, this list of conditions and the
# following disclaimer.
#
# 2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the
# following disclaimer in the documentation and/or other materials provided with the distribution.
#
# 3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote
# products derived from this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES,
# INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
# DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR
# SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY,
# WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#

import six
import zmq
import json
import redis
import socket
import logging
import eventlet
import ipaddress
import amqp.connection


from functools import partial
from collections import defaultdict
from threading import Thread, BoundedSemaphore

from .patch import patch_transport
from enteletaor_lib.libs.contrib.inetnum import get_inet_num

# Monkey patch for AMQP lib
patch_transport()
# Path thread library
eventlet.monkey_patch(socket=True, select=True, thread=True)


# Reconfigure AMQP LOGGER
logging.getLogger('amqp').setLevel(100)

log = logging.getLogger()

OPEN_SERVICES = defaultdict(dict)


# ----------------------------------------------------------------------
def _do_scan(config, sem, host):
	"""
	This function try to find brokers services open in remote servers
	"""

	handlers = {
		'Redis': handle_redis,
		'RabbitMQ': handle_amqp,
		'ZeroMQ': handle_zmq
	}

	log.warning("    > Analyzing host '%s' " % host)

	for port in config.ports.split(","):

		# Check each serve
		for server_type, handle in six.iteritems(handlers):

			log.info("      >> Trying to find %s service in '%s' port '%s'." % (server_type, host, port))

			try:

				# Try to check if port is open
				s = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
				s.settimeout(config.timeout)

				result = s.connect_ex((host, int(port)))

			except socket.gaierror as e:
				log.debug("%s : %s error: %s" % (server_type, port, e))
				continue
			finally:
				s.close()

			# Is port open?
			if result == 0:
				log.info("         <i> Port '%s' is open in '%s'" % (port, host))

				if handle(host, port, config) is True:
					log.error("      <!!> Open '%s' server found in port '%s' at '%s'" % (server_type, port, host))

					OPEN_SERVICES[host][server_type] = dict(
						state="open",
						port=port
					)
			else:
				log.debug("        <i> Port %s is closed" % port)

	sem.release()


# ----------------------------------------------------------------------
def action_scan_main(config):
	# --------------------------------------------------------------------------
	# Resolve target
	# --------------------------------------------------------------------------
	all_ips = build_targets(config)

	# --------------------------------------------------------------------------
	# Preparing scan
	# --------------------------------------------------------------------------
	target_number = len(all_ips)

	log.warning("  - Number of targets to analyze: %s" % target_number)

	# Semaphore
	sem = BoundedSemaphore(config.concurrency)
	threads = []

	# Map parameters
	_fn = partial(_do_scan, config, sem)

	log.error("  - Starting scan")

	# --------------------------------------------------------------------------
	# Do scan
	# --------------------------------------------------------------------------
	for x in all_ips:
		sem.acquire()

		t = Thread(target=_fn, args=(x,))
		threads.append(t)

		t.start()

	for t in threads:
		t.join()

	# --------------------------------------------------------------------------
	# Display results
	# --------------------------------------------------------------------------
	if OPEN_SERVICES:
		log.error("  - Open services found:")
		for host, content in six.iteritems(OPEN_SERVICES):
			log.error("    -> Host - %s" % host)
			for server_type, server_info in six.iteritems(content):
				log.error("       * %s/TCP [%s]" % (server_info['port'], server_type))

	else:
		log.error("  - No open services found")

	# --------------------------------------------------------------------------
	# Export results
	# --------------------------------------------------------------------------
	if config.output is not None:
		_output_path = "%s.json" % config.output if ".json" not in config.output else config.output

		with open(_output_path, "w") as f:
			json.dump(OPEN_SERVICES, f)

		log.error("  - Output results saved into: %s" % _output_path)


# --------------------------------------------------------------------------
def build_targets(config):

	results = set()

	# Split targets
	for t in config.target.split(","):
		try:
			results.update(str(x) for x in ipaddress.ip_network(t, strict=False))
		except ValueError:
			# --------------------------------------------------------------------------
			# If reach this, is not a IPs, is a domain
			# --------------------------------------------------------------------------

			# Try to get all assigned IP of domain
			if config.own_ips is True:

				# Extract domain
				try:
					val = get_inet_num(t.split(".")[-2])

					if val is not None:

						for v in val:
							log.debug("  -> Detected registered network '%s'. Added for scan." % v)

							results.update(str(x) for x in ipaddress.ip_network(six.u(v), strict=False))
				except KeyError:
					# Invalid domain
					log.debug("    <ii> Error while try to extract domain: '%s'" % t)

			# --------------------------------------------------------------------------
			# Get all IPs for domain
			# --------------------------------------------------------------------------

			# If target is a domain, remove CDIR
			_target_cdir = t.split("/")

			_cleaned_target = _target_cdir[0]

			try:
				# Resolve
				host_ip = socket.gethostbyname(_cleaned_target)
			except socket.gaierror:
				# Try with the hostname with "www." again
				try:
					host_ip = socket.gethostbyname("www.%s" % _cleaned_target)
				except socket.gaierror:
					log.error("    <ii> Unable to resolve '%s'" % _cleaned_target)

					continue

			# Add CDIR to result
			scan_target = "%s%s" % (host_ip, "/%s" % _target_cdir[1] if len(_target_cdir) > 1 else "")

			results.update(str(x) for x in ipaddress.ip_network(six.u(scan_target), strict=False))

	return results


# --------------------------------------------------------------------------
# These 3 functions determinate if server has listen one of these services:
# - Redis server
# - RabbitMQ server
# - ZeroMQ PUB/SUB pattern
#
# Each function try to connect or do some action and determinate if service
# is on or not.
# --------------------------------------------------------------------------
def handle_redis(host, port=6379, extra_config=None):

	# log.debug("      * Connection to Redis: %s : %s" % (host, port))

	try:
		redis.StrictRedis(host=host, port=port, socket_connect_timeout=1, socket_timeout=1).config_get()

		return True

	except Exception:
		return False


# ----------------------------------------------------------------------
def handle_amqp(host, port=5672, extra_config=None):

	host_and_port = "%s:%s" % (host, port)

	# log.debug("      * Connection to RabbitMQ: %s : %s" % (host, port))

	try:
		amqp.connection.Connection(host=host_and_port,
		                           connect_timeout=1,
		                           read_timeout=1,
		                           socket_timeout=1)
		return True

	except Exception:
		return False


# ----------------------------------------------------------------------
def handle_zmq(host, port=5555, extra_config=None):

	# log.debug("      * Connection to ZeroMQ: %s : %s" % (host, port))

	context = zmq.Context()

	# Configure
	socket = context.socket(zmq.SUB)
	socket.setsockopt(zmq.SUBSCRIBE, b"")  # All topics
	socket.setsockopt(zmq.LINGER, 0)  # All topics
	socket.RCVTIMEO = 1000  # timeout: 1 sec

	# Connect
	socket.connect("tcp://%s:%s" % (host, port))

	# Try to receive
	try:
		socket.recv()

		return True
	except Exception:
		return False
	finally:
		socket.close()