#!/usr/bin/python3
# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
# 
#   http://www.apache.org/licenses/LICENSE-2.0
# 
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

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 = "[ \033[92m%-2s\033[0m ]\n"%"OK"
    elif result is False:
        output = "[ \033[91m%-6s\033[0m ]\n"%"FAILED"
    sys.stdout.write(output)
    sys.stdout.flush()
    
def printError(msg):
    sys.stderr.write(msg)
    sys.stderr.write("\n")
    sys.stderr.flush()

def printMsg(msg):
    sys.stdout.write(msg+"\n")
    sys.stdout.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('/etc/init.d/xinetd restart')
    
    for cmd in cmds:
        if not bashWithResult(cmd): return False
        
    chkIptable = bash('chkconfig --list iptables')
    if 'on' in chkIptable.getStdout():
        printMsg("Detected iptables is running, need to open tftp port 69")
        if not bashWithResult('iptables -I INPUT 1 -p udp --dport 69 -j ACCEPT'): return False
        if not bashWithResult('/etc/init.d/iptables save'): 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))
    printMsg("")
    printMsg("Setup BareMetal PXE server successfully")
    printMsg("TFTP root directory is: %s\n"%tftpRootDir)
    sys.exit(0)