'''

 Created on Nov 19, 2013

 @author: BarthaM

 '''

 import os

 import os.path

 from subprocess import Popen, PIPE, call

 import subprocess

 import sys

 import re

 import _winreg

 from cygwin import Cygwin

 import threading

 import time

 import string

 DEBUG = False

 class USBFilter:

     vendorid = ""

     productid = ""

     revision = ""

     def __init__(self, vendorid, productid, revision):

         self.vendorid = vendorid.lower()

         self.productid = productid.lower()

         self.revision = revision.lower()

         return

     def __eq__(self, other):

         return self.vendorid == other.vendorid and self.productid == other.productid and self.revision == other.revision

     def __hash__(self):

         return hash(self.vendorid) ^ hash(self.productid) ^ hash(self.revision)

     def __repr__(self):

         return "VendorId = \'" + str(self.vendorid) + "\' ProductId = \'" + str(self.productid) + "\' Revision = \'" + str(self.revision) + "\'"

 class VMManager(object):

     vmRootName = "SecurityDVM"

     systemProperties = None

     cygwin_path = 'c:\\cygwin64\\bin\\'

     vboxManage = 'VBoxManage'

     startNotifications = list()

     _instance = None

     #def __new__(cls, *args, **kwargs):

     #    if not cls._instance:

     #        cls._instance = super(VMManager, cls).__new__(cls, *args, **kwargs)

     #    return cls._instance

     _instance = None

     #def __new__(cls, *args, **kwargs):

     #    if not cls._instance:

     #        cls._instance = super(VMManager, cls).__new__(cls, *args, **kwargs)

     #    return cls._instance

     def __init__(self):

         self.cygwin_path = os.path.join(Cygwin.root(), 'bin') + os.path.sep

         self.vboxManage = os.path.join(self.getVBoxManagePath(), 'VBoxManage')

         self.systemProperties = self.getSystemProperties()

         return

     @staticmethod

     def getInstance():

         if VMManager._instance == None:

             VMManager._instance = VMManager()

         return VMManager._instance

     def putStartNotification(self, ip):

         self.startNotifications.append(ip)

     def isSDVMStarted(self, ip):

         return self.startNotifications.contains(ip)

     def execute(self, cmd):

         if DEBUG:

             print('trying to launch: ' + cmd)

         process = Popen(cmd, stdout=PIPE, stderr=PIPE) #shell = True

         if DEBUG:

             print('launched: ' + cmd)

         result = process.wait()

         res_stdout = process.stdout.read();

         res_stderr = process.stderr.read();

         if DEBUG:

             if res_stdout != "":

                 print res_stdout

             if res_stderr != "":

                 print res_stderr

         return result, res_stdout, res_stderr

     def getVBoxManagePath(self):

         """get the path to the VirtualBox installation on this system"""

         p = None

         try:

             k = _winreg.OpenKey(_winreg.HKEY_LOCAL_MACHINE, 'SOFTWARE\Oracle\VirtualBox')

             p = _winreg.QueryValueEx(k, 'InstallDir')[0]

             _winreg.CloseKey(k)

         except:

             pass

         return p

     # return hosty system properties

     def getSystemProperties(self):

         cmd = self.vboxManage + ' list systemproperties'

         result = self.execute(cmd)

         if result[1]=='':

             return None

         props = dict((k.strip(),v.strip().strip('"')) for k,v in (line.split(':', 1) for line in result[1].strip().splitlines()))

         return props

     # return the folder containing the guest VMs     

     def getDefaultMachineFolder(self):

         return self.systemProperties["Default machine folder"]

     #list the hostonly IFs exposed by the VBox host

     def getHostOnlyIFs(self):

         cmd = 'VBoxManage list hostonlyifs'

         result = self.execute(cmd)[1]

         if result=='':

             return None

         props = dict((k.strip(),v.strip().strip('"')) for k,v in (line.split(':', 1) for line in result.strip().splitlines()))

         return props

     def listRSDS(self):

         cmd = 'VBoxManage list usbhost'

         results = self.execute(cmd)[1]

         results = results.split('Host USB Devices:')[1].strip()

         items = list( "UUID:"+result for result in results.split('UUID:') if result != '')

         rsds = dict()   

         for item in items:

             props = dict()

             for line in item.splitlines():

                 if line != "":         

                     k,v = line[:line.index(':')].strip(), line[line.index(':')+1:].strip()

                     props[k] = v;

             if 'Product' in props.keys() and props['Product'] == 'Mass Storage':

                 usb_filter = USBFilter( re.search(r"\((?P<vid>[0-9A-Fa-f]+)\)", props['VendorId']).groupdict()['vid'], 

                                         re.search(r"\((?P<pid>[0-9A-Fa-f]+)\)", props['ProductId']).groupdict()['pid'],

                                         re.search(r"\((?P<rev>[0-9A-Fa-f]+)\)", props['Revision']).groupdict()['rev'] )

                 rsds[props['UUID']] = usb_filter;

                 if DEBUG:

                     print usb_filter

         return rsds

     # list all existing VMs registered with VBox

     def listVM(self):

         cmd = 'VBoxManage list vms'

         result = self.execute(cmd)[1]

         vms = list(k.strip().strip('"') for k,_ in (line.split(' ') for line in result.splitlines()))

         return vms

     # list existing SDVMs

     def listSDVM(self):

         vms = self.listVM()

         svdms = []

         for vm in vms:

             if vm.startswith(self.vmRootName) and vm != self.vmRootName:

                 svdms.append(vm)

         return svdms

     # generate valid (not already existing SDVM name). necessary for creating a new VM

     def generateSDVMName(self):

         vms = self.listVM()

         for i in range(0,999):

             if(not self.vmRootName+str(i) in vms):

                 return self.vmRootName+str(i)

         return ''

     # return the RSDs attached to all existing SDVMs

     def getAttachedRSDs(self):

         vms = self.listSDVM()

         attached_devices = dict()

         for vm in vms:

             rsd_filter = self.getUSBFilter(vm)

             if rsd_filter != None:

                 attached_devices[vm] = rsd_filter

         return attached_devices

     # configures hostonly networking and DHCP server. requires admin rights

     def configureHostNetworking(self):

         #cmd = 'vboxmanage list hostonlyifs'

         #self.execute(cmd)

         #cmd = 'vboxmanage hostonlyif remove \"VirtualBox Host-Only Ethernet Adapter\"'

         #self.execute(cmd)

         #cmd = 'vboxmanage hostonlyif create'

         #self.execute(cmd)

         cmd = 'VBoxManage hostonlyif ipconfig \"VirtualBox Host-Only Ethernet Adapter\" --ip 192.168.56.1 --netmask 255.255.255.0'

         self.execute(cmd)

         #cmd = 'vboxmanage dhcpserver add'

         #self.execute(cmd)

         cmd = 'VBoxManage dhcpserver modify --ifname \"VirtualBox Host-Only Ethernet Adapter\" --ip 192.168.56.100 --netmask 255.255.255.0 --lowerip 192.168.56.101 --upperip 192.168.56.200'

         self.execute(cmd)

     #create new virtual machine instance based on template vm named SecurityDVM (\SecurityDVM\SecurityDVM.vmdk)

     def createVM(self, vm_name):

         hostonly_if = self.getHostOnlyIFs()

         machineFolder = self.getDefaultMachineFolder()

         cmd = 'VBoxManage createvm --name ' + vm_name + ' --ostype Debian --register'

         self.execute(cmd)

         cmd = 'VBoxManage modifyvm ' + vm_name + ' --memory 512 --vram 10 --cpus 1 --usb on --usbehci on --nic1 hostonly --hostonlyadapter1 \"' + hostonly_if['Name'] + '\" --nic2 nat' 

         self.execute(cmd)

         cmd = 'VBoxManage storagectl ' + vm_name + ' --name contr1 --add sata --portcount 2'

         self.execute(cmd)

         cmd = 'VBoxManage storageattach ' + vm_name + ' --storagectl contr1 --port 0 --device 0 --type hdd --medium \"'+ machineFolder + '\SecurityDVM\SecurityDVM.vmdk\"'

         #--mtype immutable

         self.execute(cmd)

         return

     #remove VM from the system. should be used on VMs returned by listSDVMs    

     def removeVM(self, vm_name):

         print('removing ' + vm_name)

         cmd = 'VBoxManage unregistervm ' + vm_name + ' --delete'

         print self.execute(cmd)

         machineFolder = self.getDefaultMachineFolder()

         cmd = self.cygwin_path + 'bash.exe --login -c \"rm -rf ' + machineFolder + '\\' + vm_name + '*\"'

         print self.execute(cmd)

     # start VM

     def startVM(self, vm_name):

         print('starting ' +  vm_name)

         cmd = 'VBoxManage startvm ' + vm_name + ' --type headless' 

         result = self.execute(cmd)

         while result[0] != 0:

             print "Failed to start SDVM: ", vm_name, " retrying"

             time.sleep(1)

             result = self.execute(cmd)

         #verify against (0, 'Waiting for VM "SecurityDVM0" to power on...\r\nVM "SecurityDVM0" has been successfully started.\r\n', '')

         return result[0]

     # stop VM    

     def stopVM(self, vm_name):

         print('stopping ' + vm_name)

         cmd = 'VBoxManage controlvm ' + vm_name + ' poweroff'

         print self.execute(cmd)

     # return the hostOnly IP for a running guest    

     def getHostOnlyIP(self, vm_name):

         print('gettting hostOnly IP address ' + vm_name)

         cmd = 'VBoxManage guestproperty get ' + vm_name + ' /VirtualBox/GuestInfo/Net/0/V4/IP'

         result = self.execute(cmd)

         if result=='':

             return None

         result = result[1]

         if result.startswith('No value set!'):

             return None

         return result[result.index(':')+1:].strip()

     # attach removable storage device to VM by provision of filter

     def attachRSD(self, vm_name, rsd_filter):

         cmd = 'VBoxManage usbfilter add 0 --target ' + vm_name + ' --name OpenSecurityRSD --vendorid ' + rsd_filter.vendorid + ' --productid ' + rsd_filter.productid + ' --revision ' + rsd_filter.revision

         print self.execute(cmd)

     # return the description set for an existing VM

     def getVMInfo(self, vm_name):

         cmd = 'VBoxManage showvminfo ' + vm_name + ' --machinereadable'

         results = self.execute(cmd)[1]

         props = dict((k.strip(),v.strip().strip('"')) for k,v in (line.split('=', 1) for line in results.splitlines()))

         return props

     # return the configured USB filter for an existing VM 

     def getUSBFilter(self, vm_name):

         props = self.getVMInfo(vm_name)

         keys = set(['USBFilterVendorId1', 'USBFilterProductId1', 'USBFilterRevision1'])

         keyset = set(props.keys())

         usb_filter = None

         if keyset.issuperset(keys):

             usb_filter = USBFilter(props['USBFilterVendorId1'], props['USBFilterProductId1'], props['USBFilterRevision1'])

         return usb_filter

     #generates ISO containing authorized_keys for use with guest VM

     def genCertificateISO(self, vm_name):

         machineFolder = self.getDefaultMachineFolder()

         # create .ssh folder in vm_name

         cmd = self.cygwin_path+'bash.exe --login -c \"mkdir -p \\\"' + machineFolder + '\\' + vm_name + '\\.ssh\\\"\"'

         result = self.execute(cmd)

         # generate dvm_key pair in vm_name / .ssh     

         cmd = self.cygwin_path+'bash.exe --login -c \"ssh-keygen -q -t rsa -N \\"\\" -C \\\"' + vm_name + '\\\" -f \\\"' + machineFolder + '\\' + vm_name + '\\.ssh\\dvm_key\\\"\"'   #'echo -e "y\\n" |',

         result = self.execute(cmd)

         # set permissions for keys

         #TODO: test without chmod

         cmd = self.cygwin_path+'bash.exe --login -c \"chmod 500 \\\"' + machineFolder + '\\' + vm_name + '\\.ssh\\*\\\"\"'

         result = self.execute(cmd)

         # move out private key

         cmd = self.cygwin_path+'bash.exe --login -c \"mv \\\"' + machineFolder + '\\' + vm_name + '\\.ssh\\dvm_key\\\" \\\"' + machineFolder + '\\' + vm_name + '\\\"'

         result = self.execute(cmd)

         # rename public key to authorized_keys

         cmd = self.cygwin_path+'bash.exe --login -c \"mv \\\"' + machineFolder + '\\' + vm_name + '\\.ssh\\dvm_key.pub\\\" \\\"' + machineFolder + '\\' + vm_name + '\\.ssh\\authorized_keys\\\"'

         result = self.execute(cmd)

         # generate iso image with .ssh/authorized keys

         cmd = self.cygwin_path+'bash.exe --login -c \"/usr/bin/genisoimage -J -R -o \\\"' + machineFolder + '\\' + vm_name + '\\'+ vm_name + '.iso\\\" \\\"' + machineFolder + '\\' + vm_name + '\\.ssh\\\"\"'

         result = self.execute(cmd)

     # attaches generated ssh public cert to guest vm

     def attachCertificateISO(self, vm_name):

         machineFolder = self.getDefaultMachineFolder()

         cmd = 'vboxmanage storageattach ' + vm_name + ' --storagectl contr1 --port 1 --device 0 --type dvddrive --mtype readonly --medium \"' + machineFolder + '\\' + vm_name + '\\'+ vm_name + '.iso\"'

         result = self.execute(cmd)

         return result

     handleDeviceChangeLock = threading.Lock()

     # handles device change events

     def handleDeviceChange(self):

         if VMManager.handleDeviceChangeLock.acquire(True):

             #destroy unused vms

             new_ip = None

             attached_devices = self.getAttachedRSDs()

             connected_devices = self.listRSDS()

             for vm_name in attached_devices.keys():

                 if attached_devices[vm_name] not in connected_devices.values():

                     self.unmapNetworkDrive('h:')

                     self.stopVM(vm_name)

                     self.removeVM(vm_name)

             #create new vm for attached device if any

             attached_devices = self.getAttachedRSDs()

             for connected_device in connected_devices.values():

                 if (attached_devices and False) or (connected_device not in attached_devices.values()):

                     new_sdvm = self.generateSDVMName()

                     self.createVM(new_sdvm)

                     self.attachRSD(new_sdvm, connected_device)

                     self.startVM(new_sdvm)

                     # wait for machine to come up

                     while new_ip == None:

                         time.sleep(1)

                         new_ip = self.getHostOnlyIP(new_sdvm)

                     while new_ip not in self.startNotifications:

                         time.sleep(1)

                     if new_ip != None:

                         self.mapNetworkDrive('h:', '\\\\' + new_ip + '\\USB', None, None)

                     #TODO: cleanup notifications somwhere else (eg. machine shutdown)

                     self.startNotifications.remove(new_ip)

             VMManager.handleDeviceChangeLock.release()

             return new_ip

     def handleBrowsingRequest(self):

         if VMManager.handleDeviceChangeLock.acquire(True):

             new_ip = None

             new_sdvm = self.generateSDVMName()

             self.createVM(new_sdvm)

             self.genCertificateISO(new_sdvm)

             self.attachCertificateISO(new_sdvm)

             self.startVM(new_sdvm)

             # wait for machine to come up

             while new_ip == None:

                 time.sleep(1)

                 new_ip = self.getHostOnlyIP(new_sdvm)

             while new_ip not in self.startNotifications:

                 time.sleep(1)

             VMManager.handleDeviceChangeLock.release()

         return new_sdvm

     # executes command over ssh on guest vm

     def sshGuestExecute(self, vm_name, prog, user_name='opensec'):

         # get vm ip

         address = self.getHostOnlyIP(vm_name)

         machineFolder = self.getDefaultMachineFolder()

         # run command

         cmd = self.cygwin_path+'bash.exe --login -c \"ssh -i \\\"' + machineFolder + '\\' + vm_name + '\\dvm_key\\\"  ' + user_name + '@' + address + ' ' + prog + '\"'

         return self.execute(cmd)

     # executes command over ssh on guest vm with X forwarding

     def sshGuestX11Execute(self, vm_name, prog, user_name='osecuser'):

         #TODO: verify if X server is running on user account 

         #TODO: set DISPLAY accordingly

         address = self.getHostOnlyIP(vm_name)

         machineFolder = self.getDefaultMachineFolder()

         # run command

         cmd = self.cygwin_path+'bash.exe --login -c \"DISPLAY=:0 ssh -Y -i \\\"' + machineFolder + '\\' + vm_name + '\\dvm_key\\\"  '  + user_name + '@' + address + ' ' + prog + '\"'

         return self.execute(cmd)    

     #Small function to check the availability of network resource.

     def isAvailable(self, path):

         cmd = 'IF EXIST ' + path + ' echo YES'

         result = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True).communicate()

         return string.find(str(result), 'YES',)

     #Small function to check if the mention location is a directory

     def isDirectory(self, path):

         cmd = 'dir ' + path + ' | FIND ".."'

         result = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True).communicate()

         return string.find(str(result), 'DIR',)

     def mapNetworkDrive(self, drive, networkPath, user, password):

         self.unmapNetworkDrive('h:')

         #Check for drive availability

         if self.isAvailable(drive) > -1:

             print "Drive letter is already in use: ", drive

             return -1

         #Check for network resource availability

         while self.isAvailable(networkPath) == -1:

             time.sleep(1)

             print "Path not accessible: ", networkPath, " retrying"

             #return -1

         #Prepare 'NET USE' commands

         cmd = 'NET USE ' + drive + ' ' + networkPath

         if user != None:

             cmd = cmd + ' ' + password + ' /User' + user

         print "cmd = ", cmd

         #Execute 'NET USE' command with authentication

         result = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True).communicate()

         print "Executed: ", cmd

         if string.find(str(result), 'successfully',) == -1:

             print cmd, " FAILED"

             return -1

         #Mapped with first try

         return 1

     def unmapNetworkDrive(self, drive):

         #Check if the drive is in use

         if self.isAvailable(drive) == -1:

             #Drive is not in use

             return -1

         #Prepare 'NET USE' command

         cmd = 'net use ' + drive + ' /DELETE'

         result = subprocess.Popen(cmd, stdout=subprocess.PIPE, shell=True).communicate()

         if string.find(str(result), 'successfully',) == -1:

             return -1

         return 1

 #if __name__ == '__main__':

     #man = VMManager.getInstance()

     #man.removeVM('SecurityDVM0')

     #man.netUse('192.168.56.134', 'USB\\')

     #ip = '192.168.56.139'

     #man.mapNetworkDrive('h:', '\\\\' + ip + '\USB', None, None)

     #man.cygwin_path = 'c:\\cygwin64\\bin\\'

     #man.handleDeviceChange()

     #print man.listSDVM()

     #man.configureHostNetworking()

     #new_vm = man.generateSDVMName()

     #man.createVM(new_vm)

     #man.genCertificateISO(new_vm)

     #man.attachCertificateISO(new_vm)

     #man.attachCertificateISO(vm_name)

     #man.sshGuestExecute(vm_name, "ls")

     #man.sshGuestX11Execute(vm_name, "iceweasel")

     #cmd = "c:\\cygwin64\\bin\\bash.exe --login -c \"/bin/ls\""

     #man.execute(cmd)