#
# Copyright (C) 2010 Cloud.com, Inc.  All rights reserved.
# 
# This software is licensed under the GNU General Public License v3 or later.
# 
# It is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or any later version.
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.
#

#!/usr/bin/python
import sys, os
from subprocess import PIPE, Popen
import logging
import traceback
from os.path import exists, join
from signal import alarm, signal, SIGALRM, SIGKILL

class CloudRuntimeException(Exception):
    def __init__(self, errMsg):
        self.errMsg = errMsg
    def __str__(self):
        return self.errMsg
def formatExceptionInfo(maxTBlevel=5):
    cla, exc, trbk = sys.exc_info()
    excTb = traceback.format_tb(trbk, maxTBlevel)
    msg = str(exc) + "\n"
    for tb in excTb:
        msg += tb
    return msg

class bash:
    def __init__(self, args, timeout=600):
        self.args = args
        logging.debug("execute:%s"%args)
        self.timeout = timeout
        self.process = None
        self.success = False
        self.run()

    def run(self):
        class Alarm(Exception):
            pass
        def alarm_handler(signum, frame):
            raise Alarm

        try:
            self.process = Popen(self.args, shell=True, stdout=PIPE, stderr=PIPE)
            if self.timeout != -1:
                signal(SIGALRM, alarm_handler)
                alarm(self.timeout)

            try:
                self.stdout, self.stderr = self.process.communicate()
                if self.timeout != -1:
                    alarm(0)
            except Alarm:
                os.kill(self.process.pid, SIGKILL)
                raise  CloudRuntimeException("Timeout during command execution")

            self.success = self.process.returncode == 0
        except:
            raise  CloudRuntimeException(formatExceptionInfo())

#        if not self.success: 
#            raise  CloudRuntimeException(self.getStderr())

    def isSuccess(self):
        return self.success
    
    def getStdout(self):
        return self.stdout.strip("\n")
    
    def getLines(self):
        return self.stdout.split("\n")

    def getStderr(self):
        return self.stderr.strip("\n")


def initLoging(logFile=None):
    try:
        if logFile is None:
            logging.basicConfig(level=logging.DEBUG) 
        else: 
            logging.basicConfig(filename=logFile, level=logging.DEBUG) 
    except:
        logging.basicConfig(level=logging.DEBUG) 

def writeProgressBar(msg, result=None):
    if msg is not None:
        output = "%-80s"%msg
    elif result is True:
        output = "[%-2s]\n"%"OK"
    elif result is False:
        output = "[%-6s]\n"%"Failed"
    sys.stdout.write(output)
    sys.stdout.flush()
    
def printError(msg):
    sys.stderr.write(msg)
    sys.stderr.write("\n")
    sys.stderr.flush()

def checkRpm(pkgName):
    chkPkg = bash("rpm -q %s"%pkgName)
    writeProgressBar("Checking %s"%pkgName, None)
    if not chkPkg.isSuccess():
        writeProgressBar(None, False)
        printError("%s is not found, please make sure it is installed. You may try 'yum install %s'\n"%(pkgName, pkgName))
        return False
    else:
        writeProgressBar(None, True)
        return True
      
def checkEnv():
   writeProgressBar("Checking is root")
   ret = bash("whoami")
   if ret.getStdout() != "root":
       writeProgressBar(None, False)
       printError("This script must run as root")
       return False
   else:
       writeProgressBar(None, True)
       
   pkgList = ['tftp-server', 'syslinux', 'xinetd', 'chkconfig', 'dhcp']
   for pkg in pkgList:
       if not checkRpm(pkg):
           return False
   return True

def exitIfFail(ret):
    if not ret: sys.exit(1) 
    
def bashWithResult(cmd):
    writeProgressBar("Executing '%s'"%cmd)
    ret = bash(cmd)
    if not ret.isSuccess():
        writeProgressBar(None, False)
        writeProgressBar(ret.getStderr() + '\n')
        return False
    else:
        writeProgressBar(None, True)
        return True
    
def configurePxeStuff(): 
    stuff = ['tftp', 'xinetd', 'dhcpd']
    cmds = ['chkconfig --level 345 %s on' % i for i in stuff]
    cmds.append('service xinetd restart')
    
    for cmd in cmds:
        if not bashWithResult(cmd): return False
    return True  
    
def getTftpRootDir(tftpRootDirList):
    tftpRoot = bash("cat /etc/xinetd.d/tftp | grep server_args")
    if not tftpRoot.isSuccess():
        printError("Cannot get tftp root directory from /etc/xinetd.d/tftp, here may be something wrong with your tftp-server, try reinstall it\n")
        return False
    tftpRootDir = tftpRoot.getStdout()
    index = tftpRootDir.find("/")
    if index == -1:
        printError("Wrong server_arg in /etc/xinetd.d/tftp (%s)"%tftpRootDir)
        return False
    tftpRootDir = tftpRootDir[index:]
    tftpRootDirList.append(tftpRootDir)
    return True

def preparePING(tftpRootDir):
    pingFiles = ['boot.msg', 'initrd.gz', 'kernel', 'pxelinux.0']
    pingDir = "/usr/share/PING"
    
    for f in pingFiles:
        path = join(pingDir, f)
        if not exists(path):
            printError("Cannot find %s, please make sure PING-3.01 is installed"%path)
            return False
        if not bashWithResult("cp -f %s %s"%(path, tftpRootDir)): return False
     
    if not bashWithResult("mkdir -p %s/pxelinux.cfg"%tftpRootDir): return False
    
    return True
            
        
if __name__ == "__main__":
    initLoging("/tmp/cloud-setup-baremetal.log")
    tftpRootDirList = []
    
    exitIfFail(checkEnv())
    exitIfFail(configurePxeStuff())
    exitIfFail(getTftpRootDir(tftpRootDirList))
    
    tftpRootDir = tftpRootDirList[0].strip()
    exitIfFail(preparePING(tftpRootDir))
    print ""
    writeProgressBar("Setup BareMetal PXE server successfully")
    print ""
    writeProgressBar("TFTP root directory is: %s\n"%tftpRootDir)
    print ""
    sys.exit(0)