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

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

7 8 9
class ReexecException(Exception):
    pass

10 11 12 13 14 15 16 17
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__
18

Guillaume Bury's avatar
Guillaume Bury committed
19
logging_levels = logging.WARNING, logging.INFO, logging.DEBUG, 5
20

21
class FileHandler(logging.FileHandler):
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
22

23 24 25 26 27 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
    _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)
55
    if log_level:
56
        root.setLevel(logging_levels[log_level-1])
57 58
    else:
        logging.disable(logging.CRITICAL)
Guillaume Bury's avatar
Guillaume Bury committed
59 60
    logging.addLevelName(5, 'TRACE')
    logging.trace = lambda *args, **kw: logging.log(5, *args, **kw)
61

62 63 64 65
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
66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84

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
85 86
class ArgParser(argparse.ArgumentParser):

Julien Muchembled's avatar
Julien Muchembled committed
87 88 89 90 91 92 93 94 95 96 97
    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
98
    def convert_arg_line_to_args(self, arg_line):
Julien Muchembled's avatar
Julien Muchembled committed
99
        if arg_line.split('#', 1)[0].rstrip():
Julien Muchembled's avatar
Julien Muchembled committed
100 101 102
            if arg_line.startswith('@'):
                yield arg_line
                return
103 104 105 106 107
            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
108

Julien Muchembled's avatar
Julien Muchembled committed
109 110 111 112 113 114 115
    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)


116 117 118 119 120 121 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
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()


155 156
class Popen(subprocess.Popen):

157
    def __init__(self, *args, **kw):
158
        self._args = tuple(args[0] if args else kw['args'])
159 160 161 162 163 164
        try:
            super(Popen, self).__init__(*args, **kw)
        except OSError, e:
            if e.errno != errno.ENOMEM:
                raise
            self.returncode = -1
165

166 167 168 169 170
    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)

171
    def stop(self):
172
        if self.pid and self.returncode is None:
173 174 175
            self.terminate()
            t = threading.Timer(5, self.kill)
            t.start()
176
            # PY3: use waitid(WNOWAIT) and call self.poll() after t.cancel()
177 178 179
            r = self.wait()
            t.cancel()
            return r
180 181


182
def select(R, W, T):
183
    try:
184
        r, w, _ = _select.select(R, W, (),
185 186 187 188 189 190 191
            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]()
192 193
    for w in w:
        W[w]()
194 195 196 197 198
    t = time.time()
    for next_refresh, refresh in T:
        if next_refresh <= t:
            refresh()

199
def makedirs(*args):
200
    try:
201
        os.makedirs(*args)
202 203 204 205
    except OSError, e:
        if e.errno != errno.EEXIST:
            raise

Guillaume Bury's avatar
Guillaume Bury committed
206
def binFromIp(ip):
207 208 209 210
    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
211
    return bin(ip1)[2:].rjust(64, '0') + bin(ip2)[2:].rjust(64, '0')
Guillaume Bury's avatar
Guillaume Bury committed
212

Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
213

214 215 216 217 218 219
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")
220 221
    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
222

223
def dump_address(address):
224
    return ';'.join(map(','.join, address))
225

226 227 228
def parse_address(address_list):
    for address in address_list.split(';'):
        try:
Julien Muchembled's avatar
Julien Muchembled committed
229 230 231
            a = ip, port, proto = address.split(',')
            int(port)
            yield a
232 233 234
        except ValueError, e:
            logging.warning("Failed to parse node address %r (%s)",
                            address, e)
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
235 236

def binFromSubnet(subnet):
237 238
    p, l = subnet.split('/')
    return bin(int(p))[2:].rjust(int(l), '0')
239 240 241 242 243 244 245

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()
246 247 248 249 250 251 252 253 254 255 256 257

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