# 
# WBEM communication functions
# Feb 2001 Andrew Chabokha (achabokha@teleais.com)
# 

try:
    import win32com.client
except ImportError:
    print "Python Windows extention is not installed."


#
# COM Helper function
#
def CreateObject(ProgID):
    try:
        return win32com.client.Dispatch(ProgID)
    except:
        print "OLE error" ## or raise ...
        return -1


class wbemHost:
    def __init__(self, wbemHostName="", wbemRootPath="", wbemUserName="", wbemPassword=""):
        self.connected = self.GetConnectedToWbem(wbemHostName, wbemRootPath, wbemUserName, wbemPassword)
        if not self.connected:
            raise IOError
            #self.GetSystemInfo()
        
    #
    # return an wmiService object
    #
    def ConnectToWbem(self, wbemHostName="", wbemRootPath="", wbemUserName="", wbemPassword=""):
        try:
            wmiLocator = CreateObject ("WbemScripting.SWbemLocator")
            wmiService = wmiLocator.ConnectServer (wbemHostName, wbemRootPath, wbemUserName, wbemPassword)    
            # print "\nConnetcted to WBEM!"
            return wmiService
        except:
            print "\nConnection to WBEM is failed!"
            return none

    #
    # Setup connection veriables in __dict__
    #
    def GetConnectedToWbem(self, wbemHostName="", wbemRootPath="", wbemUserName="", wbemPassword=""):
        try:
            wmiLocator = CreateObject ("WbemScripting.SWbemLocator")
            wmiService = wmiLocator.ConnectServer (wbemHostName, wbemRootPath, wbemUserName, wbemPassword)    
            self.__dict__["wmiService"] = wmiService
            #print "\nConnetcted to WBEM!"
            return 1
        except:
            print "\nConnection to WBEM is failed!"
            return 0

    def GetSystemInfo(self):
        Systems = self.ExecQuery("select * from Win32_OperatingSystem where Primary=true")

        for System in Systems:
            self.ComputerName = System.CSName
            self.SysDirectory = System.SystemDirectory 
            self.WindowsDirectory = System.WindowsDirectory
    
    #
    # Get instances of given classes  
    #
    def GetInstancesOf(self, wbemEntity):
        try:
            wmiS = self.__dict__["wmiService"]
            Entity = wmiS.InstancesOf (wbemEntity)
            return Entity
        except:
            print "Cannot get instances of: %s", wbemEntity
            return -1

    def ExecQuery(self, wbemQuery):
        try:
            wmiS = self.__dict__["wmiService"]
            Entity = wmiS.ExecQuery (wbemQuery)
            return Entity
    
        except:
            print "Cannot execute the query: %s", wbemQuery
            return -1
    
    def Get(self, wbemEntity):
        try:
            wmiS = self.__dict__["wmiService"]
            Entity = wmiS.Get (wbemEntity)
            return Entity
        except:
            print "Cannot get: %s", wbemEntity
            return -1


    def RunProcess(self, CommandLine, CurrentDirectory=None, ProcessId=0):
        try:
            wmiProcess = self.Get("Win32_Process")
            method = wmiProcess.Methods_("Create")
            inParam = method.inParameters.SpawnInstance_()
            inParam.CommandLine = CommandLine
            inParam.CurrentDirectory = CurrentDirectory
            outParam = wmiProcess.ExecMethod_("Create", inParam)
            if outParam.returnValue != 0:
                print "Failed to run: %s" % CommandLine
            ProcessId = outParam.ProcessId
            
        except:
            print "Failed to run: %s" % CommandLine


    def ShareFolder(self, folderName, shareName):
        try:
            #print "Folder Name: " + folderName + ", Share Name: " + shareName
            wmiShare = self.Get("Win32_Share")
            method = wmiShare.Methods_("Create")
            inParam = method.inParameters.SpawnInstance_()
            inParam.Path = folderName
            inParam.Name = shareName
            inParam.Type = 0
            outParam = wmiShare.ExecMethod_("Create", inParam)
            if outParam.returnValue != 0:
                print "Failed to share: %s as %s" % (folderName, shareName)
            return outParam.returnValue
        except:
            print "Share creation is not succeed! An exception has been generated."
            print "Share fodler creation return value: %d" % outParam.returnValue
            return -1

    def Reboot(self):
        try:
            Systems = self.ExecQuery("select * from Win32_OperatingSystem where Primary=true")

            for System in Systems:
                outParam = System.ExecMethod_("Reboot")
                if outParam.returnValue != 0:
                    print "Failed to reboot system: %s" % (System.CSName)
                    
                return outParam.returnValue

        except:
            print "System reboot is not succeed! An exception has been generated."
            #print "Return value: %d" % outParam.returnValue
            return -1
#
# Helper functions
#
class wbemSystem(wbemHost):

    def listEnvironment(self):
        wmiEnvirons = self.GetInstancesOf("Win32_Environment")

        for env in wmiEnvirons:
            print "env -> " + env.Name

   
    def ShowInstancesOfProcess():
        wmiS = __dict__["wmiService"]
        Processes = wmiS.InstancesOf ("win32_Process")
    
        print "\nInstances of Processes:"

        for Process in Processes:
            print Process.Handle, Process.Name, Process.Priority, Process.OSName


    def ShowInstancesOfService(self):
        wmiS = __dict__["wmiService"]
        Services = wmiS.InstancesOf ("Win32_baseService")
    
        print "\nInstances of service:"
   
        for Service in Services:
            print Service.Name, Service.StartMode, Service.Started, Service.SystemName
        

    def ShowProcesses(self):
        wmiLocator = CreateObject ("WbemScripting.SWbemLocator")    
        wmiService = wmiLocator.ConnectServer ()

        Processes = wmiService.ExecQuery ("select Name from Win32_Process")

        print "\nLocal processes:"
    
        for Process in Processes:
            print Process.Handle, Process.Name 


    def ShowAssoc(self):
        wmiLocator = CreateObject ("WbemScripting.SWbemLocator")    
        wmiService = wmiLocator.ConnectServer ()

        Processes = wmiService.ExecQuery ('ASSOCIATORS OF {Win32_DiskDrive.DeviceID="\\\\.\\PHYSICALDRIVE0"} WHERE ResultClass=Win32_DiskPartition')

        print "\nPartition:"

        for Process in Processes:
            print Process.Name 
    

if __name__ == "__main__":

    test = 1
    
    if test == 1:
        myComputer = wbemHost()
        if not myComputer.connected:
            print "return"
        else:
            print "ok!"
            Processes = myComputer.GetInstancesOf("Win32_Process")
            print Processes
            print "\nInstances of Processes:"
            try:
                for Process in Processes:
                    print Process.Handle, Process.Name, Process.Priority, Process.OSName
            except:
                print "Exception!"
       
    if test == 2:
        RemoteComputer = "virgo.dev.tor.teleias.com"
        RemoteCIM = "root\\cimv2"
        RemoteAdminAccount = "administrator"
        RemoteAdminPassword = "moonlight"
        
        rhost = wbemHost(RemoteComputer, RemoteCIM, RemoteAdminAccount, RemoteAdminPassword) # Remotely

        rhost.RunProcess("calc.exe" )

        #rhost.Reboot()