server/base_utils.py

# Copyright 2008 Google Inc, Martin J. Bligh <mbligh@google.com>,
#                Benjamin Poirier, Ryan Stutsman
# Released under the GPL v2
"""
Miscellaneous small functions.

DO NOT import this file directly - it is mixed in by server/utils.py,
import that instead
"""

import atexit, os, re, shutil, textwrap, sys, tempfile, types

from autotest_lib.client.common_lib import barrier, utils
from autotest_lib.server import subcommand


# A dictionary of pid and a list of tmpdirs for that pid
__tmp_dirs = {}


def scp_remote_escape(filename):
    """
    Escape special characters from a filename so that it can be passed
    to scp (within double quotes) as a remote file.

    Bis-quoting has to be used with scp for remote files, "bis-quoting"
    as in quoting x 2
    scp does not support a newline in the filename

    Args:
            filename: the filename string to escape.

    Returns:
            The escaped filename string. The required englobing double
            quotes are NOT added and so should be added at some point by
            the caller.
    """
    escape_chars= r' !"$&' "'" r'()*,:;<=>?[\]^`{|}'

    new_name= []
    for char in filename:
        if char in escape_chars:
            new_name.append("\\%s" % (char,))
        else:
            new_name.append(char)

    return utils.sh_escape("".join(new_name))


def get(location, local_copy = False):
    """Get a file or directory to a local temporary directory.

    Args:
            location: the source of the material to get. This source may
                    be one of:
                    * a local file or directory
                    * a URL (http or ftp)
                    * a python file-like object

    Returns:
            The location of the file or directory where the requested
            content was saved. This will be contained in a temporary
            directory on the local host. If the material to get was a
            directory, the location will contain a trailing '/'
    """
    tmpdir = get_tmp_dir()

    # location is a file-like object
    if hasattr(location, "read"):
        tmpfile = os.path.join(tmpdir, "file")
        tmpfileobj = file(tmpfile, 'w')
        shutil.copyfileobj(location, tmpfileobj)
        tmpfileobj.close()
        return tmpfile

    if isinstance(location, types.StringTypes):
        # location is a URL
        if location.startswith('http') or location.startswith('ftp'):
            tmpfile = os.path.join(tmpdir, os.path.basename(location))
            utils.urlretrieve(location, tmpfile)
            return tmpfile
        # location is a local path
        elif os.path.exists(os.path.abspath(location)):
            if not local_copy:
                if os.path.isdir(location):
                    return location.rstrip('/') + '/'
                else:
                    return location
            tmpfile = os.path.join(tmpdir, os.path.basename(location))
            if os.path.isdir(location):
                tmpfile += '/'
                shutil.copytree(location, tmpfile, symlinks=True)
                return tmpfile
            shutil.copyfile(location, tmpfile)
            return tmpfile
        # location is just a string, dump it to a file
        else:
            tmpfd, tmpfile = tempfile.mkstemp(dir=tmpdir)
            tmpfileobj = os.fdopen(tmpfd, 'w')
            tmpfileobj.write(location)
            tmpfileobj.close()
            return tmpfile


def get_tmp_dir():
    """Return the pathname of a directory on the host suitable
    for temporary file storage.

    The directory and its content will be deleted automatically
    at the end of the program execution if they are still present.
    """
    dir_name = tempfile.mkdtemp(prefix="autoserv-")
    pid = os.getpid()
    if not pid in __tmp_dirs:
        __tmp_dirs[pid] = []
    __tmp_dirs[pid].append(dir_name)
    return dir_name


def __clean_tmp_dirs():
    """Erase temporary directories that were created by the get_tmp_dir()
    function and that are still present.
    """
    pid = os.getpid()
    if pid not in __tmp_dirs:
        return
    for dir in __tmp_dirs[pid]:
        try:
            shutil.rmtree(dir)
        except OSError, e:
            if e.errno == 2:
                pass
    __tmp_dirs[pid] = []
atexit.register(__clean_tmp_dirs)
subcommand.subcommand.register_join_hook(lambda _: __clean_tmp_dirs())


