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

from pentest.parser import JarvisParser
import uuid
import pkgutil
from glob 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, 'r+')

		self.update()
		self.parser = argparse.ArgumentParser()
		subparsers = self.parser.add_subparsers(dest='command', help='Action to perform')
		start_parser = subparsers.add_parser('start', help='Start the pentest in the provided directory')
		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')
		edit_parser = subparsers.add_parser('edit', help='Runs image editor on a screenshot')
		edit_parser.add_argument('image', help='The screenshot to edit (filename)')
		subparsers.add_parser('notes', help='Starts notes edition')
		subparsers.add_parser('img', help='Lists screenhosts')
		subparsers.add_parser('hooks', help='Lists available hooks')

	def update(self):
		self.fd.seek(0, os.SEEK_END)
		if self.fd.tell() == 0:
			self.current = None
		else:
			while True:
				self.fd.seek(self.fd.tell() - 1, os.SEEK_SET)
				char = self.fd.read(1)
				self.fd.seek(self.fd.tell() - 1, os.SEEK_SET)
				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
		if self.current:
			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(self.fd.tell() - 1, os.SEEK_SET)
				char = self.fd.read(1)
				self.fd.seek(self.fd.tell() - 1, os.SEEK_SET)
				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", buildpath(self.current, self.config.userdir, self.config.imgdir, "%s.png" % screenshot_id)])
		screenshot_name = subprocess.check_output(["zenity", "--entry", "--title=Jarvis-Screen", "--text=Screenshot name"])
		screenshot_name = screenshot_name[:-1]
		while True:
			if os.path.isfile(os.path.join(self.current, self.config.userdir, self.config.imgdir, "%s.png" % screenshot_name)):
				screenshot_name = subprocess.check_output(["zenity", "--entry", "--title=Jarvis-Screen", "--text=Name already taken"])
				screenshot_name = screenshot_name[:-1]
			else:
				break

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

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

	def hooks(self):
		jarvis_path = os.path.dirname(__import__("pentest").__file__)
		hooks_path = buildpath(jarvis_path, "hooks")
		for _, name, _ in pkgutil.iter_modules([hooks_path]):
			print(name)

	def img(self):
		imgs_paths = glob("%s/*" % buildpath(self.current, self.config.userdir, self.config.imgdir))
		imgs_filenames = map(lambda x: os.path.basename(x), imgs_paths)
		for img_filename in imgs_filenames:
			print(img_filename)

	def edit(self, img):
		img_path = buildpath(self.current, self.config.userdir, self.config.imgdir, img)
		if os.path.isfile(img_path):
			subprocess.call(["pinta", img_path])
		else:
			print("[!] This image does not exists")
