utils.py 4.87 KB
Newer Older
1 2
import argparse, errno, logging, os, shlex, signal, socket
import struct, subprocess, textwrap, threading, time
3
from OpenSSL import crypto
4

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

7
class FileHandler(logging.FileHandler):
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
8

9 10 11 12 13 14 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
    _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)
41
    if log_level:
42
        root.setLevel(logging_levels[log_level-1])
43 44
    else:
        logging.disable(logging.CRITICAL)
Guillaume Bury's avatar
Guillaume Bury committed
45 46
    logging.addLevelName(5, 'TRACE')
    logging.trace = lambda *args, **kw: logging.log(5, *args, **kw)
47

Julien Muchembled's avatar
Julien Muchembled committed
48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66

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

Julien Muchembled's avatar
Julien Muchembled committed
69 70 71 72 73 74 75 76 77 78 79
    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
80
    def convert_arg_line_to_args(self, arg_line):
Julien Muchembled's avatar
Julien Muchembled committed
81
        if arg_line.split('#', 1)[0].rstrip():
Julien Muchembled's avatar
Julien Muchembled committed
82 83 84
            if arg_line.startswith('@'):
                yield arg_line
                return
85 86 87 88 89
            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
90

Julien Muchembled's avatar
Julien Muchembled committed
91 92 93 94 95 96 97
    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)


98 99 100 101 102 103 104 105 106 107 108
class Popen(subprocess.Popen):

    def stop(self):
        self.terminate()
        t = threading.Timer(5, self.kill)
        t.start()
        r = self.wait()
        t.cancel()
        return r


109 110 111 112 113 114 115
def makedirs(path):
    try:
        os.makedirs(path)
    except OSError, e:
        if e.errno != errno.EEXIST:
            raise

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

Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
120

121 122 123 124 125 126
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")
127 128
    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
129

Guillaume Bury's avatar
Guillaume Bury committed
130
def networkFromCa(ca_path):
131
    # Get network prefix from ca.crt
Guillaume Bury's avatar
Guillaume Bury committed
132
    with open(ca_path, 'r') as f:
133
        ca = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
Guillaume Bury's avatar
Guillaume Bury committed
134
        return bin(ca.get_serial_number())[3:]
Guillaume Bury's avatar
Guillaume Bury committed
135

136
def subnetFromCert(cert_path):
137
    # Get ip from cert.crt
Guillaume Bury's avatar
Guillaume Bury committed
138
    with open(cert_path, 'r') as f:
139
        cert = crypto.load_certificate(crypto.FILETYPE_PEM, f.read())
140
        return cert.get_subject().CN
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
141

142
def dump_address(address):
143
    return ';'.join(map(','.join, address))
144

145 146 147 148 149 150 151 152
def parse_address(address_list):
    for address in address_list.split(';'):
        try:
            ip, port, proto = address.split(',')
            yield ip, str(port), proto
        except ValueError, e:
            logging.warning("Failed to parse node address %r (%s)",
                            address, e)
Ulysse Beaugnon's avatar
Ulysse Beaugnon committed
153 154

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