#!/usr/bin/env python

import os
import sys
import glob
import subprocess
import json

class LtspIE:

	def __init__(self):
		
		self.ltsp_path="/opt/ltsp/"
		self.ltsp_tmp=self.ltsp_path+".tmp/"
		self.image_config="/etc/ltsp/images/"
		
	#def init
	
	def test_root(self):
		
		try:
			f=open("/run/ltsp-import-export","w")
			f.write(str(os.getpid()))
			f.close()
		except:
			return False
			
		return True
		
	#def test_root
	
	def usage(self):
		
		print("USAGE:\n\tltsp-import-export import|export TARGZ_FILE_PATH CHROOT_NAME [NEW_CHROOT_NAME]\n\t[!] Parameter order matters! [!]\n")
		sys.exit(1)
		
	#def
	
	def parse_args(self,args):
		
		if len(args)<4:
			self.usage()
			
		if args[1] not in ["import","export"]:
			self.usage()

		if not self.test_root():
			print("[!] You need root privileges to run this program [!]")
			sys.exit(1)

		file_path=args[2]
		chroot_name=args[3]
		use_xterm=False
		
		try:
			new_chroot_name=args[4]
		except:
			new_chroot_name=None
			
		if args[1]=="import":
			self.import_image(file_path,chroot_name,new_chroot_name)
			
		if args[1]=="export":
				
			self.export_image(file_path,chroot_name,False)
				
	#def parse_args

	def import_image(self,file_path,chroot_name,new_chroot_name=None):

		if not os.path.exists(file_path):
			self.usage()
		
		
		if new_chroot_name==None:
		
			cmd="echo;\
			echo [LtspIE] Extracting file...;\
			tar -xvzf " + file_path + " -C " + self.ltsp_path + " ;\
			mv "+self.ltsp_path+chroot_name+"/"+chroot_name+".json "+self.image_config+";\
			echo [LtspIE] Updating kernels...;\
			ltsp-update-kernels "+chroot_name+";\
			echo [LtspIE] Generating keys...;\
			ltsp-update-sshkeys;\
			echo [LtspIE]Updating image...;\
			ltsp-update-image "+chroot_name+";\
			echo [LtspIE] DONE;\
			echo;\
			"
			
			os.system(cmd)
			
		else:
			
			cmd="echo;\
			mkdir -p " + self.ltsp_tmp + ";\
			echo [LtspIE] Extracting file...;\
			tar -xvzf " + file_path + " -C " + self.ltsp_tmp + " ;\
			mv "+self.ltsp_tmp+chroot_name+" "+self.ltsp_path+new_chroot_name+ " ;\
			mv "+self.ltsp_path+new_chroot_name+"/"+chroot_name+".json "+self.image_config+new_chroot_name+".json ;\
			echo [LtspIE] Updating kernels...;\
			ltsp-update-kernels "+new_chroot_name+";\
			echo [LtspIE] Generating keys...;\
			ltsp-update-sshkeys;\
			echo [LtspIE]Updating image...;\
			ltsp-update-image "+new_chroot_name+";\
			echo [LtspIE] DONE;\
			echo;\
			"
			
			os.system(cmd)
		
			file_name=self.image_config+new_chroot_name+".json"
		
			f=open(file_name,"r")
			tmp=json.load(f)
			f.close()
			tmp["id"]=new_chroot_name
			tmp["name"]=new_chroot_name
			#tmp["status"]="working"
			
			f=open(file_name,"w")
			data=unicode(json.dumps(tmp,indent=4,encoding="utf-8",ensure_ascii=False)).encode("utf-8")
			f.write(data)
			f.close()
			
		
		return 0

	#def import_image
	
	def export_image(self,file_path,chroot_name,use_xterm=False):
		
		cmd="echo;\
		echo [LtspIE] Copying configuration file...;\
		cp " + self.image_config + chroot_name + ".json " + self.ltsp_path +chroot_name+";\
		echo [LtspIE] Creating tar.gz file...;\
		echo [LtspIE] tar -cvzf "+ file_path + " --one-file-system --exclude=/lost+found -C " + self.ltsp_path  + " " + chroot_name + " ;\
		tar -cvzf "+ file_path + " --one-file-system --exclude=/lost+found -C " + self.ltsp_path  + " " + chroot_name + " ;\
		echo [LtspIE] Image saved : " + file_path + " \
		echo [LtspIE] Finished;\
		echo;\
		"
		
		if use_xterm:
			cmd='xterm -T "Export LTSP Image" -e "' + cmd + 'read;"'
		
		try:
			output=subprocess.check_output(["chroot", self.ltsp_path+chroot_name, "apt-get", "clean"])
			for fl in glob.glob(self.ltsp_path + chroot_name + "/jdk-*.tar.gz"):
				try:
					os.remove(fl)
				except:
					pass
		except:
			pass

		os.system(cmd)
		
		return 0
		
	#def  export_image
	
	
	
if __name__=="__main__":
	
	ltsp=LtspIE()
	ltsp.parse_args(sys.argv)