utils.py 7.95 KB
Newer Older
1
import argparse, errno, fcntl, hashlib, logging, os, select as _select
2 3
import shlex, signal, socket, sqlite3, struct, subprocess
import sys, textwrap, threading, time, traceback
4

5 6 7 8 9 10
# PY3: It will be even better to use Popen(pass_fds=...),
#      and then socket.SOCK_CLOEXEC will be useless.
#      (We already follow the good practice that consists in not
#      relying on the GC for the closing of file descriptors.)
socket.SOCK_CLOEXEC = 0x80000

11 12
HMAC_LEN = len(hashlib.sha1('').digest())

13 14 15
class ReexecException(Exception):
    pass

16 17 18 19 20 21 22 23
try:
    subprocess.CalledProcessError(0, '', '')
except TypeError: # BBB: Python < 2.7
    def __init__(self, returncode, cmd, output=None):
        self.returncode = returncode
        self.cmd = cmd
        self.output = output
    subprocess.CalledProcessError.__init__ = __init__
24

Guillaume Bury's avatar
Guillaume Bury committed
25
logging_levels = logging.WARNING, logging.INFO, logging.DEBUG, 5
26

27
class FileHandler(logging.FileHandler):
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
28

29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60
    _reopen = False

    def release(self):
        try:
            if self._reopen:
                self._reopen = False
                self.close()
                self._open()
        finally:
            self.lock.release()
        # In the rare case _reopen is set just before the lock was released
        if self._reopen and self.lock.acquire(0):
            self.release()

    def async_reopen(self, *_):
        self._reopen = True
        if self.lock.acquire(0):
            self.release()

def setupLog(log_level, filename=None, **kw):
    if log_level and filename:
        makedirs(os.path.dirname(filename))
        handler = FileHandler(filename)
        sig = handler.async_reopen
    else:
        handler = logging.StreamHandler()
        sig = signal.SIG_IGN
    handler.setFormatter(logging.Formatter(
        '%(asctime)s %(levelname)-9s %(message)s', '%d-%m-%Y %H:%M:%S'))
    root = logging.getLogger()
    root.addHandler(handler)
    signal.signal(signal.SIGUSR1, sig)
61
    if log_level:
62
        root.setLevel(logging_levels[log_level-1])
63 64
    else:
        logging.disable(logging.CRITICAL)
Guillaume Bury's avatar
Guillaume Bury committed
65 66
    logging.addLevelName(5, 'TRACE')
    logging.trace = lambda *args, **kw: logging.log(5, *args, **kw)
67

68 69 70 71
def log_exception():
    f = traceback.format_exception(*sys.exc_info())
    logging.error('%s%s', f.pop(), ''.join(f))

Julien Muchembled's avatar
Julien Muchembled committed
72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90

class HelpFormatter(argparse.ArgumentDefaultsHelpFormatter):

    def _get_help_string(self, action):
        return super(HelpFormatter, self)._get_help_string(action) \
            if action.default else action.help

    def _split_lines(self, text, width):
        """Preserves new lines in option descriptions"""
        lines = []
        for text in text.splitlines():
            lines += textwrap.wrap(text, width)
        return lines

    def _fill_text(self, text, width, indent):
        """Preserves new lines in other descriptions"""
        kw = dict(width=width, initial_indent=indent, subsequent_indent=indent)
        return '\n'.join(textwrap.fill(t, **kw) for t in text.splitlines())

Julien Muchembled's avatar
Julien Muchembled committed
91 92
class ArgParser(argparse.ArgumentParser):

Julien Muchembled's avatar
Julien Muchembled committed
93 94 95 96 97 98 99 100 101 102 103
    class _HelpFormatter(HelpFormatter):

        def _format_actions_usage(self, actions, groups):
            r = HelpFormatter._format_actions_usage(self, actions, groups)
            if actions and actions[0].option_strings:
                r = '[@OPTIONS_FILE] ' + r
            return r

    _ca_help = "Certificate authority (CA) file in .pem format." \
               " Serial number defines the prefix of the network."

Julien Muchembled's avatar
Julien Muchembled committed
104
    def convert_arg_line_to_args(self, arg_line):
Julien Muchembled's avatar
Julien Muchembled committed
105
        if arg_line.split('#', 1)[0].rstrip():
Julien Muchembled's avatar
Julien Muchembled committed
106 107 108
            if arg_line.startswith('@'):
                yield arg_line
                return
109 110 111 112 113
            arg_line = shlex.split(arg_line)
            arg = '--' + arg_line.pop(0)
            yield arg[arg not in self._option_string_actions:]
            for arg in arg_line:
                yield arg
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
114

Julien Muchembled's avatar
Julien Muchembled committed
115 116 117 118 119 120 121
    def __init__(self, **kw):
        super(ArgParser, self).__init__(formatter_class=self._HelpFormatter,
            epilog="""Options can be read from a file. For example:
  $ cat OPTIONS_FILE
  ca /etc/re6stnet/ca.crt""", **kw)


