#!/usr/bin/python
#
# $Id$
# $URL$

import sys
import os.path
import pdb
import subprocess
from optparse import OptionParser
import socket

#
# Thierry Parmentelat - INRIA
#
# class for issuing commands on a box, either local or remote
#
# the notion of 'buildname' is for providing each test run with a dir of its own
# buildname is generally the name of the build being tested, and can be considered unique
#
# thus 'run_in_buildname' mostly :
# (*) either runs locally in . - as on a local node we are already in a dedicated directory
# (*) or makes sure that there's a remote dir called 'buildname' and runs in it
#
# also, the copy operations
# (*) either do nothing if ran locally
# (*) or copy a local file into the remote 'buildname' 
# 

########################################
class Logger:

    def __init__(self):
        pass

    def log (self,msg):
        print >>sys.stderr, msg

logger=Logger()

########################################
# bool_to_exit_code
def b2e (b):
    if b: return 0
    else: return 1
# exit_to_bool_code
def e2b (b):
    return b==0

verbose=False
verbose=True
########################################
class Shell:
    def __init__(self,dry_run=False):
        self.dry_run=dry_run

    def header (self,message):
        logger.log("==============="+message)

    # return an exit code
    def run (self, argv, trash_stdout=False, trash_stderr=False, message="", force_dry_run=False):
        if self.dry_run:
            print 'DRY_RUN',message+'RUNNING '+" ".join(argv)
            return 0
        else:
            if message: self.header(message),
            if verbose: logger.log ('RUNNING '+" ".join(argv))
            if trash_stdout and trash_stderr:
                return subprocess.call(argv,stdout=file('/dev/null','w'),stderr=subprocess.STDOUT)
            elif trash_stdout and not trash_stderr:
                return subprocess.call(argv,stdout=file('/dev/null','w'))
            elif not trash_stdout and trash_stderr:
                return subprocess.call(argv,stderr=file('/dev/null','w'))
            elif not trash_stdout and not trash_stderr:
                return subprocess.call(argv)
            
    # like shell's $(...) which used to be `...`
    def backquote (self, argv, trash_stderr=False,message=""):
        if self.dry_run or verbose:
            if verbose: logger.log("WARNING: " + message + " backquote is actually running "+" ".join(argv))
        if not trash_stderr:
            return subprocess.Popen(argv,stdout=subprocess.PIPE).communicate()[0]
        else:
            return subprocess.Popen(argv,stdout=subprocess.PIPE,stderr=file('/dev/null','w')).communicate()[0]
            
####################
class Ssh:
    
    #################### the class
    def __init__(self,hostname=None,key=None, username='root',dry_run=False):
        self.hostname=hostname
        if not hostname: self.hostname='localhost'
        self.key=key
        self.username=username
        self.dry_run=dry_run

    # inserts a backslash before each occurence of the following chars
    # \ " ' < > & | ; ( ) $ * ~ 
    @staticmethod
    def backslash_shell_specials (command):
        result=''
        for char in command:
            if char in "\\\"'<>&|;()$*~":
                result +='\\'+char
            else:
                result +=char
        return result

    # check main IP address against the provided hostname
    @staticmethod
    def is_local_hostname (hostname):
        if hostname == "localhost":
            return True
        try:
            local_ip = socket.gethostbyname(socket.gethostname())
            remote_ip = socket.gethostbyname(hostname)
            return local_ip==remote_ip
        except:
            logger.log("WARNING : something wrong in is_local_hostname with hostname=%s"%hostname)
            return False

    @staticmethod
    def canonical_name (hostname):
        try:
            ip=socket.gethostbyname(hostname)
            return socket.gethostbyaddr(ip)[0]
        except:
            logger.log( 'Warning - could not find canonical name for host %s'%hostname)
            return hostname


    ########## building argv
    def is_local(self):
        return Ssh.is_local_hostname(self.hostname)
     
    def key_argv (self):
        if not self.key:
            return []
        return ["-i",self.key]

    def hostname_arg (self):
        if not self.username:
            return self.hostname
        else:
            return "%s@%s"%(self.username,self.hostname)
    
    std_options="-o BatchMode=yes -o StrictHostKeyChecking=no -o CheckHostIP=no -o ConnectTimeout=5"
    
    # command gets run on the right box
    def argv (self, argv, keep_stdin=False):
        if self.is_local():
            return argv
        ssh_argv = ["ssh"]
        if not keep_stdin: ssh_argv.append('-n')
        ssh_argv += Ssh.std_options.split()
        ssh_argv += self.key_argv()
        ssh_argv.append (self.hostname_arg())
        # using subprocess removes the need for special chars at this point
        #ssh_argv += [Ssh.backslash_shell_specials(x) for x in argv]
        ssh_argv += argv
        return ssh_argv

    # argv is a list
    def run(self, argv,*args, **kwds):
        return Shell(self.dry_run).run(self.argv(argv),*args, **kwds)
    # command is a string
    def run_string(self, command, *args, **kwds):
        return Shell(self.dry_run).run(self.argv(command.split()),*args, **kwds)

    # argv is a list
    def backquote(self, argv, *args, **kwds):
        return Shell(self.dry_run).backquote(self.argv(argv),*args, **kwds)
    # command is a string
    def backquote_string(self, command, *args, **kwds):
        return Shell(self.dry_run).backquote(self.argv(command.split()),*args, **kwds)

