utils.py 8.33 KB
Newer Older
1 2
import argparse, calendar, errno, logging, os, select as _select, shlex, signal
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 160 161 162 163 164 165
        if self.pid:
            self.terminate()
            t = threading.Timer(5, self.kill)
            t.start()
            r = self.wait()
            t.cancel()
            return r
166 167


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

183 184 185 186 187 188 189
def makedirs(path):
    try:
        os.makedirs(path)
    except OSError, e:
        if e.errno != errno.EEXIST:
            raise

Guillaume Bury's avatar
Guillaume Bury committed
190 191
def binFromIp(ip):
    ip1, ip2 = struct.unpack('>QQ', socket.inet_pton(socket.AF_INET6, ip))
Guillaume Bury's avatar
Guillaume Bury committed
192
    return bin(ip1)[2:].rjust(64, '0') + bin(ip2)[2:].rjust(64, '0')
Guillaume Bury's avatar
Guillaume Bury committed
193

Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
194

195 196 197 198 199 200
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")
201 202
    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
203

204 205 206 207 208
def networkFromCa(ca):
    return bin(ca.get_serial_number())[3:]

def subnetFromCert(cert):
    return cert.get_subject().CN
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
209

210 211 212
def notAfter(cert):
    return calendar.timegm(time.strptime(cert.get_notAfter(),'%Y%m%d%H%M%SZ'))

213
def dump_address(address):
214
    return ';'.join(map(','.join, address))
215

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

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

230 231 232 233 234 235 236 237 238 239 240 241 242 243 244 245 246 247 248 249 250 251 252 253 254 255 256
if 1:
    def _iterRoutes():
        with open('/proc/net/ipv6_route') as f:
            routing_table = f.read()
        for line in routing_table.splitlines():
            line = line.split()
            iface = line[-1]
            if 0 < int(line[5], 16) < 1 << 31: # positive metric
                yield (iface, bin(int(line[0], 16))[2:].rjust(128, '0'),
                              int(line[1], 16))

_iterRoutes.__doc__ = """Iterates over all routes

    Amongst all returned routes starting with re6st prefix:
    - one is the local one with our prefix
    - any route with null prefix will be ignored
    - other are reachable routes installed by babeld
    """

def iterRoutes(network, exclude_prefix=None):
    a = len(network)
    for iface, ip, prefix_len in _iterRoutes():
        if ip[:a] == network:
            prefix = ip[a:prefix_len]
            if prefix and prefix != exclude_prefix:
                yield iface, prefix

257
def decrypt(key_path, data):
258 259
    p = Popen(('openssl', 'rsautl', '-decrypt', '-inkey', key_path),
              stdin=subprocess.PIPE, stdout=subprocess.PIPE)
260 261
    out, err = p.communicate(data)
    if p.returncode:
262
        raise subprocess.CalledProcessError(p.returncode, 'openssl', err)
263 264 265 266 267 268
    return out

def encrypt(cert, data):
    r, w = os.pipe()
    try:
        threading.Thread(target=os.write, args=(w, cert)).start()
269
        p = Popen(('openssl', 'rsautl', '-encrypt', '-certin',
270
                              '-inkey', '/proc/self/fd/%u' % r),
271
                  stdin=subprocess.PIPE, stdout=subprocess.PIPE)
272 273 274 275 276
        out, err = p.communicate(data)
    finally:
        os.close(r)
        os.close(w)
    if p.returncode:
277
        raise subprocess.CalledProcessError(p.returncode, 'openssl', err)
278
    return out