#!/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, read_user_db, write_user_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', 'users'], 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)')

def process_db(action, file, key, read_db, write_db, query):
    """
    The logic of insertion/deletion/initialization is the same regardless of
    the DB.  This function encapsulates that control flow.
    """

    if action != 'init' and action != 'initall':
	keys = read_db(file)
    else:
	keys = set()

    # init action falls through to write an empty DB
    if action == 'initall':
	# Add all keys returned by the query from the Emulab DB
	try:
	    db = MySQLdb.connect(db="tbdb")
	except:
	    sys.exit("Cannot access the Emulab database")

	c = db.cursor()
	c.execute(query)
	for k in c.fetchall():
	    keys.add(k)
	c.close()
	db.close()
    elif action == 'add':
	keys.add(key)
    elif action == 'delete': 
	if key in keys:
	    keys.remove(key)
	else:
	    print >>sys.stderr, "Cannot delete %s: not in db" % (key,)

    write_db(file, keys)


parser = key_opts()
action = sys.argv[1]
init_actions = set(['init', 'initall'])

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

(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 not in init_actions:
	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")
    else: 
	key = None
	user = None
    process_db(action, opts.file, (user, key), read_key_db, write_key_db,
	    "SELECT uid, pubkey FROM user_pubkeys")

elif opts.type == 'projects':
    if action not in init_actions:
	if opts.project: project = opts.project
	else: sys.exit("Must specify project")
    else: project = None
    process_db(action, opts.file, project, read_project_db, write_project_db, 
	    "SELECT pid FROM projects")
elif opts.type == 'users':
    if action not in init_actions:
	if opts.user: user = opts.user
	else: sys.exit("Must specify user")
    else: user = None
    process_db(action, opts.file, user, read_user_db, write_user_db, 
	    "SELECT uid FROM users")
else:
    sys.exit("Invalid --type field (how'd you do that?)")

sys.exit(0)