############################################################
class Vserver:

    def __init__ (self, name, dry_run=False):
        try:
            (self.name,self.hostname)=name.split("@")
        except:
            (self.name,self.hostname)=(name,'localhost')
        self.dry_run=dry_run
        self.ssh=Ssh(self.hostname,username='root',dry_run=self.dry_run)
        self._is_running=None

    # cache run mode
    def is_running (self):
        if self._is_running is not None: return self._is_running
        self._is_running = e2b(self.ssh.run_string('vserver %s status'%self.name,
                                                   trash_stdout=True,trash_stderr=True,
                                                   message='retrieving status for %s'%self.name))
    def running_char (self):
        if self._is_running is None:    return '?'
        elif not self._is_running:      return '-'
        else:                           return '+'

    def stop (self):
        # uncache
        self._is_running=None
        return e2b(self.ssh.run_string("vserver %s stop"%self.name,message='Stopping %s'%self.name))

    def start (self):
        # uncache
        self._is_running=None
        return e2b(self.ssh.run_string("vserver %s start"%self.name,message='Starting %s'%self.name))

    def printable_simple(self): return self.name + '@' + self.hostname
    def printable_dbg(self): return "%s@%s (%s)"%(self.name,self.hostname,self.running_char())
    def printable(self): 
        # force fetch
        self.is_running()
        return "%s@%s (%s)"%(self.name,self.hostname,self.running_char())

    def is_local (self):
        return self.ssh.is_local()

    @staticmethod
    def probe_host (hostname, dry_run=False):
        canonical=Ssh.canonical_name(hostname)
        ssh=Ssh(hostname,dry_run=dry_run)
        ls=ssh.backquote_string("ls -d1 /vservers/*",trash_stderr=True)
        return [ '%s@%s'%(path.replace('/vservers/',''),canonical) for path in ls.split()]

    # renaming
    def rename (self, newname):
        target=Vserver(newname)
        if self.hostname != target.hostname:
            # make user's like easier, no need to mention hostname again
            if target.is_local():
                target.hostname=self.hostname
            else:
                logger.log('rename cannot rename (%s->%s) - must be on same box'%(self.hostname, target.hostname))
                return False
        logger.log("Renaming %s into %s"%(self.printable(),target.printable_simple()))

        ssh=self.ssh
        oldname=self.name
        newname=target.name
        xid=ssh.backquote(['cat','/etc/vservers/%(oldname)s/context'%locals()])
        currently_running=self.is_running()
        result=True
        if currently_running:
            result = result and self.stop()
        result = result \
            and e2b(ssh.run_string("mv /vservers/%(oldname)s /vservers/%(newname)s"%locals())) \
            and e2b(ssh.run_string("mv /etc/vservers/%(oldname)s /etc/vservers/%(newname)s"%locals())) \
            and e2b(ssh.run_string("rm /etc/vservers/%(newname)s/run"%locals())) \
            and e2b(ssh.run_string("ln -s /var/run/vservers/%(newname)s /etc/vservers/%(newname)s/run"%locals())) \
            and e2b(ssh.run_string("rm /etc/vservers/%(newname)s/vdir"%locals())) \
            and e2b(ssh.run_string("ln -s /etc/vservers/.defaults/vdirbase/%(newname)s /etc/vservers/%(newname)s/vdir"%locals())) \
            and e2b(ssh.run_string("rm /etc/vservers/%(newname)s/cache"%locals())) \
            and e2b(ssh.run_string("ln -s /etc/vservers/.defaults/cachebase/%(newname)s /etc/vservers/%(newname)s/cache"%locals())) \
            and e2b(ssh.run_string("rm /var/run/vservers.rev/%(xid)s"%locals())) \
            and e2b(ssh.run_string("ln -s /etc/vserver/%(newname)s /var/run/vservers.rev/%(xid)s"%locals()))
        if currently_running:
            result = result and self.start()
        return result
            
    def ok(self,message): logger.log('OK %s: %s'%(self.printable(),message))
    def error(self,message): logger.log('ERROR %s: %s'%(self.printable(),message))
    def warning(self,message): logger.log( 'WARNING %s: %s'%(self.printable(),message))

    def netcheck(self):
        result=True
        ssh=self.ssh
        name=self.name
        hostname=ssh.backquote_string("cat /etc/vservers/%(name)s/uts/nodename"%locals()).strip()
        if not hostname: 
            self.error('no uts/nodename')
            result=False
        ip=ssh.backquote_string("cat /etc/vservers/%(name)s/interfaces/0/ip"%locals()).strip()
        if not ip: 
            self.error('no interfaces/0/ip')
            result=False
        try:
            expected_ip=socket.gethostbyname(hostname)
            if expected_ip == ip:
                self.ok('hostname %s maps to IP %s'%(hostname,ip))
            else:
                self.error('IP mismatch, current %s -- expected %s'%(ip,expected_ip))
                result=False
        except:
            self.warning ('could not resolve hostname <%s>'%hostname)
            expected_ip='255.255.255.255'
            result=False
        mask_path="/etc/vservers/%(name)s/interfaces/mask"%locals()
        mask_exists = e2b(ssh.run_string('ls %s'%mask_path,trash_stderr=True))
        if mask_exists:
            self.error('found a mask file in %s'%(self.printable(),mask_path))
            result=False
        else:
            self.ok('no mask file %s'%mask_path)
        return result
        
