zvapp

#!/usr/bin/python

import shutil
import sys
import re
import tarfile
from tempfile import mkstemp, mkdtemp

from eventlet.green.subprocess import Popen, PIPE
from eventlet.green import os
from eventlet import GreenPool
from zvshlib.zvsh import ZvRunner, ZvArgs, ZvConfig


try:
    import simplejson as json
except ImportError:
    import json

try:
    from zerocloud.configparser import ClusterConfigParser, \
        ClusterConfigParsingError, NodeEncoder
    from zerocloud.common import parse_location, \
        CLUSTER_CONFIG_FILENAME, NODE_CONFIG_FILENAME, \
        ACCESS_WRITABLE, ACCESS_READABLE, ACCESS_CDR, ACCESS_NETWORK, \
        SwiftPath, ImagePath
    from zerocloud.proxyquery import NameService
except ImportError:
    sys.stderr.write('Please install python-zerocloud before running %s\n'
                     % sys.argv[0])
    sys.exit(1)


class AppRunner(ZvRunner):

    def __init__(self, command_line, report_file):
        self.command = command_line
        self.process = None
        self.report = ''
        self.report_file = report_file

    def run(self):
        try:
            self.process = Popen(self.command, stdout=PIPE)
            rep_reader = self.spawn(True, self.report_reader)
            self.process.wait()
            rep_reader.join()
        except (KeyboardInterrupt, Exception):
            pass
        finally:
            if self.process:
                self.process.wait()
                fd = open(self.report_file, 'wb')
                fd.write(self.report)
                fd.close()


class AppArgs(ZvArgs):

    def add_agruments(self):
        self.parser.add_argument('exec_file',
                                 help='ZeroVM application archive '
                                      'or map file\n')
        self.parser.add_argument('--swift-root-path',
                                 help='Root path for resolving swift:// urls, '
                                      'ex.\n'
                                      'swift://account/container/object ->'
                                      ' swift-root-path/account/container/'
                                      'object\n')
        self.parser.add_argument('--swift-account-path',
                                 help='Account path for all urls\n'
                                      'to serve all swift://account/ '
                                      'urls from one path')
        self.parser.add_argument('--sysimage-root-path',
                                 help='Root path for system image files\n')
        self.parser.add_argument('--zvm-save-dir',
                                 help='Save ZeroVM environment files into '
                                      'provided directory,\n'
                                      'directory will be created/re-created\n')
        self.parser.add_argument('--dry-run',
                                 help='Print the resulting job descriptions '
                                      'and exit\n',
                                 action='store_true')


