#!/usr/bin/python

"""
    This module contains the deployment script to install a bootable OS of the first available disk
"""

import string, os, os.path, sys, glob, systemMgr
from logMgr import addLogRecord,CRITICAL,ERROR,WARNING,INFO,DEBUG
from config import DeploymentStore, DeployPath
from config import getstatusoutput
from deploymentData import ATSTATE,FAILURE,UPDATE,REFRESH,REBOOT,UNSUPPORTED
try:
    __file__
except:
    __file__ = __name__

def installSlax(device,restore):
    if restore:
        # creates an empty partition table
        os.system('fdisk %s << fend \no\n\w\nfend'%device)
        return ATSTATE
    
    SFDISK = 4
    SPARTITIONED = 5
    STATEVAR='DISKState'
    DISKstate = os.system("persistVariable read %s"%STATEVAR)
    if DISKstate < 0 or DISKstate == 0xff00:
        os.system("persistVariable update %s=%x"%(STATEVAR,SFDISK))
        DISKstate = SFDISK
    elif sys.platform == 'linux2':
        DISKstate = DISKstate >> 8
    
    if DISKstate == SFDISK:
        # creates 2 partitions, 2nd being a default Linux partition and is set active
        os.system('fdisk %s << fend \no\nn\np\n1\n\n+100M\nt\nc\nn\np\n2\n\n+5000M\na\n2\nw\nfend'%device)
        os.system("persistVariable update %s=%x"%(STATEVAR,SPARTITIONED))
        print "device %s was partitioned with fdisk"%device
        addLogRecord(0,__file__, INFO,"device %s was partitioned with fdisk"%device)
        return REBOOT
    try:
        if os.path.exists('/mnt/%s1'%os.path.basename(device)):
            os.removedirs('/mnt/%s1'%os.path.basename(device))
            os.system('umount /mnt/%s1'%os.path.basename(device))
        os.system('mkfs -t msdos %s1'%device)
        if os.path.exists('/mnt/%s2'%os.path.basename(device)):
            os.removedirs('/mnt/%s2'%os.path.basename(device))
            os.system('umount /mnt/%s2'%os.path.basename(device))
        os.system('mkfs %s2'%device)
    except OSError, e:
        print "Warning: the install of SLAX may fail. error: %s"%e
        addLogRecord(0,__file__, WARNING,"Warning: the install of SLAX may fail. error: %s"%e)
    cwd = os.getcwd()
    os.chdir('/boot')
    os.system('./make_disk.sh %s2'%device)
    print "SLAX is installed on device %s2"%device
    addLogRecord(0,__file__, INFO,"SLAX is installed on device %s2"%device)
    os.chdir(cwd)
    os.system("persistVariable delete %s"%(STATEVAR))
    return ATSTATE

