various cosmetic changes
[nodemanager.git] / tools.py
1 """A few things that didn't seem to fit anywhere else."""
2
3 import cPickle
4 import errno
5 import os
6 import pwd
7 import tempfile
8 import threading
9 import fcntl
10
11 import logger
12
13
14 PID_FILE = '/var/run/nm.pid'
15
16 def as_daemon_thread(run):
17     """Call function <run> with no arguments in its own thread."""
18     thr = threading.Thread(target=run)
19     thr.setDaemon(True)
20     thr.start()
21
22 def close_nonstandard_fds():
23     """Close all open file descriptors other than 0, 1, and 2."""
24     _SC_OPEN_MAX = 4
25     for fd in range(3, os.sysconf(_SC_OPEN_MAX)):
26         try: os.close(fd)
27         except OSError: pass  # most likely an fd that isn't open
28
29 # after http://www.erlenstar.demon.co.uk/unix/faq_2.html
30 def daemon():
31     """Daemonize the current process."""
32     if os.fork() != 0: os._exit(0)
33     os.setsid()
34     if os.fork() != 0: os._exit(0)
35     os.chdir('/')
36     os.umask(0)
37     devnull = os.open(os.devnull, os.O_RDWR)
38     os.dup2(devnull, 0)
39     # xxx fixme - this is just to make sure that nothing gets stupidly lost - should use devnull
40     crashlog = os.open('/var/log/nm.daemon', os.O_RDWR | os.O_APPEND | os.O_CREAT, 0644)
41     os.dup2(crashlog, 1)
42     os.dup2(crashlog, 2)
43
44 def fork_as(su, function, *args):
45     """fork(), cd / to avoid keeping unused directories open, close all nonstandard file descriptors (to avoid capturing open sockets), fork() again (to avoid zombies) and call <function> with arguments <args> in the grandchild process.  If <su> is not None, set our group and user ids appropriately in the child process."""
46     child_pid = os.fork()
47     if child_pid == 0:
48         try:
49             os.chdir('/')
50             close_nonstandard_fds()
51             if su:
52                 pw_ent = pwd.getpwnam(su)
53                 os.setegid(pw_ent[3])
54                 os.seteuid(pw_ent[2])
55             child_pid = os.fork()
56             if child_pid == 0: function(*args)
57         except:
58             os.seteuid(os.getuid())  # undo su so we can write the log file
59             os.setegid(os.getgid())
60             logger.log_exc()
61         os._exit(0)
62     else: os.waitpid(child_pid, 0)
63
64 def pid_file():
65     """We use a pid file to ensure that only one copy of NM is running at a given time.  If successful, this function will write a pid file containing the pid of the current process.  The return value is the pid of the other running process, or None otherwise."""
66     other_pid = None
67     if os.access(PID_FILE, os.F_OK):  # check for a pid file
68         handle = open(PID_FILE)  # pid file exists, read it
69         other_pid = int(handle.read())
70         handle.close()
71         # check for a process with that pid by sending signal 0
72         try: os.kill(other_pid, 0)
73         except OSError, e:
74             if e.errno == errno.ESRCH: other_pid = None  # doesn't exist
75             else: raise  # who knows
76     if other_pid == None:
77         # write a new pid file
78         write_file(PID_FILE, lambda f: f.write(str(os.getpid())))
79     return other_pid
80
81 def write_file(filename, do_write, **kw_args):
82     """Write file <filename> atomically by opening a temporary file, using <do_write> to write that file, and then renaming the temporary file."""
83     os.rename(write_temp_file(do_write, **kw_args), filename)
84
85 def write_temp_file(do_write, mode=None, uidgid=None):
86     fd, temporary_filename = tempfile.mkstemp()
87     if mode: os.chmod(temporary_filename, mode)
88     if uidgid: os.chown(temporary_filename, *uidgid)
89     f = os.fdopen(fd, 'w')
90     try: do_write(f)
91     finally: f.close()
92     return temporary_filename
93
94
95 class NMLock:
96     def __init__(self, file):
97         self.fd = os.open(file, os.O_RDWR|os.O_CREAT, 0600)
98         flags = fcntl.fcntl(self.fd, fcntl.F_GETFD)
99         flags |= fcntl.FD_CLOEXEC
100         fcntl.fcntl(self.fd, fcntl.F_SETFD, flags)
101     def __del__(self):
102         os.close(self.fd)
103     def acquire(self):
104         fcntl.lockf(self.fd, fcntl.LOCK_EX)
105     def release(self):
106         fcntl.lockf(self.fd, fcntl.LOCK_UN)