class Main:

    modes={ 
        # expected arg numbers can be either
        # negative int, e.g. -1 means at least one arg
        # positive int, exactly that number of args
        # list of ints : allowed number of args
        'probe' : "hostname(s)\n\tprobes all hostnames and returns the found vservers",
        'netcheck' : "vsname\n\tchecks the network config for a vserver",
        'rename' : "oldname newname\n\trename a vserver - both named must be local",
        'migrate' : "oldname hostname@newname\n\tmigrate a vserver over to a new box; must be run on the current box\n\tnot implemented yet",
        }

    def read_args_from_file (self,args_from_file):
        if not args_from_file: return []
        try:
            raw=file(args_from_file).read().split()
            return [x for x in raw if x.find('#')<0]
        except:
            logger.log('ERROR: Could not open args file %s'%args_from_file)
            return []

    def run (self):

        mode=None
        for function in Main.modes.keys():
            if sys.argv[0].find(function) >= 0:
                mode = function
                mode_desc = Main.modes[mode]
                break
        if not mode:
            print "Unsupported command",sys.argv[0]
            print "Supported commands:" + " ".join(Main.modes.keys())
            sys.exit(1)

        parser=OptionParser(usage="%prog " + mode_desc)
        parser.add_option('-n','--dry-run',dest='dry_run',action='store_true',default=False,
                          help="dry run")
        if mode == 'probe' or mode == 'netcheck':
            parser.add_option ('-f','--file',action='store',dest='args_from_file',default=None,
                               help='read args from file')

        (options,args) = parser.parse_args()
        parser.set_usage("%prog " + mode_desc)

        # check number of arguments
        def bail_out (msg):
            parser.print_help(msg)
            sys.exit(1)

        # 
        if mode == 'probe':
            names=[]
            if not options.args_from_file and len(args)==0: bail_out('no arg specified')
            for arg in args + self.read_args_from_file(options.args_from_file):
                names += Vserver.probe_host(arg,options.dry_run)
            names.sort()
            for name in names: print name
                
        elif mode=='netcheck':
            result=True
            if not options.args_from_file and len(args)==0: bail_out('no arg specified')
            for arg in args + self.read_args_from_file(options.args_from_file):
                if not b2e(Vserver(arg,options.dry_run).netcheck()): result=False
            return result
        elif mode=='rename':
            [x,y]=args
            return b2e(Vserver(x,options.dry_run).rename(y))
        elif mode=='migrate':
            print 'migrate not yet implemented'
            return 1
        else:
            print 'unknown mode %s'%mode
            return b2e(False)

if __name__ == '__main__':
    sys.exit( Main().run() )