source: fedd/fixed_resource.py @ e087a7a

#!/usr/local/bin/python

import sys
import MySQLdb
from federation.fixed_resource import read_key_db, write_key_db, \
        read_project_db, write project_db
from optparse import OptionParser

actions = set(("initall", "init", "add", "delete"))

def read_key_file(file):
    """
    Read a single line from a keyfile and trim the whitespace
    """
    f = open(file, 'r')
    key = f.readline().rstrip()
    f.close()
    return key

class key_opts(OptionParser):
    """
    Options to the command line, pretty self describing
    """
    def __init__(self):
        OptionParser.__init__(self, usage="%prog " + \
                "(init|add|delete|initall) " + \
                "[opts] (--help for details)", version="0.1")
        self.add_option('-t','--type', dest='type', type='choice', 
                choices=['keys','projects'], help="database type")
        self.add_option('-u','--user', dest='user', default=None,
                action='store', help='user to add/delete')
        self.add_option('-p','--project', dest='project', default=None,
                action='store', help='project to add/delete')
        self.add_option('-k','--key', dest='key', default=None,
                action='store', help='key to add/delete (string)')
        self.add_option('-K','--keyfile', dest='keyfile', default=None,
                action='store', help='key to add/delete (file)')
        self.add_option('-d', '--database', dest='file', default=None,
                action='store', help='database file')
        self.add_option('-f', '--file', dest='file', action='store',
                help='database file (synonym for --database)')


parser = key_opts()
action = sys.argv[1]

# Check the action
if action in actions: del sys.argv[1]
else: sys.exit("Bad action, must be one of %s" % ", ".join(actions))

# Parse and check the the options for consistency
(opts, args) = parser.parse_args()
if not opts.file:
    sys.exit("Must specify DB file")
if not opts.type:
    sys.exit("Must specify database type (--type)")
elif opts.type == 'keys':
    if action != "init" and action != "initall" and action != "addall":
        if opts.user and (opts.key or opts.keyfile):
            user = opts.user
            if opts.key:
                key = opts.key
            elif opts.keyfile: 
                try:
                    key = read_key_file(opts.keyfile)
                except IOError, e:
                    sys.exit("Error reading keyfile: %s" % e)
        else:
            sys.exit("Must specify user and key")
elif opts.type == 'projects':
    if action != "init" and action != "initall" and action != "addall":
        if opts.project: project = opts.project
        else: sys.exit("Must specify project")
else:
    sys.exit("Invalid --type field (how'd you do that?)")

if opts.type == 'keys':
    if action != 'init' and action != 'initall':
        keys = read_key_db(opts.file)
    else:
        keys = set()
    
    if action == 'initall':
        # Add all users from the Emulab DB
        try:
            db = MySQLdb.connect(db="tbdb")
        except:
            sys.exit("Cannot access the Emulab database")

        c = db.cursor()
        c.execute("SELECT uid, pubkey FROM user_pubkeys")
        for u, k in c.fetchall():
            keys.add((u, k))
        c.close()
        db.close()
    elif action == 'add':
        keys.add((user, key))
    elif action == 'delete': 
        if (user, key) in keys:
            keys.remove((user, key))
        else:
            print >>sys.stderr, "Cannot delete (%s, %s): not in db" % \
                    (user, key)
    # init action falls through to write an empty DB
    write_key_db(opts.file, keys)
else:
    if action != 'init' and action != 'initall':
        projects = read_project_db(opts.file)
    else:
        projects = set()
    
    if action == 'initall':
        # Add all projects from the Emulab DB
        try:
            db = MySQLdb.connect(db="tbdb")
        except:
            sys.exit("Cannot access the Emulab database")

        c = db.cursor()
        c.execute("SELECT pid FROM projects")
        for p in c.fetchall():
            projects.add(p)
        c.close()
        db.close()
    elif action == 'add':
        projects.add(project)
    elif action == 'delete': 
        if project in projects:
            projects.remove(project)
        else:
            print >>sys.stderr, "Cannot delete %s: not in db" % project

    # init action falls through to write an empty DB
    write_project_db(opts.file, projects)

sys.exit(0)