def unarchive(host, source_material):
    """Uncompress and untar an archive on a host.

    If the "source_material" is compresses (according to the file
    extension) it will be uncompressed. Supported compression formats
    are gzip and bzip2. Afterwards, if the source_material is a tar
    archive, it will be untarred.

    Args:
            host: the host object on which the archive is located
            source_material: the path of the archive on the host

    Returns:
            The file or directory name of the unarchived source material.
            If the material is a tar archive, it will be extracted in the
            directory where it is and the path returned will be the first
            entry in the archive, assuming it is the topmost directory.
            If the material is not an archive, nothing will be done so this
            function is "harmless" when it is "useless".
    """
    # uncompress
    if (source_material.endswith(".gz") or
            source_material.endswith(".gzip")):
        host.run('gunzip "%s"' % (utils.sh_escape(source_material)))
        source_material= ".".join(source_material.split(".")[:-1])
    elif source_material.endswith("bz2"):
        host.run('bunzip2 "%s"' % (utils.sh_escape(source_material)))
        source_material= ".".join(source_material.split(".")[:-1])

    # untar
    if source_material.endswith(".tar"):
        retval= host.run('tar -C "%s" -xvf "%s"' % (
                utils.sh_escape(os.path.dirname(source_material)),
                utils.sh_escape(source_material),))
        source_material= os.path.join(os.path.dirname(source_material),
                retval.stdout.split()[0])

    return source_material


def get_server_dir():
    path = os.path.dirname(sys.modules['autotest_lib.server.utils'].__file__)
    return os.path.abspath(path)


def find_pid(command):
    for line in utils.system_output('ps -eo pid,cmd').rstrip().split('\n'):
        (pid, cmd) = line.split(None, 1)
        if re.search(command, cmd):
            return int(pid)
    return None


def nohup(command, stdout='/dev/null', stderr='/dev/null', background=True,
                                                                env = {}):
    cmd = ' '.join(key+'='+val for key, val in env.iteritems())
    cmd += ' nohup ' + command
    cmd += ' > %s' % stdout
    if stdout == stderr:
        cmd += ' 2>&1'
    else:
        cmd += ' 2> %s' % stderr
    if background:
        cmd += ' &'
    utils.system(cmd)


def default_mappings(machines):
    """
    Returns a simple mapping in which all machines are assigned to the
    same key.  Provides the default behavior for
    form_ntuples_from_machines. """
    mappings = {}
    failures = []

    mach = machines[0]
    mappings['ident'] = [mach]
    if len(machines) > 1:
        machines = machines[1:]
        for machine in machines:
            mappings['ident'].append(machine)

    return (mappings, failures)


def form_ntuples_from_machines(machines, n=2, mapping_func=default_mappings):
    """Returns a set of ntuples from machines where the machines in an
       ntuple are in the same mapping, and a set of failures which are
       (machine name, reason) tuples."""
    ntuples = []
    (mappings, failures) = mapping_func(machines)

    # now run through the mappings and create n-tuples.
    # throw out the odd guys out
    for key in mappings:
        key_machines = mappings[key]
        total_machines = len(key_machines)

        # form n-tuples
        while len(key_machines) >= n:
            ntuples.append(key_machines[0:n])
            key_machines = key_machines[n:]

        for mach in key_machines:
            failures.append((mach, "machine can not be tupled"))

    return (ntuples, failures)


def parse_machine(machine, user = 'root', port = 22, password = ''):
    """
    Parse the machine string user:pass@host:port and return it separately,
    if the machine string is not complete, use the default parameters
    when appropriate.
    """

    user = user
    port = port
    password = password

    if re.search('@', machine):
        machine = machine.split('@')

        if re.search(':', machine[0]):
            machine[0] = machine[0].split(':')
            user = machine[0][0]
            password = machine[0][1]

        else:
            user = machine[0]

        if re.search(':', machine[1]):
            machine[1] = machine[1].split(':')
            hostname = machine[1][0]
            port = int(machine[1][1])

        else:
            hostname = machine[1]

    elif re.search(':', machine):
        machine = machine.split(':')
        hostname = machine[0]
        port = int(machine[1])

    else:
        hostname = machine

    return hostname, user, password, port


def get_public_key():
    """
    Return a valid string ssh public key for the user executing autoserv or
    autotest. If there's no DSA or RSA public key, create a DSA keypair with
    ssh-keygen and return it.
    """

    ssh_conf_path = os.path.expanduser('~/.ssh')

    dsa_public_key_path = os.path.join(ssh_conf_path, 'id_dsa.pub')
    dsa_private_key_path = os.path.join(ssh_conf_path, 'id_dsa')

    rsa_public_key_path = os.path.join(ssh_conf_path, 'id_rsa.pub')
    rsa_private_key_path = os.path.join(ssh_conf_path, 'id_rsa')

    has_dsa_keypair = os.path.isfile(dsa_public_key_path) and \
        os.path.isfile(dsa_private_key_path)
    has_rsa_keypair = os.path.isfile(rsa_public_key_path) and \
        os.path.isfile(rsa_private_key_path)

    if has_dsa_keypair:
        print 'DSA keypair found, using it'
        public_key_path = dsa_public_key_path

    elif has_rsa_keypair:
        print 'RSA keypair found, using it'
        public_key_path = rsa_public_key_path

    else:
        print 'Neither RSA nor DSA keypair found, creating DSA ssh key pair'
        utils.system('ssh-keygen -t dsa -q -N "" -f %s' % dsa_private_key_path)
        public_key_path = dsa_public_key_path

    public_key = open(public_key_path, 'r')
    public_key_str = public_key.read()
    public_key.close()

    return public_key_str