def installWinPE(device,restore):
    #find disk no
    os.system("echo list disk > disk.txt")
    status, output = getstatusoutput("diskpart /s disk.txt")
    if status:
        print "'diskpart /s - find disk no' failed: status = 0x%x, %s"%(status,output)
        addLogRecord(0,__file__, ERROR,"'diskpart /s - find disk no' failed: status = 0x%x, %s"%(status,output))
        return FAILURE
    for line in output:
        if line.find("Disk") >= 0:
            parts = line.split()
            if parts[1].isdigit():
                No = parts[1]
                break
    # deletes partition table
    os.system("echo select Disk %s > disk.txt"%(No))
    os.system("echo select partition 1 >> disk.txt")
    os.system("echo delete partition >> disk.txt")
    os.system("echo select partition 1 >> disk.txt")
    os.system("echo delete partition >> disk.txt")
    os.system("echo select partition 1 >> disk.txt")
    os.system("echo delete partition >> disk.txt")
    os.system("echo select partition 1 >> disk.txt")
    os.system("echo delete partition >> disk.txt")
    status, output = getstatusoutput("diskpart /s disk.txt")
    if restore:
        os.remove("disk.txt")
        return ATSTATE
    
    #create partiton
    os.system("echo select Disk %s > disk.txt"%(No))
    os.system("echo clean >> disk.txt")
    os.system("echo create partition primary size=1024 >> disk.txt")
    os.system("echo active >> disk.txt")
    os.system("echo assign >> disk.txt")
    status, output = getstatusoutput("diskpart /s disk.txt")
    if status:
        print "'diskpart /s - create partition' failed: status = 0x%x, %s"%(status,output)
        addLogRecord(0,__file__, ERROR,"'diskpart /s - create partition' failed: status = 0x%x, %s"%(status,output))
        return FAILURE
    else:
        print "diskpart: create partition was successful for disk %s"%No
        addLogRecord(0,__file__, INFO,"diskpart: create partition was successful for disk %s"%No)
    # if DeploymentStore was on the disk just repartitioned - reboot!
    if not os.path.exists(DeploymentStore):
        os.system('reboot')
    # get volume drive letter
    os.system("echo list volume > disk.txt")
    status, output = getstatusoutput("diskpart /s disk.txt")
    os.remove("disk.txt")
    if status:
        print "'diskpart /s - list volume' failed: status = 0x%x, %s"%(status,output)
        addLogRecord(0,__file__, ERROR,"'diskpart /s - list volume' failed: status = 0x%x, %s"%(status,output))
        return FAILURE
    for line in output:
        if line.find("RAW") >= 0:
            parts = line.split()
            drive = parts[2]
            break
    else:
        print "RAW partition not found\n%s"%(output)
        addLogRecord(0,__file__, ERROR,"RAW partition not found\n%s"%(output))
        return FAILURE
    
    # format new partition
    os.system("format %s: /FS:FAT32 /V:WinPE /Y /Q"%drive)
    import shutil
    for drv in string.ascii_uppercase:
        if os.path.exists(drv+':'+os.sep+'winPDK.SDI'):
            for name in glob.glob(drv+':'+os.sep+'*'):
                if os.path.isdir(name):
                    if not name == drv+':'+os.sep+'i386':
                        continue
                    shutil.copytree(name,drive+':'+os.sep+os.path.basename(name))
                else:
                    shutil.copy2(name,drive+':'+os.sep)
            state = ATSTATE
            print "WinPE is installed on a drive letter %s"%drive
            addLogRecord(0,__file__, INFO,"WinPE is installed on a drive letter %s"%drive)
            break
    else:
        print "failed to find OS to install"
        addLogRecord(0,__file__, ERROR,"failed to find OS to install")
        state = FAILURE
    return state

def runState(restore=0):
    ret, device = systemMgr.queryProperty(0,'Disk','device','',{"instance":"1"})
    if not device or ret:
        print 'Disk device unavailable: ret = %d'%ret
        addLogRecord(0,__file__, ERROR,'Disk device unavailable: ret = %d'%ret)
        return FAILURE
    state = ATSTATE
    if sys.platform == 'linux2':
        state = installSlax(device,restore)
    elif sys.platform == 'win32':
        state = installWinPE(device,restore)
    return state

if __name__ == '__main__':
    """
    this script assumed these enviroment variables have been set:
    (see deploy.sh or deploy.cmd)
        DeployPath
        DeploymentStore
        PATH
        LD_LIBRARY_PATH (Linux)
        PYTHONPATH
    and this script has been run:
        prepareDeployment
    """
    
    if len(sys.argv) == 1:
        restore = 0
    elif len(sys.argv) == 2 and sys.argv[1] == '0':
        restore = 0
    else:
        restore = 1
    while 1:
        print "This script will repartition the first disk found in the system"
        print "Data will be lost. Do you want to continue? (Y)"
        next = sys.stdin.read(1)            # read a one-character string
        if next == 'Y' or next == 'y':
            break
        sys.exit()
    state = runState(restore)
    if state == REBOOT:
        print "System must be reset to effect changes to the partition table"
    sys.exit(state)
