Always kill child processes gracefully
[re6stnet.git] / re6st / utils.py
1 import argparse, errno, logging, os, shlex, signal, socket
2 import struct, subprocess, textwrap, threading, time
3 from OpenSSL import crypto
4
5 logging_levels = logging.WARNING, logging.INFO, logging.DEBUG, 5
6
7 class FileHandler(logging.FileHandler):
8
9 _reopen = False
10
11 def release(self):
12 try:
13 if self._reopen:
14 self._reopen = False
15 self.close()
16 self._open()
17 finally:
18 self.lock.release()
19 # In the rare case _reopen is set just before the lock was released
20 if self._reopen and self.lock.acquire(0):
21 self.release()
22
23 def async_reopen(self, *_):
24 self._reopen = True
25 if self.lock.acquire(0):
26 self.release()
27
28 def setupLog(log_level, filename=None, **kw):
29 if log_level and filename:
30 makedirs(os.path.dirname(filename))
31 handler = FileHandler(filename)
32 sig = handler.async_reopen
33 else:
34 handler = logging.StreamHandler()
35 sig = signal.SIG_IGN
36 handler.setFormatter(logging.Formatter(
37 '%(asctime)s %(levelname)-9s %(message)s', '%d-%m-%Y %H:%M:%S'))
38 root = logging.getLogger()
39 root.addHandler(handler)
40 signal.signal(signal.SIGUSR1, sig)
41 if log_level:
42 root.setLevel(logging_levels[log_level-1])
43 else:
44 logging.disable(logging.CRITICAL)
45 logging.addLevelName(5, 'TRACE')
46 logging.trace = lambda *args, **kw: logging.log(5, *args, **kw)
47
48
49 class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter):
50
51 def _get_help_string(self, action):
52 return super(HelpFormatter, self)._get_help_string(action) \
53 if action.default else action.help
54
55 def _split_lines(self, text, width):
56 """Preserves new lines in option descriptions"""
57 lines = []
58 for text in text.splitlines():
59 lines += textwrap.wrap(text, width)
60 return lines
61
62 def _fill_text(self, text, width, indent):
63 """Preserves new lines in other descriptions"""
64 kw = dict(width=width, initial_indent=indent, subsequent_indent=indent)
65 return '\n'.join(textwrap.fill(t, **kw) for t in text.splitlines())
66
67 class ArgParser(argparse.ArgumentParser):
68
69 class _HelpFormatter(HelpFormatter):
70
71 def _format_actions_usage(self, actions, groups):
72 r = HelpFormatter._format_actions_usage(self, actions, groups)
73 if actions and actions[0].option_strings:
74 r = '[@OPTIONS_FILE] ' + r
75 return r
76
77 _ca_help = "Certificate authority (CA) file in .pem format." \
78 " Serial number defines the prefix of the network."
79
80 def convert_arg_line_to_args(self, arg_line):
81 arg_line = arg_line.split('#', 1)[0].rstrip()
82 if arg_line:
83 if arg_line.startswith('@'):
84 yield arg_line
85 return
86 arg_line = shlex.split(arg_line)
87 arg = '--' + arg_line.pop(0)
88 yield arg[arg not in self._option_string_actions:]
89 for arg in arg_line:
90 yield arg
91
92 def __init__(self, **kw):
93 super(ArgParser, self).__init__(formatter_class=self._HelpFormatter,
94 epilog="""Options can be read from a file. For example:
95 $ cat OPTIONS_FILE
96 ca /etc/re6stnet/ca.crt""", **kw)
97
98
99 class Popen(subprocess.Popen):
100
101 def stop(self):
102 self.terminate()
103 t = threading.Timer(5, self.kill)
104 t.start()
105 r = self.wait()
106 t.cancel()
107 return r
108
109
110 def makedirs(path):
111 try:
112 os.makedirs(path)
113 except OSError, e:
114 if e.errno != errno.EEXIST:
115 raise
116
117 def binFromIp(ip):
118 ip1, ip2 = struct.unpack('>QQ', socket.inet_pton(socket.AF_INET6, ip))
119 return bin(ip1)[2:].rjust(64, '0') + bin(ip2)[2:].rjust(64, '0')
120
121
122 def ipFromBin(ip, suffix=''):
123 suffix_len = 128 - len(ip)
124 if suffix_len > 0:
125 ip += suffix.rjust(suffix_len, '0')
126 elif suffix_len:
127 sys.exit("Prefix exceeds 128 bits")
128 return socket.inet_ntop(socket.AF_INET6,
129 struct.pack('>QQ', int(ip[:64], 2), int(ip[64:], 2)))
130
131 def networkFromCa(ca_path):
132 # Get network prefix from ca.crt
133 with open(ca_path, 'r') as f:
134 ca = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
135 return bin(ca.get_serial_number())[3:]
136
137 def subnetFromCert(cert_path):
138 # Get ip from cert.crt
139 with open(cert_path, 'r') as f:
140 cert = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
141 return cert.get_subject().CN
142
143 def address_str(address):
144 return ';'.join(map(','.join, address))
145
146
147 def address_list(address_list):
148 return list(tuple(address.split(','))
149 for address in address_list.split(';'))
150
151
152 def binFromSubnet(subnet):
153 p, l = subnet.split('/')
154 return bin(int(p))[2:].rjust(int(l), '0')