hackedteam/test-av

import argparse
import sys
from time import sleep

from lib.config import Config
from lib.command import Command            

class WinUpdate:
    def __init__(self):
        """ Configure command interface through Config class
        """
        self.config_file = "c:/test-av/conf/vmware.conf"
        self.conf = Config(self.config_file)
        #
        # you can redefine the Command class providing vmrun full path:
        # self.cmd = Command("vmrun full path")
        #
        # TODO:
        #    - Command class with one argument
        self.cmd = Command(self.conf.path, self.conf.host, self.conf.user, self.conf.passwd)
        
        # full paths of script neede for update
        self.cscriptPath="c:/windows/system32/cscript.exe"
        self.netENScript="c:/Users/avtest/Desktop/EnableIF.bat"
        self.netDISScript="c:/Users/avtest/Desktop/DisableIF.bat"
        self.scriptPath="c:/Users/avtest/Desktop/WUA_SearchDownloadInstall.vbs"
        

    def doUpdate(self, vmx):
        if vmx == 'all':
            #update all vms
            vms = self.conf.getMachines()
            
            sys.stdout.write("[*] Starting all guests.\n")
            for vm in vms:
                self.cmd.startup(self.conf.getVmx(vm))
                sleep(5)
                
            sys.stdout.write("[*] Enabling Networking on guests.\n")
            for vm in vms:
                self.cmd.executeCmd(self.conf.getVmx(vm), self.netENScript)

            sys.stdout.write("[*] Updating operatings system on guests.\n")
            for vm in vms:
                self.cmd.executeCmd(self.conf.getVmx(vm), self.cscriptPath, self.netDISScript)
            sys.stdout.write("[*] Done.\n")
                    
        else:
            sys.stdout.write("[*] Starting %s.\n" % vmx)
            self.cmd.startup(self.conf.getVmx(vmx))
            sleep(60)
            sys.stdout.write("[*] Enabling Networking on %s.\n" % vmx)
            self.cmd.executeCmd(self.conf.getVmx(vmx), self.netENScript)
            sleep(10)
            sys.stdout.write("[*] Updating operatings system on %s.\n" % vmx)
            self.cmd.executeCmd(self.conf.getVmx(vmx), self.cscriptPath, self.scriptPath)
            sys.stdout.write("[*] Done.\n")
            
    
    def doReboot(self, vmx):
        if vmx == "all":
            vms = self.conf.getMachines()
            sys.stdout.write("[*] Disabling network on guests.\n")
            for vm in vms:
                self.cmd.executeCmd(self.conf.getVmx(vm), self.netDISScript)
                sleep(5)
            sys.stdout.write("[*] Rebooting guests.\n")
            for vm in vms:
                self.cmd.reboot(self.conf.getVmx(vm))
                sleep(5)    
        else:
            vm = self.conf.getVmx(vmx)
            sys.stdout.write("[*] Disabling network on %s.\n" % vmx)
            self.cmd.executeCmd(vm, self.netDISScript)
            sleep(20)
            sys.stdout.write("[*] Rebooting %s.\n" % vmx)
            self.cmd.reboot(vm)


    def refreshShot(self, vmx):
        if vmx == "all":
            sys.stdout.write("[*] Refresh snapshots on guests.\n")
            for vm in vms:
                self.cmd.refreshSnapshot(self.conf.getVmx(vm), 'current')
                sleep(30)
                self.cmd.suspend(self.conf.getVmx(vmx))
        else:
            sys.stdout.write("[*] Refresh snapshot of %s" % vmx)
            self.cmd.refreshSnapshot(self.conf.getVmx(vmx), 'current')
            self.cmd.suspend(self.conf.getVmx(vmx))
    
'''
MAIN
    
    parser = argparse.ArgumentParser(description='VMs updater script')
    parser.add_argument('op', metavar='operation', type=str,
                        help='operation to perform (start,reboot,refresh)')
    args = parser.parse_args()
    
    if args.op == "start":
        """ Stage 1:
        For each VM startup and launch update script
        """
        sys.stdout.write('[*] Startin Virtual machines\n\n')
        for vm in vms:

            sys.stdout.write("[*] Updating %s\n" % vm)    
            vmx = conf.getVmx(vm)
            cmd = Command(vmx, exe)
            cmd.revertSnapshot("current")
            sleep(10)
            cmd.startup()
            sleep(40)
            cmd.executeCmd(netENScript,None)
            sleep(10)
            cmd.update()

    elif args.op == "reboot":
        """ Stage 2:
        Reboot for updates
        """
        sys.stdout.write('[*] Performing reboot')
        for vm in vms:
            vmx = conf.getVmx(vm)
            cmd = Command(vmx, exe)
            #cmd.shutdown()
            #cmd.startup()
            cmd.reboot()
            sleep(5)
                
    elif args.op == "refresh":
        """ Stage 3:
        Refresh all VMs' snapshots and shut them down (remember reboot to apply all updates!!)
        """
        sys.stdout.write('[*] Refreshing vms snapshots')
        for vm in vms:
            vmx = conf.getVmx(vm)
            cmd = Command(vmx, exe)
            cmd.executeCmd(netDISScript,None)
            #cmd.refreshSnapshot("current")
            sleep(30)
            cmd.shutdown()
            sleep(1)
    
    else:
        print '[*] no action will be performed'
        

'''    
sys.stdout.write("Lets start!\n\n")
wu = WinUpdate()
#wu.doUpdate('norman')
#wu.doReboot('norman')
wu.refreshShot('norman')

sys.stdout.write("Everything is done!\n\n")