#!/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 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()
		self.logger = JarvisLogger(self.manager)
		self.args_config = None
		self.record = False
		self._output_directory = self.manager.outdir

		self.clean_env = os.environ.copy()
		dirs = self.clean_env['PATH'].split(':')
		dir_to_exclude = os.path.dirname(self.args[0])
		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: "'%s'" % x, self.args))

	def __repr__(self):
		return ' '.join(self.args)

	@property
	def output_directory(self):
		if self.manager.current:
			return buildpath(self.manager.current, self.manager.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:
			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(unicode(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)
