#!/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_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_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()

# This is the command line utility
if __name__ =='__main__':
    import sys
    import MySQLdb
    from optparse import OptionParser

    actions = set(("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) " + \
		    "[opts] (--help for details)", version="0.1")
	    self.add_option('-A', '--all-users', dest='all', default=False,
		    action='store_true', help='add all users to the database')
	    self.add_option('-u','--user', dest='user', default=None,
		    action='store', help='user 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 we're initializing, clear the DB unless we're adding all from the 
    # Emulab DB
    if action == 'init':
	if not opts.all:
	    write_db(opts.file, keys)
	    sys.exit(0)
    else:
	keys = read_db(opts.file)

    if opts.all:
	if action == 'delete':
	    sys.exit("delete and --all don't mix")
	# 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()
	write_db(opts.file, keys)
    else:
	# Add/delete a single user
	if opts.user: user = opts.user
	else: sys.exit("Must provide user (--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 provide key (--key or --keyfile)")
	
	if action == 'add':
	    keys.add((user, key))
	    write_db(opts.file, keys)
	elif action == 'delete':
	    if (user, key) in keys:
		keys.remove((user, key))
		write_db(opts.file, keys)
	    else:
		print >>sys.stderr, "Cannot delete (%s, %s) not in db" % \
			(user, key)
	else:
	    sys.exit("How did we get here??")
    sys.exit(0)