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

from pentest.parser import JarvisParser
import ConfigParser
import json
import os
from os.path import join as buildpath
from os.path import expanduser
import argparse


class JarvisManager(object):
	"""Top-level Jarvis class that manages penetration tests"""
	def __init__(self):
		config = ConfigParser.ConfigParser()
		config.read("/etc/jarvis.conf")
		self.userdir = config.get('jarvis', 'user_directory')
		self.outdir = config.get('jarvis', 'output_directory')
		self.histfile = config.get('jarvis', 'pentests_history')
		self.tree = json.loads(config.get('jarvis', 'pentest_tree'))
		self.notes_ = config.get('jarvis', 'notesfiles')
		if not os.path.isfile(self.histfile):
			open(self.histfile, 'w').close()
		self.fd = open(self.histfile, 'rb+')

		self.update()
		self.parser = argparse.ArgumentParser("Jarvis")
		subparsers = self.parser.add_subparsers(dest='command', help='Action to perform')
		start_parser = subparsers.add_parser('start', help='Edits a picture')
		start_parser.add_argument('directory', help='The directory name')
		subparsers.add_parser('stop', help='Stop the current pentest')
		subparsers.add_parser('init', help='Creates working tree for current pentest')
		subparsers.add_parser('screenshot', help='Takes a screenshot')
		subparsers.add_parser('notes', help='Starts notes edition')
		subparsers.add_parser('img', help='Lists screenhosts')

	def update(self):
		self.fd.seek(0, os.SEEK_END)
		if self.fd.tell() == 0:
			self.current = None
		else:
			while True:
				self.fd.seek(-1, os.SEEK_CUR)
				char = self.fd.read(1)
				self.fd.seek(-1, os.SEEK_CUR)
				if char == "\n" or self.fd.tell() == 0:
					if char == "\n":
						self.fd.read(1)
					self.current = self.fd.read()
					break
		self.fd.seek(0, os.SEEK_SET)

	def shortcutize(self):
		try:
			os.unlink("%s/pentest" % expanduser("~"))
		except Exception:
			pass
		os.symlink(self.current, "%s/pentest" % expanduser("~"))

	def start(self, path):
		self.fd.seek(0, os.SEEK_END)
		if os.path.isdir(path):
			if os.path.getsize(self.histfile) == 0:
				self.fd.write("%s" % os.path.abspath(path))
			else:
				self.fd.write("\n%s" % os.path.abspath(path))
		self.fd.seek(0, os.SEEK_SET)
		self.current = os.path.abspath(path)
		self.shortcutize()
		print '[i] A pentest located at %s has been stacked' % os.path.abspath(path)

	def stop(self):
		self.fd.seek(0, os.SEEK_END)
		if self.fd.tell() != 0:
			while True:
				self.fd.seek(-1, os.SEEK_CUR)
				char = self.fd.read(1)
				self.fd.seek(-1, os.SEEK_CUR)
				if char == '\n' or self.fd.tell() == 0:
					self.fd.truncate()
					break
		self.fd.seek(0, os.SEEK_SET)
		self.update()
		self.shortcutize()

	def init(self):
		if self.current:
			for directory in self.tree:
				if not os.path.isdir(os.path.join(self.current, directory)):
					os.mkdir(os.path.join(self.current, directory))
				else:
					print '[!] %s already exists, skipping' % os.path.join(self.current, directory)

	def screenshot(self):
		screenshot_id = str(uuid.uuid4())
		subprocess.check_output(["import", "%s/%s/img/%s.png" % (self.current, self.userdir, screenshot_id)])
		screenshot_name = subprocess.check_output(["zenity", "--entry", "--title=Jarvis-Screen", "--text=Nom de la capture d'écran"])
		screenshot_name = screenshot_name[:-1]
		while True:
			if os.path.isfile(os.path.join(self.current, self.userdir, "img", "%s.png" % screenshot_name)):
				screenshot_name = subprocess.check_output(["zenity", "--entry", "--title=Jarvis-Screen", "--text=Nom déjà utilisé"])
				screenshot_name = screenshot_name[:-1]
			else:
				break

		os.rename(
			os.path.join(self.current, self.userdir, "img", "%s.png" % screenshot_id),
			os.path.join(self.current, self.usedir, "img", "%s.png" % screenshot_name)
		)

	def notes(self):
		os.system('vim %s/%s/%s' % (self.current, self.userdir, self.notes_))
