utils.py 6.55 KB
Newer Older
1
import argparse, errno, logging, os, select as _select, shlex, signal
2
import socket, struct, subprocess, sys, textwrap, threading, time, traceback
3 4 5 6 7 8 9 10
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__
11

Guillaume Bury's avatar
Guillaume Bury committed
12
logging_levels = logging.WARNING, logging.INFO, logging.DEBUG, 5
13

14
class FileHandler(logging.FileHandler):
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
15

16 17 18 19 20 21 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
    _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)
48
    if log_level:
49
        root.setLevel(logging_levels[log_level-1])
50 51
    else:
        logging.disable(logging.CRITICAL)
Guillaume Bury's avatar
Guillaume Bury committed
52 53
    logging.addLevelName(5, 'TRACE')
    logging.trace = lambda *args, **kw: logging.log(5, *args, **kw)
54

55 56 57 58
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
59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77

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

Julien Muchembled's avatar
Julien Muchembled committed
80 81 82 83 84 85 86 87 88 89 90
    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
91
    def convert_arg_line_to_args(self, arg_line):
Julien Muchembled's avatar
Julien Muchembled committed
92
        if arg_line.split('#', 1)[0].rstrip():
Julien Muchembled's avatar
Julien Muchembled committed
93 94 95
            if arg_line.startswith('@'):
                yield arg_line
                return
96 97 98 99 100
            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
101

Julien Muchembled's avatar
Julien Muchembled committed
102 103 104 105 106 107 108
    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)


109 110 111 112 113 114 115 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
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()


148 149
class Popen(subprocess.Popen):

150
    def __init__(self, *args, **kw):
151 152 153 154 155 156
        try:
            super(Popen, self).__init__(*args, **kw)
        except OSError, e:
            if e.errno != errno.ENOMEM:
                raise
            self.returncode = -1
157

158
    def stop(self):
159
        if self.pid and self.returncode is None:
160 161 162
            self.terminate()
            t = threading.Timer(5, self.kill)
            t.start()
163
            # PY3: use waitid(WNOWAIT) and call self.poll() after t.cancel()
164 165 166
            r = self.wait()
            t.cancel()
            return r
167 168


169
def select(R, W, T):
170
    try:
171
        r, w, _ = _select.select(R, W, (),
172 173 174 175 176 177 178
            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]()
179 180
    for w in w:
        W[w]()
181 182 183 184 185
    t = time.time()
    for next_refresh, refresh in T:
        if next_refresh <= t:
            refresh()

186 187 188 189 190 191 192
def makedirs(path):
    try:
        os.makedirs(path)
    except OSError, e:
        if e.errno != errno.EEXIST:
            raise

Guillaume Bury's avatar
Guillaume Bury committed
193
def binFromIp(ip):
194 195 196 197
    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
198
    return bin(ip1)[2:].rjust(64, '0') + bin(ip2)[2:].rjust(64, '0')
Guillaume Bury's avatar
Guillaume Bury committed
199

Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
200

201 202 203 204 205 206
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")
207 208
    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
209

210
def dump_address(address):
211
    return ';'.join(map(','.join, address))
212

213 214 215
def parse_address(address_list):
    for address in address_list.split(';'):
        try:
Julien Muchembled's avatar
Julien Muchembled committed
216 217 218
            a = ip, port, proto = address.split(',')
            int(port)
            yield a
219 220 221
        except ValueError, e:
            logging.warning("Failed to parse node address %r (%s)",
                            address, e)
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
222 223

def binFromSubnet(subnet):
224 225
    p, l = subnet.split('/')
    return bin(int(p))[2:].rjust(int(l), '0')