class ZvLocalFilesystem(object):

    SYSIMAGE_MASK = re.compile(r'(.*?)(\.[^.]+)?$')

    def __init__(self, sysimage_root_path=None,
                 root_path=None, account_path=None, config=None, savedir=None):
        if not root_path:
            root_path = config.get('root_path', None) or ''
        self.image_path = None
        self.root_path = os.path.abspath(root_path)
        self.account_path = account_path
        if not self.account_path:
            self.account_path = config.get('account_path', None)
        if self.account_path:
            self.account_path = os.path.abspath(self.account_path)
        self.savedir = savedir
        if self.savedir:
            self.tempdir = os.path.abspath(self.savedir)
            if os.path.isdir(self.tempdir):
                shutil.rmtree(self.tempdir)
            os.makedirs(self.tempdir)
        else:
            self.tempdir = mkdtemp()
        self.sysimage_devices = {}
        if not sysimage_root_path:
            sysimage_root_path = config.get('sysimage_path', None)
        if sysimage_root_path:
            self.sysimage_devices = \
                self._list_sysimage_devices(
                    os.path.abspath(sysimage_root_path))
        self.immediate_responses = {}

    def list_account(self, account, mask=None):
        account_path = self.account_path
        if not account_path:
            account_path = os.path.join(self.root_path, account)
        ret = []
        for fname in sorted(os.listdir(account_path)):
            if not mask or mask.match(fname):
                ret.append(fname)
        return ret

    def list_container(self, account, container, mask=None):
        account_path = self.account_path
        if not account_path:
            account_path = os.path.join(self.root_path, account)
        container_path = os.path.join(account_path, container)
        ret = []
        for dirpath, _junk, filenames in os.walk(container_path):
            dirpath = dirpath.replace(container_path, '', 1)[1:]
            for fname in sorted(filenames):
                objpath = os.path.join(dirpath, fname)
                if not mask or mask.match(objpath):
                    ret.append(objpath)
        return ret

    def _list_sysimage_devices(self, sysimage_path):
        result = {}
        if os.path.isdir(sysimage_path):
            for fname in os.listdir(sysimage_path):
                path = os.path.join(sysimage_path, fname)
                if os.path.isfile(path):
                    base = os.path.basename(path)
                    sysimage_name = self.SYSIMAGE_MASK.search(base).group(1)
                    result[sysimage_name] = path
        return result

    def get_local_path(self, device, path, access, node_name=None):
        sysimage = self.sysimage_devices.get(device, None)
        if sysimage:
            return sysimage
        if 'image' == device and not path:
            return self.image_path
        loc = parse_location(path)
        if not loc and access & ACCESS_WRITABLE:
            temp_file = self.create_temp_file()
            self.immediate_responses[node_name] = temp_file
            return temp_file
        elif isinstance(loc, SwiftPath):
            account_path = self.account_path
            if not account_path:
                account_path = os.path.join(self.root_path, loc.account)
            local_path = os.path.join(account_path, loc.container, loc.obj)
            return local_path
        elif isinstance(loc, ImagePath) and access & (ACCESS_READABLE |
                                                      ACCESS_CDR):
            if 'image' == loc.image:
                if os.path.isdir(self.image_path):
                    return os.path.join(os.path.abspath(self.image_path),
                                        loc.path)
                else:
                    return self._extract_file(self.image_path, loc.path)
            sysimage = self.sysimage_devices.get(loc.image, None)
            return self._extract_file(sysimage, loc.path)
        elif access & ACCESS_NETWORK:
            return path

    def _extract_file(self, image, file_name):
        tar = tarfile.open(name=image)
        efile = None
        try:
            efile = tar.extractfile(file_name)
        except KeyError:
            pass
        if not efile:
            return None
        (fd, fn) = mkstemp(dir=self.tempdir)
        reader = iter(lambda: efile.read(4096), '')
        for chunk in reader:
            os.write(fd, chunk)
        os.close(fd)
        return fn

    def resolve_local_paths(self, node_config):
        for ch in node_config['channels']:
            ch['lpath'] = self.get_local_path(ch['device'],
                                              ch['path'],
                                              ch['access'],
                                              node_config['name'])

    def cleanup(self):
        if not self.savedir:
            shutil.rmtree(self.tempdir)

    def create_temp_file(self):
        (fd, fn) = mkstemp(dir=self.tempdir)
        os.close(fd)
        return fn

    def get_responses(self):
        data = ''
        for fname in sorted(self.immediate_responses.keys()):
            data += open(self.immediate_responses[fname]).read()
        return data

    def create_temp_files(self, node_name):
        session_dir = os.path.join(self.tempdir, node_name)
        os.makedirs(session_dir)
        result = []
        for n in ['nvram', 'manifest', 'report']:
            result.append(os.path.join(session_dir, n))
        #result.append(session_dir)
        return result

