#!/usr/bin/env python
# -*- coding: utf-8 -*-

import os
from os.path import join as buildpath
from pentest.manager import JarvisManager
from pentest.utils import JarvisLogger
from datetime import datetime
from slugify import slugify
import shlex
import sys
import subprocess


class JarvisCmd(object):
	"""Top-level Jarvis class that manages penetration tests"""
	def __init__(self, *args):
		self.args = list(*args)
		self.name = os.path.basename(self.args[0])
		self.manager = JarvisManager()
		if self.manager.current:
			self.logger = JarvisLogger(self.manager)
		self.args_config = None
		self.record = False
		self._output_directory = self.manager.config.outdir

		self.clean_env = os.environ.copy()
		dirs = self.clean_env["PATH"].split(':')
		dir_to_exclude = os.path.dirname(self.args[0])
		while dir_to_exclude in dirs:
			del dirs[dirs.index(dir_to_exclude)]
		self.clean_env["SHELL"] = "/bin/bash"
		self.clean_env["PATH"] = ':'.join(dirs)

	def __str__(self):
		return ' '.join(map(lambda x: shlex.quote(x), self.args))

	def __repr__(self):
		return ' '.join(map(lambda x: shlex.quote(x), self.args))

	@property
	def output_directory(self):
		if self.manager.current:
			return buildpath(self.manager.current, self.manager.config.userdir, self._output_directory)
		else:
			return False

	@output_directory.setter
	def output_directory(self, value):
		self._output_directory = value

	def set_args_config(self, args_config):
		self.args_config = args_config

	def set_recordable(self, recordable):
		self.record = recordable

	def run(self, marker=None):
		if not self.manager.current:
			self.args[0] = self.name
			self.execute()
		elif "--nojarvis" in self.args:
			self.args[0] = self.name
			del self.args[self.args.index("--nojarvis")]
			self.execute()
		else:
			for passthrough_arg in self.args_config["passthrough_args"]:
				if passthrough_arg in self.args:
					self.args[0] = self.name
					self.execute()
					break

			self.marker = marker
			self.prepare_command()
			self.logger.register(str(self))
			if self.record:
				JarvisRecordableCommand(self).run()
			else:
				self.execute()

	def prepare_command(self):
		self.args[0] = self.name

		# if no marker is provided, hook will not run
		if self.marker and not self.record:
			for output_arg in self.args_config["output_args"]:
				self.add_arguments(
					output_arg.get("pos", None) or "begin",
					output_arg["arg"],
					output_arg.get("filename", None) or	self.generate_file_name(
						output_arg['format'],
						self.marker
					)
				)

	def generate_file_name(self, format_, marker, record=False):
		filename = self.name + '-'
		filename += (marker.replace('/', '_').replace(':', '_') + '-') if marker else ''
		filename += self.timestamp()
		filename += ('.' + format_) if format_ else ''
		return buildpath(self.output_directory, filename)

	def timestamp(self):
		return slugify(str(datetime.now().replace(microsecond=0)))

	def execute(self):
		"""
		Execute the original program by patching the PATH
		environment variable
		"""
		subprocess.call(self.args, env=self.clean_env)
		sys.exit(0)

	def add_arguments(self, pos, *args):
		"""
		Method used to add safely add arguments into the commandline
		"""
		if pos == "begin":
			for arg in args[::-1]:
				self.args.insert(1, arg)
		elif pos == "end":
			for arg in args:
				self.args.insert(len(self.args), arg)
		else:
			print("[!] Invalid argument position : %s" % pos)
			sys.exit(0)


class JarvisRecordableCommand(object):

	def __init__(self, JarvisCommand):
		self.cmd = JarvisCommand
		self.recordable_cmd = ["script", "-e", "-q", "-c"]
		self.recordable_cmd.append(str(self.cmd))
		self.recordable_cmd.append(self.record_filename)

	@property
	def record_filename(self):
		filename = self.cmd.name + '-'
		filename += "record-"
		filename += (self.cmd.marker.replace('/', '_').replace(':', '_') + '-') if self.cmd.marker else ''
		filename += self.cmd.timestamp()
		filename += ".txt"
		return buildpath(self.cmd.output_directory, filename)

	def run(self):
		subprocess.call(self.recordable_cmd, env=self.cmd.clean_env)
		sys.exit(0)
