ceph/teuthology/locker/api.py

import json
import web
import subprocess

from config import DB

def load_machine(name):
    results = list(DB.select('machine', what='*',
                             where='name = $name',
                             vars=dict(name=name)))
    if not results:
        raise web.NotFound()
    return results[0]

def get_sshkey(name):
    if '@' in name:
        _, name = name.rsplit('@')
    args = ['ssh-keyscan']
    args.append(name)
    p = subprocess.Popen(
        args=args,
        stdout=subprocess.PIPE,
        )
    out, _ = p.communicate()
    pubkey = None
    for key_entry in out.splitlines():
        hostname, pubkey = key_entry.split(' ', 1)
    if not pubkey:
        status = 1
    else:
        status = 0
    return (pubkey), status

def update_sshkey(name, key, type):
    if type == 'vps':
        return
    res = DB.update('machine', where='name = $name AND locked = false',
                    vars=dict(name=name),
                    sshpubkey=key,)
    assert res == 1, 'Failed to update key of machine {name}'.format(name=name)
    print 'Updated key on ', name

class MachineLock:
    def GET(self, name):
        row = load_machine(name)
        row.locked_since = row.locked_since.isoformat()
        web.header('Content-type', 'text/json')
        return json.dumps(row)

    def DELETE(self, name):
        user = web.input('user')['user']
        machine = load_machine(name)
        if not machine.locked:
            raise web.BadRequest()
        if machine.locked_by != user:
            raise web.Forbidden()

        res = DB.update('machine',
                        where='locked = true AND name = $name AND locked_by = $user',
                        vars=dict(name=name, user=user),
                        locked=False, locked_by=None, description=None)
        assert res == 1, 'Failed to unlock machine {name}'.format(name=name)
        print user, 'unlocked', name

    def POST(self, name):
        user = web.input('user')['user']
        desc = web.input(desc=None)['desc']
        machine = load_machine(name)
        if machine.locked:
            raise web.Forbidden()

        if machine.type == 'vps':
            curkey = machine.sshpubkey
        else:
            curkey, getstatus = get_sshkey(name)
            if getstatus != 0:
                curkey = machine.sshpubkey
        if machine.sshpubkey != curkey:
            newkey = curkey
        else:
            newkey = machine.sshpubkey
        res = DB.update('machine', where='name = $name AND locked = false',
                        vars=dict(name=name),
                        locked=True,
                        description=desc,
                        sshpubkey=newkey,
                        locked_by=user,
                        locked_since=web.db.SQLLiteral('NOW()'))
        assert res == 1, 'Failed to lock machine {name}'.format(name=name)
        print user, 'locked single machine', name, 'desc', desc

    def PUT(self, name):
        desc = web.input(desc=None)['desc']
        status = web.input(status=None)['status']
        sshpubkey = web.input(sshpubkey=None)['sshpubkey']

        updated = {}
        if desc is not None:
            updated['description'] = desc
        if status is not None:
            updated['up'] = (status == 'up')
        if sshpubkey is not None:
            updated['sshpubkey'] = sshpubkey

        if not updated:
            raise web.BadRequest()
        DB.update('machine', where='name = $name',
                  vars=dict(name=name), **updated)
        print 'updated', name, 'with', updated, 'desc', desc

class Lock:
    def GET(self):
        rows = list(DB.select('machine', what='*'))
        if not rows:
            raise web.NotFound()
        for row in rows:
            row.locked_since = row.locked_since.isoformat()
        web.header('Content-type', 'text/json')
        return json.dumps(rows)

    def POST(self):
        user = web.input('user')['user']
        desc = web.input(desc=None)['desc']
        num = int(web.input('num')['num'])
        machinetype = dict(machinetype=(web.input(machinetype='plana')['machinetype']))

        if num < 1:
            raise web.BadRequest()

        tries = 0
        check_existing = True
        while True:
            try:
                # transaction will be rolled back if an exception is raised
                with DB.transaction():
                    if desc is not None and check_existing:
                        # if a description is provided, treat it as a
                        # key for locking in case the same run locked
                        # machines in the db successfully before, but
                        # the web server reported failure to it
                        # because the request took too long. Only try
                        # this once per request.
                        check_existing = False
                        results = list(DB.select('machine',
                                                 machinetype, desc, user,
                                                 what='name, sshpubkey',
                                                 where='locked = true AND up = true AND type = $machinetype AND description = $desc AND locked_by = $user',
                                                 limit=num))
                        if len(results) == num:
                            name_keys = {}
                            for row in results:
                                name_keys[row.name] = row.sshpubkey
                            print 'reusing machines', name_keys.keys()
                            break

                    results = list(DB.select('machine', machinetype,
                                             what='name, sshpubkey, type',
                                             where='locked = false AND up = true AND type = $machinetype',
                                             limit=num))
                    if len(results) < num:
                        raise web.HTTPError(status='503 Service Unavailable')
                    name_keys = {}
                    for row in results:
                        if row.type == 'vps':
                            curkey = row.sshpubkey
                        else:
                            curkey, getstatus = get_sshkey(row.name)
                            if getstatus != 0:
                                curkey = row.sshpubkey
                        if row.sshpubkey != curkey:
                            newkey = curkey
                            update_sshkey(row.name, curkey, row.type)
                        else:
                            newkey = row.sshpubkey
                        name_keys[row.name] = newkey
                    where_cond = web.db.sqlors('name = ', name_keys.keys()) \
                        + ' AND locked = false AND up = true'
                    num_locked = DB.update('machine',
                                           where=where_cond,
                                           locked=True,
                                           locked_by=user,
                                           description=desc,
                                           locked_since=web.db.SQLLiteral('NOW()'))
                    assert num_locked == num, 'Failed to lock machines'
            except Exception:
                log.exception("Saw exception")
                tries += 1
                if tries < 10:
                    continue
                raise
            else:
                break

        print user, 'locked', name_keys.keys(), 'desc', desc

        web.header('Content-type', 'text/json')
        return json.dumps(name_keys)