[afa43a8] | 1 | #!/usr/local/bin/python |
---|
| 2 | |
---|
| 3 | # Encode the set of fixed keys as a colon-separated user/key pairs. This also |
---|
| 4 | # includes a command line utility to manipulate the DB. |
---|
| 5 | |
---|
[159a447] | 6 | def read_key_db(file): |
---|
[afa43a8] | 7 | """ |
---|
| 8 | Read the set of fixed keys fom the file |
---|
| 9 | """ |
---|
| 10 | keys = set() |
---|
| 11 | |
---|
| 12 | f = open(file, "r") |
---|
| 13 | for line in f: |
---|
| 14 | u, k = line.rstrip().split(':', 2) |
---|
| 15 | keys.add((u, k)) |
---|
| 16 | f.close() |
---|
| 17 | return keys |
---|
| 18 | |
---|
[159a447] | 19 | def write_key_db(file, keys): |
---|
[afa43a8] | 20 | """ |
---|
| 21 | Write the set of keys to the given file |
---|
| 22 | """ |
---|
| 23 | |
---|
| 24 | f = open(file, 'w') |
---|
| 25 | for t in keys: |
---|
| 26 | print >>f, "%s:%s" % t |
---|
| 27 | f.close() |
---|
| 28 | |
---|
[159a447] | 29 | def read_project_db(file): |
---|
| 30 | """ |
---|
| 31 | Read the set of fixed keys fom the file |
---|
| 32 | """ |
---|
| 33 | projects = set() |
---|
| 34 | |
---|
| 35 | f = open(file, "r") |
---|
| 36 | for line in f: |
---|
| 37 | projects.add(line.rstrip()) |
---|
| 38 | f.close() |
---|
| 39 | return projects |
---|
| 40 | |
---|
| 41 | def write_project_db(file, projects): |
---|
| 42 | """ |
---|
| 43 | Write the set of keys to the given file |
---|
| 44 | """ |
---|
| 45 | |
---|
| 46 | f = open(file, 'w') |
---|
| 47 | for p in projects: |
---|
| 48 | print >>f, "%s" % p |
---|
| 49 | f.close() |
---|
| 50 | |
---|
| 51 | |
---|
[afa43a8] | 52 | # This is the command line utility |
---|
| 53 | if __name__ =='__main__': |
---|
| 54 | import sys |
---|
| 55 | import MySQLdb |
---|
| 56 | from optparse import OptionParser |
---|
| 57 | |
---|
[159a447] | 58 | actions = set(("initall", "init", "add", "delete")) |
---|
[afa43a8] | 59 | |
---|
| 60 | def read_key_file(file): |
---|
| 61 | """ |
---|
| 62 | Read a single line from a keyfile and trim the whitespace |
---|
| 63 | """ |
---|
| 64 | f = open(file, 'r') |
---|
| 65 | key = f.readline().rstrip() |
---|
| 66 | f.close() |
---|
| 67 | return key |
---|
| 68 | |
---|
| 69 | class key_opts(OptionParser): |
---|
| 70 | def __init__(self): |
---|
[159a447] | 71 | OptionParser.__init__(self, usage="%prog " + \ |
---|
| 72 | "(init|add|delete|initall) " + \ |
---|
[afa43a8] | 73 | "[opts] (--help for details)", version="0.1") |
---|
[159a447] | 74 | self.add_option('-t','--type', dest='type', type='choice', |
---|
| 75 | choices=['keys','projects'], help="database type") |
---|
[afa43a8] | 76 | self.add_option('-u','--user', dest='user', default=None, |
---|
| 77 | action='store', help='user to add/delete') |
---|
[159a447] | 78 | self.add_option('-p','--project', dest='project', default=None, |
---|
| 79 | action='store', help='project to add/delete') |
---|
[afa43a8] | 80 | self.add_option('-k','--key', dest='key', default=None, |
---|
| 81 | action='store', help='key to add/delete (string)') |
---|
| 82 | self.add_option('-K','--keyfile', dest='keyfile', default=None, |
---|
| 83 | action='store', help='key to add/delete (file)') |
---|
| 84 | self.add_option('-d', '--database', dest='file', default=None, |
---|
| 85 | action='store', help='database file') |
---|
| 86 | self.add_option('-f', '--file', dest='file', action='store', |
---|
| 87 | help='database file (synonym for --database)') |
---|
| 88 | |
---|
| 89 | |
---|
| 90 | parser = key_opts() |
---|
| 91 | action = sys.argv[1] |
---|
| 92 | keys = set() |
---|
| 93 | |
---|
| 94 | # Check the action |
---|
| 95 | if action in actions: del sys.argv[1] |
---|
| 96 | else: sys.exit("Bad action, must be one of %s" % ", ".join(actions)) |
---|
| 97 | |
---|
| 98 | # Parse the options |
---|
| 99 | (opts, args) = parser.parse_args() |
---|
| 100 | if not opts.file: |
---|
| 101 | sys.exit("Must specify DB file") |
---|
[159a447] | 102 | if not opts.type: |
---|
| 103 | sys.exit("Must specify database type (--type)") |
---|
| 104 | elif opts.type == 'keys': |
---|
| 105 | if action != "init" and action != "initall" and action != "addall": |
---|
| 106 | if opts.user and (opts.key or opts.keyfile): |
---|
| 107 | user = opts.user |
---|
| 108 | if opts.key: |
---|
| 109 | key = opts.key |
---|
| 110 | elif opts.keyfile: |
---|
| 111 | try: |
---|
| 112 | key = read_key_file(opts.keyfile) |
---|
| 113 | except IOError, e: |
---|
| 114 | sys.exit("Error reading keyfile: %s" % e) |
---|
| 115 | else: |
---|
| 116 | sys.exit("Must specify user and key") |
---|
| 117 | elif opts.type == 'projects': |
---|
| 118 | if action != "init" and action != "initall" and action != "addall": |
---|
| 119 | if opts.project: project = opts.project |
---|
| 120 | else: sys.exit("Must specify project") |
---|
[afa43a8] | 121 | else: |
---|
[159a447] | 122 | sys.exit("Invalid --type field (how'd you do that?)") |
---|
[afa43a8] | 123 | |
---|
[159a447] | 124 | if opts.type == 'keys': |
---|
| 125 | if action != 'init' and action != 'initall': |
---|
| 126 | keys = read_key_db(opts.file) |
---|
[afa43a8] | 127 | else: |
---|
[159a447] | 128 | keys = set() |
---|
[afa43a8] | 129 | |
---|
[159a447] | 130 | if action == 'initall': |
---|
| 131 | # Add all users from the Emulab DB |
---|
| 132 | try: |
---|
| 133 | db = MySQLdb.connect(db="tbdb") |
---|
| 134 | except: |
---|
| 135 | sys.exit("Cannot access the Emulab database") |
---|
| 136 | |
---|
| 137 | c = db.cursor() |
---|
| 138 | c.execute("SELECT uid, pubkey FROM user_pubkeys") |
---|
| 139 | for u, k in c.fetchall(): |
---|
| 140 | keys.add((u, k)) |
---|
| 141 | c.close() |
---|
| 142 | db.close() |
---|
| 143 | elif action == 'add': |
---|
[afa43a8] | 144 | keys.add((user, key)) |
---|
[159a447] | 145 | elif action == 'delete': |
---|
[afa43a8] | 146 | if (user, key) in keys: |
---|
| 147 | keys.remove((user, key)) |
---|
| 148 | else: |
---|
[159a447] | 149 | print >>sys.stderr, "Cannot delete (%s, %s): not in db" % \ |
---|
[afa43a8] | 150 | (user, key) |
---|
[159a447] | 151 | write_key_db(opts.file, keys) |
---|
| 152 | else: |
---|
| 153 | if action != 'init' and action != 'initall': |
---|
| 154 | projects = read_project_db(opts.file) |
---|
[afa43a8] | 155 | else: |
---|
[159a447] | 156 | projects = set() |
---|
| 157 | |
---|
| 158 | if action == 'initall': |
---|
| 159 | # Add all projects from the Emulab DB |
---|
| 160 | try: |
---|
| 161 | db = MySQLdb.connect(db="tbdb") |
---|
| 162 | except: |
---|
| 163 | sys.exit("Cannot access the Emulab database") |
---|
| 164 | |
---|
| 165 | c = db.cursor() |
---|
| 166 | c.execute("SELECT pid FROM projects") |
---|
| 167 | for p in c.fetchall(): |
---|
| 168 | projects.add(p) |
---|
| 169 | c.close() |
---|
| 170 | db.close() |
---|
| 171 | elif action == 'add': |
---|
| 172 | projects.add(project) |
---|
| 173 | elif action == 'delete': |
---|
| 174 | if project in projects: |
---|
| 175 | projects.remove(project) |
---|
| 176 | else: |
---|
| 177 | print >>sys.stderr, "Cannot delete %s: not in db" % project |
---|
| 178 | write_project_db(opts.file, projects) |
---|
| 179 | |
---|
[afa43a8] | 180 | sys.exit(0) |
---|