import asyncio, inspect, os
from typing import final
from autorecon.config import config
from autorecon.io import e, info

class Target:

	def __init__(self, address, ipversion, type, autorecon):
		self.address = address
		self.ipversion = ipversion
		self.type = type
		self.autorecon = autorecon
		self.basedir = ''
		self.reportdir = ''
		self.scandir = ''
		self.lock = asyncio.Lock()
		self.ports = None
		self.pending_services = []
		self.services = []
		self.scans = []
		self.running_tasks = {}

	async def add_service(self, service):
		async with self.lock:
			self.pending_services.append(service)

	def extract_service(self, line, regex=None):
		return self.autorecon.extract_service(line, regex)

	async def extract_services(self, stream, regex=None):
		return await self.autorecon.extract_services(stream, regex)

	async def execute(self, cmd, blocking=True, outfile=None, errfile=None):
		target = self

		# Create variables for command references.
		address = target.address
		addressv6 = target.address
		scandir = target.scandir

		nmap_extra = target.autorecon.args.nmap
		if target.autorecon.args.nmap_append:
			nmap_extra += ' ' + target.autorecon.args.nmap_append

		if target.ipversion == 'IPv6':
			nmap_extra += ' -6'
			addressv6 = '[' + addressv6 + ']'

		plugin = inspect.currentframe().f_back.f_locals['self']

		cmd = e(cmd)

		tag = plugin.slug

		if config['verbose'] >= 1:
			info('Port scan {bblue}' + plugin.name + ' {green}(' + tag + '){rst} is running the following command against {byellow}' + address + '{rst}: ' + cmd)

		if outfile is not None:
			outfile = os.path.join(target.scandir, e(outfile))

		if errfile is not None:
			errfile = os.path.join(target.scandir, e(errfile))

		async with target.lock:
			with open(os.path.join(target.scandir, '_commands.log'), 'a') as file:
				file.writelines(cmd + '\n\n')

		process, stdout, stderr = await target.autorecon.execute(cmd, target, tag, patterns=plugin.patterns, outfile=outfile, errfile=errfile)

		target.running_tasks[tag]['processes'].append({'process': process, 'stderr': stderr, 'cmd': cmd})

		# If process should block, sleep until stdout and stderr have finished.
		if blocking:
			while (not (stdout.ended and stderr.ended)):
				await asyncio.sleep(0.1)
			await process.wait()

		return process, stdout, stderr

class Service:

	def __init__(self, protocol, port, name, secure=False):
		self.target = None
		self.protocol = protocol.lower()
		self.port = int(port)
		self.name = name
		self.secure = secure
		self.manual_commands = {}

	@final
	def tag(self):
		return self.protocol + '/' + str(self.port) + '/' + self.name

	@final
	def full_tag(self):
		return self.protocol + '/' + str(self.port) + '/' + self.name + '/' + ('secure' if self.secure else 'insecure')

	@final
	def add_manual_commands(self, description, commands):
		if not isinstance(commands, list):
			commands = [commands]
		if description not in self.manual_commands:
			self.manual_commands[description] = []

		# Merge in new unique commands, while preserving order.
		[self.manual_commands[description].append(m) for m in commands if m not in self.manual_commands[description]]

	@final
	def add_manual_command(self, description, command):
		self.add_manual_commands(description, command)

	@final
	async def execute(self, cmd, blocking=True, outfile=None, errfile=None):
		target = self.target

		# Create variables for command references.
		address = target.address
		addressv6 = target.address
		scandir = target.scandir
		protocol = self.protocol
		port = self.port
		name = self.name

		if config['create_port_dirs']:
			scandir = os.path.join(scandir, protocol + str(port))
			os.makedirs(scandir, exist_ok=True)
			os.makedirs(os.path.join(scandir, 'xml'), exist_ok=True)

		# Special cases for HTTP.
		http_scheme = 'https' if 'https' in self.name or self.secure is True else 'http'

		nmap_extra = target.autorecon.args.nmap
		if target.autorecon.args.nmap_append:
			nmap_extra += ' ' + target.autorecon.args.nmap_append

		if protocol == 'udp':
			nmap_extra += ' -sU'

		if target.ipversion == 'IPv6':
			nmap_extra += ' -6'
			addressv6 = '[' + addressv6 + ']'

		plugin = inspect.currentframe().f_back.f_locals['self']

		cmd = e(cmd)

		tag = self.tag() + '/' + plugin.slug

		if config['verbose'] >= 1:
			info('Service scan {bblue}' + plugin.name + ' {green}(' + tag + '){rst} is running the following command against {byellow}' + address + '{rst}: ' + cmd)

		if outfile is not None:
			outfile = os.path.join(scandir, e(outfile))

		if errfile is not None:
			errfile = os.path.join(scandir, e(errfile))

		async with target.lock:
			with open(os.path.join(target.scandir, '_commands.log'), 'a') as file:
				file.writelines(cmd + '\n\n')

		process, stdout, stderr = await target.autorecon.execute(cmd, target, tag, patterns=plugin.patterns, outfile=outfile, errfile=errfile)

		target.running_tasks[tag]['processes'].append({'process': process, 'stderr': stderr, 'cmd': cmd})

		# If process should block, sleep until stdout and stderr have finished.
		if blocking:
			while (not (stdout.ended and stderr.ended)):
				await asyncio.sleep(0.1)
			await process.wait()

		return process, stdout, stderr