accounts.py

   1 import Queue
   2 import os
   3 import pwd
   4 import threading
   5
   6 import logger
   7 import tools
   8
   9
  10 # shell path -> account class association
  11 shell_acct_class = {}
  12 # account type -> account class association
  13 type_acct_class = {}
  14
  15 def register_class(acct_class):
  16     """Call once for each account class.  This method adds the class to the dictionaries used to ook up account classes by shell and type."""
  17     shell_acct_class[acct_class.SHELL] = acct_class
  18     type_acct_class[acct_class.TYPE] = acct_class
  19
  20
  21 # private account name -> worker object association and associated lock
  22 _name_worker_lock = threading.Lock()
  23 _name_worker = {}
  24
  25 def all():
  26     """Returns a list of all NM accounts on the system.  Format is (type, username)."""
  27     pw_ents = pwd.getpwall()
  28     for pw_ent in pw_ents:
  29         if pw_ent[6] in shell_acct_class:
  30             yield shell_acct_class[pw_ent[6]].TYPE, pw_ent[0]
  31
  32 def get(name):
  33     """Return the worker object for a particular username.  If no such object exists, create it first."""
  34     _name_worker_lock.acquire()
  35     try:
  36         if name not in _name_worker: _name_worker[name] = Worker(name)
  37         return _name_worker[name]
  38     finally: _name_worker_lock.release()
  39
  40
  41 def install_ssh_keys(rec):
  42     """Write <rec['ssh_keys']> to <rec['name']>'s authorized_keys file."""
  43     dot_ssh = '/home/%s/.ssh' % rec['name']
  44     def do_installation():
  45         if not os.access(dot_ssh, os.F_OK): os.mkdir(dot_ssh)
  46         tools.write_file(dot_ssh + '/authorized_keys',
  47                          lambda thefile: thefile.write(rec['ssh_keys']))
  48     logger.log('%s: installing ssh keys' % rec['name'])
  49     tools.fork_as(rec['name'], do_installation)
  50
  51
  52 class Worker:
  53     # these semaphores are acquired before creating/destroying an account
  54     _create_sem = threading.Semaphore(1)
  55     _destroy_sem = threading.Semaphore(1)
  56
  57     def __init__(self, name):
  58         # username
  59         self.name = name
  60         # the account object currently associated with this worker
  61         self._acct = None
  62         # task list
  63         # outsiders request operations by putting (fn, args...) tuples on _q
  64         # the worker thread (created below) will perform these operations in order
  65         self._q = Queue.Queue()
  66         tools.as_daemon_thread(self._run)
  67
  68     def ensure_created(self, rec):
  69         """Caused the account specified by <rec> to exist if it doesn't already."""
  70         self._q.put((self._ensure_created, tools.deepcopy(rec)))
  71
  72     def _ensure_created(self, rec):
  73         curr_class = self._get_class()
  74         next_class = type_acct_class[rec['account_type']]
  75         if next_class != curr_class:
  76             self._destroy(curr_class)
  77             self._create_sem.acquire()
  78             try: next_class.create(self.name)
  79             finally: self._create_sem.release()
  80         self._make_acct_obj()
  81         self._acct.configure(rec)
  82         if next_class != curr_class: self._acct.start()
  83
  84     def ensure_destroyed(self): self._q.put((self._ensure_destroyed,))
  85     def _ensure_destroyed(self): self._destroy(self._get_class())
  86
  87     def start(self): self._q.put((self._start,))
  88     def _start(self):
  89         self._make_acct_obj()
  90         self._acct.start()
  91
  92     def stop(self): self._q.put((self._stop,))
  93     def _stop(self):
  94         self._make_acct_obj()
  95         self._acct.stop()
  96
  97     def _destroy(self, curr_class):
  98         self._acct = None
  99         if curr_class:
 100             self._destroy_sem.acquire()
 101             try: curr_class.destroy(self.name)
 102             finally: self._destroy_sem.release()
 103
 104     def _get_class(self):
 105         try: shell = pwd.getpwnam(self.name)[6]
 106         except KeyError: return None
 107         return shell_acct_class[shell]
 108
 109     def _make_acct_obj(self):
 110         curr_class = self._get_class()
 111         if not isinstance(self._acct, curr_class):
 112             self._acct = curr_class(self.name)
 113
 114     def _run(self):
 115         while True:
 116             try:
 117                 cmd = self._q.get()
 118                 cmd[0](*cmd[1:])
 119             except: logger.log_exc()