#!/usr/bin/python
#

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,level=0):
        if level>0: print >> sys.stderr,level*4*'=',
        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

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

    def header (self,message):
        logger.log(message,level=1)

    # convenience : if argv is a string, we just split it to get a list
    # if args may contain spaces, it's safer to build the list yourself
    def normalize (self, argv):
        if isinstance (argv,list): return argv
        if isinstance (argv,str): return argv.split()
        else: raise Exception, "Unsupported command type in Shell - %r"%argv

    # return an exit code
    # in general dry_run ... does not run
    # there are cases where the command is harmless and required for the rest to proceed correctly
    # in this case you can set force_dry_run
    def run (self, argv, trash_stdout=False, trash_stderr=False, message="", force_dry_run=False):
        argv=self.normalize(argv)
        if self.verbose and message: self.header(message),
        if self.dry_run and not force_dry_run:
            print 'WOULD RUN '+" ".join(argv)
            return 0
        else:
            helper='RUNNING'
            if force_dry_run and self.verbose: helper += ' (forced)'
            if self.verbose: logger.log (helper+' '+" ".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 `...` - cannot trash stdout of course
    # ditto for the dry_run mode
    def backquote (self, argv, trash_stderr=False,message="", force_dry_run=False):
        argv=self.normalize (argv)
        if self.verbose and message: self.header(message),
        if self.dry_run and not force_dry_run:
            print 'WOULD RUN $('+" ".join(argv) + ')'
            return ""
        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 (Shell):
    
    #################### the class
    def __init__(self,hostname=None,key=None, username='root',dry_run=False,verbose=False):
        self.hostname=hostname
        if not hostname: self.hostname='localhost'
        self.key=key
        self.username=username
        Shell.__init__(self,dry_run=dry_run,verbose=verbose)

    # 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

    def run(self, argv,*args, **kwds):
        return Shell.run(self,self.argv(self.normalize(argv)),*args, **kwds)
    def backquote(self, argv, *args, **kwds):
        return Shell.backquote(self,self.argv(self.normalize(argv)),*args, **kwds)

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

    def __init__ (self, name, dry_run=False,verbose=False):
        self.dry_run=dry_run
        self.verbose=verbose
        self._running_state=None
        self._xid=None
        self._autostart=None
        try:
            (self.name,hostname)=name.split("@")
        except:
            (self.name,hostname)=(name,'localhost')
        self.set_hostname(hostname)

    def set_hostname (self, hostname):
        self.hostname=hostname
        self.ssh=Ssh(self.hostname,username='root',dry_run=self.dry_run,verbose=self.verbose)

    def ok(self,message): 
        if self.verbose: logger.log('OK %s: %s'%(self.printable(),message))
    def warning(self,message): 
        if self.verbose: logger.log( 'WARNING %s: %s'%(self.printable(),message))
    def error(self,message): 
        logger.log('ERROR %s: %s'%(self.printable(),message))

    # cache run mode
    def running_state (self):
        if self._running_state is not None: return self._running_state
        status_retcod= self.ssh.run('vserver %s status'%self.name,
                                    trash_stdout=True,trash_stderr=True,
                                    message='retrieving status for %s'%self.name,
                                    force_dry_run=True)
        if status_retcod == 0:          self._running_state = 'start'
        elif status_retcod == 3:        self._running_state = 'stop'
        else:                           self._running_state = 'unknown'
        return self._running_state

    def running_char (self):
        if self._running_state is None:         return '?'
        elif self._running_state == 'unknown' : return '!'
        elif self._running_state == 'stop':     return '-'
        else:                                   return '+'

    # cache xid
    def xid (self):
        if self._xid is not None: return self._xid
        self._xid = self.ssh.backquote(['cat','/etc/vservers/%s/context'%self.name],
                                       force_dry_run=True,trash_stderr=True,
                                       message='Retrieving xid for %s'%self.name).strip()
        return self._xid

    def autostart (self):
        if self._autostart is not None: return self._autostart
        self._autostart = self.ssh.backquote(['cat','/etc/vservers/%s/apps/init/mark'%self.name],
                                             force_dry_run=True,trash_stderr=True,
                                             message='Retrieving autostart status for %s'%self.name).strip()=='default'
        return self._autostart

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

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

    def printable_simple(self): return self.name + '@' + self.hostname
    def printable_dbg(self): 
        autostart_part='<*>'
        if not self.autostart(): autostart_part='< >'
        return autostart_part+" (%s) [%s] %s@%s"%(self.running_char(),self.xid(),self.name,self.hostname)
    def printable(self): 
        # force fetch cache info
        self.running_state()
        self.xid()
        self.autostart()
        return self.printable_dbg()

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

    # does it have a pattern
    def has_pattern (self):
        return self.name.find('*') >= 0

    # probe it
    def probe_pattern (self):
        pattern=self.name
        ls=self.ssh.backquote("ls -d1 /vservers/%s"%pattern,
                              trash_stderr=True,force_dry_run=True,
                              message='scanning /vservers with pattern %s on %s'%(pattern,self.hostname))
        return [ '%s@%s'%(path.replace('/vservers/',''),self.hostname) for path in ls.split()]

    #################### toplevel methods for main
    def show (self, long_format):
        if long_format:
            print self.printable()
        else:
            print self.printable_simple()

    def check(self):
        result=True
        ssh=self.ssh
        name=self.name
        hostname=ssh.backquote("cat /etc/vservers/%(name)s/uts/nodename"%locals()).strip()
        if not hostname: 
            self.warning('no uts/nodename')
        ip=ssh.backquote("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('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
        
    def delete (self):
        # ommitting --silent just silently returns - stdin must be closed or something
        command=["vserver","--silent",self.name,"delete"]
        return e2b (self.ssh.run (command))

    # 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('ERROR:rename cannot rename (%s->%s) - must be on same box'%(self.hostname, target.hostname))
                return False
        logger.log("%s into %s"%(self.printable(),target.printable_simple()),level=2)
        
        ssh=self.ssh
        oldname=self.name
        newname=target.name
        xid=self.xid()
        currently_running=self.running_state()=='start'
        result=True
        if currently_running:
            result = result and self.stop()
        result = result and e2b(ssh.run("mv /vservers/%(oldname)s /vservers/%(newname)s"%locals(),
                                        message="tweaking /vservers/<>"))
        result = result and e2b(ssh.run("mv /etc/vservers/%(oldname)s /etc/vservers/%(newname)s"%locals(),
                                        message="tweaking /etc/vservers/<>")) 
        result = result and e2b(ssh.run("mv /vservers/.pkg/%(oldname)s /vservers/.pkg/%(newname)s"%locals(),
                                        message="renaming /vservers/.pkg/<>")) 
        result = result and e2b(ssh.run("rm /etc/vservers/%(newname)s/run"%locals(),
                                        message="cleaning /etc/vservers/<>/run")) 
        result = result and e2b(ssh.run("ln -s /var/run/vservers/%(newname)s /etc/vservers/%(newname)s/run"%locals(),
                                        message="adjusting /etc/vservers/<>/run")) 
        result = result and e2b(ssh.run("rm /etc/vservers/%(newname)s/vdir"%locals(),
                                        message="cleaning /etc/vservers/<>/vdir")) 
        result = result and e2b(ssh.run("ln -s /etc/vservers/.defaults/vdirbase/%(newname)s /etc/vservers/%(newname)s/vdir"%locals(),
                                        message="adjusting /etc/vservers/<>/vdir")) 
        result = result and e2b(ssh.run("rm /etc/vservers/%(newname)s/cache"%locals(),
                                        message="cleaning /etc/vservers/<>/cache")) 
        result = result and e2b(ssh.run("ln -s /etc/vservers/.defaults/cachebase/%(newname)s /etc/vservers/%(newname)s/cache"%locals(),
                                        message="adjusting /etc/vservers/<>/cache")) 
        result = result and e2b(ssh.run("rm /var/run/vservers.rev/%(xid)s"%locals(),
                                        message="cleaning /var/run/vservers.rev/<>")) 
        result = result and e2b(ssh.run("ln -s /etc/vservers/%(newname)s /var/run/vservers.rev/%(xid)s"%locals(),
                                        message="adjusting /var/run/vservers.rev/<>")) 
        result = result and e2b(ssh.run("sed -i -e s,%(oldname)s,%(newname)s,g /etc/vservers/%(newname)s/name"%locals(),
                                        message="adjusting /etc/vservers/<>/name"))
        result = result and e2b(ssh.run("sed -i -e s,%(oldname)s,%(newname)s,g /etc/vservers/%(newname)s/uts/nodename"%locals(),
                                        message="adjusting /etc/vservers/<>/uts/nodename"))
        # ignore rsult
        e2b(ssh.run("sed -i -e s,\\[%(oldname)s\\],\\[%(newname)s\\], /vservers/%(newname)s/root/.profile"%locals(),
                                        message="adjusting /vservers/<>/root/.profile"))

#        # refreshing target instance
        target=Vserver(newname,dry_run=self.dry_run,verbose=self.verbose)
        target.set_hostname(self.hostname)
        if currently_running and not self.dry_run: 
            result = result and target.start()
        print target.printable()
        return result
            
class VserverTools:

    modes={ 
        'show' : "vsname(s)\n\tdisplay status & details",
        'start': "vsname(s)\n\tstart vservers",
        'stop' : "vsname(s)\n\tstop vservers",
        'check' : "vsname(s)\n\tsanity check - ip from hostname, netmask, ...",
        'delete' : "vsname(s)\n\tdelete vservers",
        '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\n\tnot implemented yet",
        }

    def __init__ (self):
        self.parser=OptionParser()
        self.args=None
        self.options=None

    def vserver(self, vsname):
        return Vserver(vsname,dry_run=self.options.dry_run,verbose=self.options.verbose)

    def probe_star_arg (self, arg):
        vs=self.vserver(arg)
        if vs.has_pattern():    return vs.probe_pattern()
        else:                   return [arg,]

    def read_args_from_file (self):
        if not self.options.args_from_file: return []
        try:
            result=[]
            raw_args=file(self.options.args_from_file).read().split()
            for arg in raw_args:
                result += self.probe_star_arg(arg)
            return result
        except:
            logger.log('ERROR: Could not open args file %s'%self.options.args_from_file)
            return []

    def check_usual_args(self):
        if not self.options.args_from_file and len(self.args)==0: 
            print ('no arg specified')
            self.parser.print_help()
            sys.exit(1)
        result=[]
        for arg in self.args: result += self.probe_star_arg(arg)
        result += self.read_args_from_file()
        result.sort()
        return result
        
    # return exit code
    def run (self):

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

        self.parser.set_usage("%prog " + mode_desc)
        self.parser.add_option('-n','--dry-run',dest='dry_run',action='store_true',default=False,
                          help="dry run")
        self.parser.add_option('-v','--verbose',dest='verbose',action='store_true',default=False,
                          help="be verbose")
        self.parser.add_option ('-l','--long',action='store_true', dest='long_format',default=False,
                                help="show more detailed result")
        ##########
        if mode not in ['rename','migrate'] :
            self.parser.add_option ('-f','--file',action='store',dest='args_from_file',default=None,
                               help="read args from file")

        (self.options,self.args) = self.parser.parse_args()

        if mode == 'show':
            for arg in self.check_usual_args():
                self.vserver(arg).show(self.options.long_format)
        elif mode == 'stop':
            for arg in self.check_usual_args():
                vs = self.vserver(arg)
                vs.stop()
                vs.show(True)
        elif mode == 'start':
            for arg in self.check_usual_args():
                vs = self.vserver(arg)
                vs.start()
                vs.show(True)
        elif mode == 'check':
            result=True
            for arg in self.check_usual_args():
                if not b2e(self.vserver(arg).check()): result=False
            return result
        elif mode == 'delete':
            result=True
            for arg in self.check_usual_args():
                if not b2e(self.vserver(arg).delete()): result=False
                self.vserver(arg).show(True)
            return result
        elif mode == 'rename':
            [x,y]=self.args
            return b2e(self.vserver(x).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( VserverTools().run() )