if __name__ == '__main__':
    threadpool = GreenPool()
    nspool = GreenPool(1)
    app_args = AppArgs()
    app_args.parse(sys.argv[1:])
    zvsh_config = ['zvsh.cfg',
                   os.path.expanduser('~/.zvsh.cfg'),
                   '/etc/zvsh.cfg']
    zvconfig = ZvConfig()
    zvconfig.read(zvsh_config)
    cluster_config = None
    local_fs = ZvLocalFilesystem(app_args.args.sysimage_root_path,
                                 app_args.args.swift_root_path,
                                 app_args.args.swift_account_path,
                                 zvconfig['zvapp'])
    image_path = None
    try:
        if os.path.isdir(app_args.args.exec_file):
            # user tries to run an extracted application
            # we will need to create image on the fly
            # to load it in zerovm as a channel
            app_dir = app_args.args.exec_file
            tar_file = local_fs.create_temp_file()
            tar = tarfile.TarFile.open(name=tar_file, mode='w')
            try:
                tar.add(os.path.join(app_dir, CLUSTER_CONFIG_FILENAME),
                        arcname=CLUSTER_CONFIG_FILENAME)
                boot = CLUSTER_CONFIG_FILENAME
            except OSError:
                try:
                    tar.add(os.path.join(app_dir, NODE_CONFIG_FILENAME),
                            arcname=NODE_CONFIG_FILENAME)
                    boot = NODE_CONFIG_FILENAME
                except OSError:
                    sys.stderr.write('Cannot find boot map anywhere in %s\n'
                                     % app_dir)
                    sys.exit(1)

            def tar_filter(tar_info):
                if tar_info.name in [CLUSTER_CONFIG_FILENAME,
                                     NODE_CONFIG_FILENAME]:
                    return None
                return tar_info

            for item in os.listdir(app_dir):
                tar.add(os.path.join(app_dir, item), arcname=item,
                        filter=tar_filter)
            tar.close()
            image_path = os.path.abspath(tar_file)
            cluster_config = json.load(open(os.path.join(app_dir, boot), 'rb'))
        else:
            try:
                cluster_config_fd = open(app_args.args.exec_file, 'rb')
            except IOError, e:
                sys.stderr.write(str(e) + '\n')
                sys.exit(1)
            try:
                # let's load the file as a cluster map
                cluster_config = json.load(cluster_config_fd)
                image_path = os.path.dirname(
                    os.path.abspath(app_args.args.exec_file))
            except Exception:
                # it's not a cluster map file
                # try to load it as zvm app archive
                try:
                    tar = tarfile.open(app_args.args.exec_file)
                    for name in tar.getnames():
                        if name in [CLUSTER_CONFIG_FILENAME,
                                    NODE_CONFIG_FILENAME]:
                            cluster_config = json.load(tar.extractfile(name))
                            break
                    image_path = os.path.abspath(app_args.args.exec_file)
                except tarfile.ReadError:
                    # it's not a tar file, bail out for now
                    sys.stderr.write('Cannot parse the input file %s\n'
                                     % app_args.args.exec_file)
                    sys.exit(1)

        local_fs.image_path = image_path
        threads = {}
        parser = ClusterConfigParser(local_fs.sysimage_devices,
                                     'application/octet-stream',
                                     zvconfig,
                                     local_fs.list_account,
                                     local_fs.list_container)
        try:
            parser.parse(cluster_config,
                         os.path.isfile(local_fs.image_path),
                         account_name='local')
        except ClusterConfigParsingError, e:
            sys.stderr.write(str(e) + '\n')
            exit(1)

        ns_server = None
        if parser.total_count > 1:
            ns_server = NameService(parser.total_count)
            ns_server.start(nspool)
            for node in parser.node_list:
                node.name_service = 'udp:%s:%d' % ('127.0.0.1', ns_server.port)
                parser.build_connect_string(node)
        if app_args.args.dry_run:
            print json.dumps(parser.node_list, cls=NodeEncoder, indent=2)
            exit(0)
        for node in parser.node_list:
            node_config = json.loads(json.dumps(node, cls=NodeEncoder))

            local_fs.resolve_local_paths(node_config)
            nexe_path = local_fs.get_local_path('boot',
                                                node_config['exe'],
                                                ACCESS_READABLE)
            nvram_file, manifest_file, report_file = \
                local_fs.create_temp_files(node_config['name'])
            manifest = parser.prepare_for_standalone(node_config, nvram_file,
                                                     nexe_path, None)
            with open(manifest_file, 'wb') as fd:
                fd.write(manifest)
            command_line = ['zerovm', '-PQ', manifest_file]
            runner = AppRunner(command_line, report_file)
            threads[node_config['name']] = (report_file, node_config['id'],
                                            runner)
        for name in sorted(threads.keys()):
            runner = threads[name][2]
            threadpool.spawn_n(runner.run)
        threadpool.waitall()
        if ns_server:
            ns_server.stop()
        for name in sorted(threads.keys()):
            rfile, node_id, runner = threads[name]
            print "---------- Node: %s id: %s ---------" % (name, node_id)
            print open(rfile).read()
        print "========== Result =========="
        print local_fs.get_responses()
    finally:
        local_fs.cleanup()