Changeset 159a447 for fedd

Timestamp:

Nov 18, 2008 10:55:47 AM (16 years ago)

Author:

Ted Faber <faber@…>

Branches:

axis_example, compt_changes, info-ops, master, version-1.30, version-2.00, version-3.01, version-3.02

Children:

fd729b9

Parents:

7583a62

Message:

pre rename commit

File:

• fedd/fixed_key.py (modified) (6 diffs)

fedd/fixed_key.py

@@ -4 +4 @@
 # includes a command line utility to manipulate the DB.
 
-def read_db(file):
+def read_key_db(file):
     """
     Read the set of fixed keys fom the file
@@ -17 +17 @@
     return keys
 
-def write_db(file, keys):
+def write_key_db(file, keys):
     """
     Write the set of keys to the given file
@@ -27 +27 @@
     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__':
@@ -33 +56 @@
     from optparse import OptionParser
 
-    actions = set(("init", "add", "delete"))
+    actions = set(("initall", "init", "add", "delete"))
 
     def read_key_file(file):
@@ -46 +69 @@
     class key_opts(OptionParser):
         def __init__(self):
-            OptionParser.__init__(self, usage="%prog (init|add|delete) " + \
+            OptionParser.__init__(self, usage="%prog " + \
+                    "(init|add|delete|initall) " + \
                     "[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('-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)')
@@ -74 +100 @@
     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 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.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")
 
-    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':
+            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))
-            write_db(opts.file, keys)
-        elif action == 'delete':
+        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" % \
+                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:
-            sys.exit("How did we get here??")
+            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)