def get_sync_control_file(control, host_name, host_num,
                          instance, num_jobs, port_base=63100):
    """
    This function is used when there is a need to run more than one
    job simultaneously starting exactly at the same time. It basically returns
    a modified control file (containing the synchronization code prepended)
    whenever it is ready to run the control file. The synchronization
    is done using barriers to make sure that the jobs start at the same time.

    Here is how the synchronization is done to make sure that the tests
    start at exactly the same time on the client.
    sc_bar is a server barrier and s_bar, c_bar are the normal barriers

                      Job1              Job2         ......      JobN
     Server:   |                        sc_bar
     Server:   |                        s_bar        ......      s_bar
     Server:   |      at.run()         at.run()      ......      at.run()
     ----------|------------------------------------------------------
     Client    |      sc_bar
     Client    |      c_bar             c_bar        ......      c_bar
     Client    |    <run test>         <run test>    ......     <run test>

    @param control: The control file which to which the above synchronization
            code will be prepended.
    @param host_name: The host name on which the job is going to run.
    @param host_num: (non negative) A number to identify the machine so that
            we have different sets of s_bar_ports for each of the machines.
    @param instance: The number of the job
    @param num_jobs: Total number of jobs that are going to run in parallel
            with this job starting at the same time.
    @param port_base: Port number that is used to derive the actual barrier
            ports.

    @returns The modified control file.
    """
    sc_bar_port = port_base
    c_bar_port = port_base
    if host_num < 0:
        print "Please provide a non negative number for the host"
        return None
    s_bar_port = port_base + 1 + host_num # The set of s_bar_ports are
                                          # the same for a given machine

    sc_bar_timeout = 180
    s_bar_timeout = c_bar_timeout = 120

    # The barrier code snippet is prepended into the conrol file
    # dynamically before at.run() is called finally.
    control_new = []

    # jobid is the unique name used to identify the processes
    # trying to reach the barriers
    jobid = "%s#%d" % (host_name, instance)

    rendv = []
    # rendvstr is a temp holder for the rendezvous list of the processes
    for n in range(num_jobs):
        rendv.append("'%s#%d'" % (host_name, n))
    rendvstr = ",".join(rendv)

    if instance == 0:
        # Do the setup and wait at the server barrier
        # Clean up the tmp and the control dirs for the first instance
        control_new.append('if os.path.exists(job.tmpdir):')
        control_new.append("\t system('umount -f %s > /dev/null"
                           "2> /dev/null' % job.tmpdir,"
                           "ignore_status=True)")
        control_new.append("\t system('rm -rf ' + job.tmpdir)")
        control_new.append(
            'b0 = job.barrier("%s", "sc_bar", %d, port=%d)'
            % (jobid, sc_bar_timeout, sc_bar_port))
        control_new.append(
        'b0.rendezvous_servers("PARALLEL_MASTER", "%s")'
         % jobid)

    elif instance == 1:
        # Wait at the server barrier to wait for instance=0
        # process to complete setup
        b0 = barrier.barrier("PARALLEL_MASTER", "sc_bar", sc_bar_timeout,
                     port=sc_bar_port)
        b0.rendezvous_servers("PARALLEL_MASTER", jobid)

        if(num_jobs > 2):
            b1 = barrier.barrier(jobid, "s_bar", s_bar_timeout,
                         port=s_bar_port)
            b1.rendezvous(rendvstr)

    else:
        # For the rest of the clients
        b2 = barrier.barrier(jobid, "s_bar", s_bar_timeout, port=s_bar_port)
        b2.rendezvous(rendvstr)

    # Client side barrier for all the tests to start at the same time
    control_new.append('b1 = job.barrier("%s", "c_bar", %d, port=%d)'
                    % (jobid, c_bar_timeout, c_bar_port))
    control_new.append("b1.rendezvous(%s)" % rendvstr)

    # Stick in the rest of the control file
    control_new.append(control)

    return "\n".join(control_new)