#!/usr/bin/python

WhatThisIS="""

This is a hacked together script to connect to all of your running EC2 instances, and to provide a single SSH prompt amongst all of them, and to assist in uploading, splitting, and downloading files. email me at winniningham at and email server called gmail.com :P

"""

import paramiko
from numpy import *
from libcloud.types import Provider
from libcloud.providers import get_driver

EC2 = get_driver(Provider.EC2)
driver = EC2(YOURKEY, YOURSECRET)

class myConnection():
	def _connect(self):
		ssh = paramiko.SSHClient()
		ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
		ssh.connect(self._host, self._port, self._user, key_filename=self._keypath)
		self.ssh=ssh
	def _command(self, command):
		inp,out,err = self.ssh.exec_command(command)
		inp.close()
		return (out.readlines(), err.readlines())

import cmd

class bot(cmd.Cmd):
	def __init__(self):
		cmd.Cmd.__init__(self)
		self.bots=[]
		self.prompt="cluster> "
	def append(self,h,u,k,name='remote'):
		a=myConnection()
		a._host=h
		a._user=u
		a._port=22
		a._keypath=k
		a._connect()
		a.name=name
		self.bots.append(a)
		print(repr([h,u]))
	def mdo(self,command):
		print "----------------------------------------"
		for b in self.bots:
			out, err = b._command(command)
			if len(out) == 0 and len(err) == 0:
				print b.name + ':\t no output'
			for x in out:
				print b.name + ':\t' + x.strip('\n')
			for x in err:
				print b.name + ': ERR: \t' + x.strip('\n')
		print "----------------------------------------"
	def onecmd(self, line):
		if line.strip() == 'exit':
			exit()
		elif line.strip()[:4] == 'PUT ':
			options=line.strip().split('PUT ')[1].split(' ')
			self.put_all(options[0],options[1])
		elif line.strip()[:4] == 'GET ':
			options=line.strip().split('GET ')[1].split(' ')
		elif line.strip()[:8] == 'GETSAFE ':
			options=line.strip().split('GETSAFE ')[1].split(' ')
			self.get_all_safe(options[0],options[1],options[2])
		elif line.strip()[:6] == 'SPLIT ':
			options=line.strip().split('SPLIT ')[1].split(' ')
			self.split_file(options[0],options[1])
		elif line.strip()[:9] == 'POSITION ':
			options=line.strip().split('POSITION ')[1].split(' ')
			self.position_files(options[0],options[1])
		else:
			self.mdo(line)
	def put_all(self, filename,target_path):
		for b in self.bots:
			sftp=b.ssh.open_sftp()
			sftp.put(filename,target_path+filename)
			print "PUT on " + b.name + ': ' + target_path+filename
	def get_all(self, filename, target_path):
		for b in self.bots:
			sftp=b.ssh.open_sftp()
			target=b.name.replace(' ','_') + '_' + filename
			sftp.get(filename,target_path+target)
	def get_all_safe(self, src_path, filename, target_path):
		for b in self.bots:
			sftp=b.ssh.open_sftp()
			target=b.name.replace(' ','_') + '_' + filename
			sftp.get(src_path+filename,target_path+target)
	def position_files(self, pad, target_path):
		Count=0
		for b in self.bots:
			filename=str(Count).rjust(int(pad),'0')+'.txt'
			sftp=b.ssh.open_sftp()
			sftp.put(filename,target_path)
			print "POSITIONED on " + b.name + ': ' + target_path
			Count+=1
	def split_file(self,inputfile,target_path):
		f=open(inputfile,'r').readlines()
		pieces=array_split(f,len(self.bots))
		#print repr(pieces)
		Count=0
		for b in self.bots:
			sftp=b.ssh.open_sftp()
			piece=[x for x in pieces[Count]]
			g=sftp.open(target_path,'w')
			g.write(''.join(piece))
			g.close()
			print "SPLIT on " + b.name + ': ' + target_path + ' (' + str(len(piece))+ ')'
			Count +=1
	def run_ec2(self):
		Count=0
		for x in driver.list_nodes():
			if x.state==0:
				self.append(x.public_ip[0],'root','/home/thomas/Downloads/mini9.pem',x.name + ' ' + str(Count).rjust(2,'0'))
				Count += 1
		self.mdo('uptime')
	def interactive(self):
		self.cmdloop()


if __name__ == "__main__":
        b=bot()
	b.run_ec2()
	b.interactive()