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

from pentest.parser import JarvisParser
import uuid
import glob
from pentest.utils import Config
import subprocess
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):
		self.config = Config()
		if not os.path.isfile(self.config.histfile):
			open(self.config.histfile, 'w').close()
		self.fd = open(self.config.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('edit', help='Runs editor on current notes file')
		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.config.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.config.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.config.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.config.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.config.userdir, "img", "%s.png" % screenshot_id),
			os.path.join(self.current, self.config.usedir, "img", "%s.png" % screenshot_name)
		)

	def notes(self):
		subprocess.call([self.config.editor, self.current, self.config.userdir, self.config.notes])
