+++ /dev/null
-import json
-import web
-import subprocess
-
-from config import DB
-
-import logging
-log = logging.getLogger(__name__)
-
-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)