122 123 124 125 126 127 128 129 130 131 132 133 134 135 136 137 138 139 140 141 142 143 144 145 146 147 148 149 150 151 152 153 154 155 156 157 158 159 160
class exit(object):

    status = None

    def __init__(self):
        l = threading.Lock()
        self.acquire = l.acquire
        r = l.release
        def release():
            try:
                if self.status is not None:
                    self.release = r
                    sys.exit(self.status)
            finally:
                r()
        self.release = release

    def __enter__(self):
        self.acquire()

    def __exit__(self, t, v, tb):
        self.release()

    def kill_main(self, status):
        self.status = status
        os.kill(os.getpid(), signal.SIGTERM)

    def signal(self, status, *sigs):
        def handler(*args):
            if self.status is None:
                self.status = status
            if self.acquire(0):
                self.release()
        for sig in sigs:
            signal.signal(sig, handler)

exit = exit()


161 162
class Popen(subprocess.Popen):

163
    def __init__(self, *args, **kw):
164
        self._args = tuple(args[0] if args else kw['args'])
165 166 167 168 169 170
        try:
            super(Popen, self).__init__(*args, **kw)
        except OSError, e:
            if e.errno != errno.ENOMEM:
                raise
            self.returncode = -1
171

172 173 174 175 176
    def send_signal(self, sig):
        logging.info('Sending signal %s to pid %s %r',
                     sig, self.pid, self._args)
        super(Popen, self).send_signal(sig)

177
    def stop(self):
178
        if self.pid and self.returncode is None:
179 180 181
            self.terminate()
            t = threading.Timer(5, self.kill)
            t.start()
182
            # PY3: use waitid(WNOWAIT) and call self.poll() after t.cancel()
183 184 185
            r = self.wait()
            t.cancel()
            return r
186 187


188 189 190 191
def setCloexec(fd):
    flags = fcntl.fcntl(fd, fcntl.F_GETFD)
    fcntl.fcntl(fd, fcntl.F_SETFD, flags | fcntl.FD_CLOEXEC)

192
def select(R, W, T):
193
    try:
194
        r, w, _ = _select.select(R, W, (),
195 196 197 198 199 200 201
            max(0, min(T)[0] - time.time()) if T else None)
    except _select.error as e:
        if e.args[0] != errno.EINTR:
            raise
        return
    for r in r:
        R[r]()
202 203
    for w in w:
        W[w]()
204 205 206 207 208
    t = time.time()
    for next_refresh, refresh in T:
        if next_refresh <= t:
            refresh()

209
def makedirs(*args):
210
    try:
211
        os.makedirs(*args)
212 213 214 215
    except OSError, e:
        if e.errno != errno.EEXIST:
            raise

Guillaume Bury's avatar
Guillaume Bury committed
216
def binFromIp(ip):
217 218 219 220
    return binFromRawIp(socket.inet_pton(socket.AF_INET6, ip))

def binFromRawIp(ip):
    ip1, ip2 = struct.unpack('>QQ', ip)
Guillaume Bury's avatar
Guillaume Bury committed
221
    return bin(ip1)[2:].rjust(64, '0') + bin(ip2)[2:].rjust(64, '0')
Guillaume Bury's avatar
Guillaume Bury committed
222

Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
223

224 225 226 227 228 229
def ipFromBin(ip, suffix=''):
    suffix_len = 128 - len(ip)
    if suffix_len > 0:
        ip += suffix.rjust(suffix_len, '0')
    elif suffix_len:
        sys.exit("Prefix exceeds 128 bits")
230 231
    return socket.inet_ntop(socket.AF_INET6,
        struct.pack('>QQ', int(ip[:64], 2), int(ip[64:], 2)))
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
232

233
def dump_address(address):
234
    return ';'.join(map(','.join, address))
235

236 237 238
def parse_address(address_list):
    for address in address_list.split(';'):
        try:
Julien Muchembled's avatar
Julien Muchembled committed
239 240 241
            a = ip, port, proto = address.split(',')
            int(port)
            yield a
242 243 244
        except ValueError, e:
            logging.warning("Failed to parse node address %r (%s)",
                            address, e)
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
245 246

def binFromSubnet(subnet):
247 248
    p, l = subnet.split('/')
    return bin(int(p))[2:].rjust(int(l), '0')
249 250 251 252 253 254 255

def newHmacSecret():
    from random import getrandbits as g
    pack = struct.Struct(">QQI").pack
    assert len(pack(0,0,0)) == HMAC_LEN
    return lambda x=None: pack(g(64) if x is None else x, g(64), g(32))
newHmacSecret = newHmacSecret()
256 257 258 259 260 261 262 263 264 265 266 267

def sqliteCreateTable(db, name, *columns):
    sql = "CREATE TABLE %s (%s)" % (name, ','.join('\n  ' + x for x in columns))
    for x, in db.execute(
            "SELECT sql FROM sqlite_master WHERE type='table' and name=?""",
            (name,)):
        if x == sql:
            return
        raise sqlite3.OperationalError(
            "table %r already exists with unexpected schema" % name)
    db.execute(sql)
    return True