- forgot nm.py
[nodemanager.git] / accounts.py
1 """Functionality common to all account classes.
2
3 Each subclass of Account must provide five methods: create() and
4 destroy(), which are static; configure(), start(), and stop(), which
5 are not.  configure(), which takes a record as its only argument, does
6 things like set up ssh keys.  In addition, an Account subclass must
7 provide static member variables SHELL, which contains the unique shell
8 that it uses; and TYPE, a string that is used by the account creation
9 code.  For no particular reason, TYPE is divided hierarchically by
10 periods; at the moment the only convention is that all sliver accounts
11 have type that begins with sliver.
12
13 There are any number of race conditions that may result from the fact
14 that account names are not unique over time.  Moreover, it's a bad
15 idea to perform lengthy operations while holding the database lock.
16 In order to deal with both of these problems, we use a worker thread
17 for each account name that ever exists.  On 32-bit systems with large
18 numbers of accounts, this may cause the NM process to run out of
19 *virtual* memory!  This problem may be remedied by decreasing the
20 maximum stack size.
21 """
22
23 import Queue
24 import os
25 import pwd
26 import threading
27
28 import logger
29 import tools
30
31
32 # shell path -> account class association
33 shell_acct_class = {}
34 # account type -> account class association
35 type_acct_class = {}
36
37 def register_class(acct_class):
38     """Call once for each account class.  This method adds the class to the dictionaries used to look up account classes by shell and type."""
39     shell_acct_class[acct_class.SHELL] = acct_class
40     type_acct_class[acct_class.TYPE] = acct_class
41
42
43 # private account name -> worker object association and associated lock
44 name_worker_lock = threading.Lock()
45 name_worker = {}
46
47 def all():
48     """Return the names of all accounts on the system with recognized shells."""
49     return [pw_ent[0] for pw_ent in pwd.getpwall() if pw_ent[6] in shell_acct_class]
50
51 def get(name):
52     """Return the worker object for a particular username.  If no such object exists, create it first."""
53     name_worker_lock.acquire()
54     try:
55         if name not in name_worker: name_worker[name] = Worker(name)
56         return name_worker[name]
57     finally: name_worker_lock.release()
58
59
60 class Account:
61     def __init__(self, rec):
62         self.name = rec['name']
63         self.keys = ''
64         self.configure(rec)
65
66     @staticmethod
67     def create(name): abstract
68     @staticmethod
69     def destroy(name): abstract
70
71     def configure(self, rec):
72         """Write <rec['keys']> to my authorized_keys file."""
73         new_keys = rec['keys']
74         if new_keys != self.keys:
75             self.keys = new_keys
76             dot_ssh = '/home/%s/.ssh' % self.name
77             def do_installation():
78                 if not os.access(dot_ssh, os.F_OK): os.mkdir(dot_ssh)
79                 tools.write_file(dot_ssh + '/authorized_keys', lambda f: f.write(new_keys))
80             logger.log('%s: installing ssh keys' % self.name)
81             tools.fork_as(self.name, do_installation)
82
83     def start(self, delay=0): pass
84     def stop(self): pass
85
86
87 class Worker:
88     # these semaphores are acquired before creating/destroying an account
89     _create_sem = threading.Semaphore(1)
90     _destroy_sem = threading.Semaphore(1)
91
92     def __init__(self, name):
93         self.name = name  # username
94         self._acct = None  # the account object currently associated with this worker
95         # task list
96         # outsiders request operations by putting (fn, args...) tuples on _q
97         # the worker thread (created below) will perform these operations in order
98         self._q = Queue.Queue()
99         tools.as_daemon_thread(self._run)
100
101     def ensure_created(self, rec):
102         """Cause the account specified by <rec> to exist if it doesn't already."""
103         self._q.put((self._ensure_created, rec.copy()))
104
105     def _ensure_created(self, rec):
106         curr_class = self._get_class()
107         next_class = type_acct_class[rec['type']]
108         if next_class != curr_class:
109             self._destroy(curr_class)
110             self._create_sem.acquire()
111             try: next_class.create(self.name, rec['vref'])
112             finally: self._create_sem.release()
113         if not isinstance(self._acct, next_class): self._acct = next_class(rec)
114         else: self._acct.configure(rec)
115         if next_class != curr_class: self._acct.start()
116
117     def ensure_destroyed(self): self._q.put((self._ensure_destroyed,))
118     def _ensure_destroyed(self): self._destroy(self._get_class())
119
120     def start(self, delay=0): self._q.put((self._start, delay))
121     def _start(self, d): self._acct.start(delay=d)
122
123     def stop(self): self._q.put((self._stop,))
124     def _stop(self): self._acct.stop()
125
126     def _destroy(self, curr_class):
127         self._acct = None
128         if curr_class:
129             self._destroy_sem.acquire()
130             try: curr_class.destroy(self.name)
131             finally: self._destroy_sem.release()
132
133     def _get_class(self):
134         try: shell = pwd.getpwnam(self.name)[6]
135         except KeyError: return None
136         return shell_acct_class[shell]
137
138     def _run(self):
139         """Repeatedly pull commands off the queue and execute.  If memory usage becomes an issue, it might be wise to terminate after a while."""
140         while True:
141             try:
142                 cmd = self._q.get()
143                 cmd[0](*cmd[1:])
144             except: logger.log_exc()