source: fedd/fixed_key.py @ 159a447

#!/usr/local/bin/python

# Encode the set of fixed keys as a colon-separated user/key pairs.  This also
# includes a command line utility to manipulate the DB.

def read_key_db(file):
    """
    Read the set of fixed keys fom the file
    """
    keys = set()

    f = open(file, "r")
    for line in f:
        u, k = line.rstrip().split(':', 2)
        keys.add((u, k))
    f.close()
    return keys

def write_key_db(file, keys):
    """
    Write the set of keys to the given file
    """

    f = open(file, 'w')
    for t in keys:
        print >>f, "%s:%s" % t
    f.close()

def read_project_db(file):
    """
    Read the set of fixed keys fom the file
    """
    projects = set()

    f = open(file, "r")
    for line in f:
        projects.add(line.rstrip())
    f.close()
    return projects

def write_project_db(file, projects):
    """
    Write the set of keys to the given file
    """

    f = open(file, 'w')
    for p in projects:
        print >>f, "%s" % p
    f.close()


# This is the command line utility
if __name__ =='__main__':
    import sys
    import MySQLdb
    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):
        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]
    keys = set()

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

    # Parse the options
    (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)
        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
        write_project_db(opts.file, projects)

    sys.exit(0)