<commit_before>
from __future__ import print_function

import collections
import sys
from peeringdb.util import pretty_speed

class WhoisFormat(object):
    def __init__(self, fobj=sys.stdout):
        self.fobj = fobj

        self.display_names = {
            'fac_set': 'Facilities',
            }

    def mk_fmt(self, *widths):
        return '%-' + 's %-'.join(map(str, widths)) + 's'

    def mk_set_headers(self, data, columns):
        """ figure out sizes and create header fmt """
        columns = tuple(columns)
        lens = []

        for key in columns:
            value_len = max(len(str(each.get(key, ''))) for each in data)
            # account for header lengths
            lens.append(max(value_len, len(self._get_name(key))))

        fmt = self.mk_fmt(*lens)
        return fmt

    def _get_name(self, key):
        """ get display name for a key, or mangle for display """
        if key in self.display_names:
            return self.display_names[key]

        return key.capitalize()

    def _get_val(self, data, key):
        """ get value from a dict, format if necessary """
        return data.get(key, '')

    def _get_columns(self, data):
        """ get columns from a dict """
        return data.keys()

    def display_section(self, name):
        print(name, file=self.fobj)
        print('=' * len(name), file=self.fobj)
        print("", file=self.fobj)

    def display_headers(self, fmt, headers):
        print(fmt % headers, file=self.fobj)
        print(fmt % tuple('-' * len(x) for x in headers), file=self.fobj)

    def display_set(self, typ, data, columns):
        """ display a list of dicts """
        self.display_section("%s (%d)" % (self._get_name(typ), len(data)))
        headers = tuple(map(self._get_name, columns))
        fmt = self.mk_set_headers(data, columns)
        self.display_headers(fmt, headers)

        for each in data:
            row = tuple(self._get_val(each, k) for k,v in each.items())
            print(fmt % row, file=self.fobj)

        print("\n", file=self.fobj)

    def display_field(self, fmt, obj, field, display=None):
        if not display:
            display = self._get_name(field)
        print(fmt % (display, obj[field]), file=self.fobj)

    def check_set(self, data, name):
        if data.get(name, None):
            if hasattr(self, 'print_' + name):
                getattr(self, 'print_' + name)(data[name])

    def print_net(self, data):
        self.display_section("Network Information")
        fmt = "%-21s: %s"
        self.display_field(fmt, data, 'name', 'Name')
        self.display_field(fmt, data, 'asn', 'Primary ASN')
        self.display_field(fmt, data, 'aka', 'Also Known As')
        self.display_field(fmt, data, 'website', 'Website')
        self.display_field(fmt, data, 'irr_as_set', 'IRR AS-SET')
        self.display_field(fmt, data, 'info_type', 'Network Type')
        self.display_field(fmt, data, 'info_prefixes6', 'Approx IPv6 Prefixes')
        self.display_field(fmt, data, 'info_prefixes4', 'Approx IPv6 Prefixes')
        self.display_field(fmt, data, 'looking_glass', 'Looking Glass')
        self.display_field(fmt, data, 'route_server', 'Route Server')
        self.display_field(fmt, data, 'created', 'Created at')
        self.display_field(fmt, data, 'updated', 'Updated at')
        print("\n", file=self.fobj)

        self.display_section("Peering Policy Information")
        self.display_field(fmt, data, 'policy_url', 'URL')
        self.display_field(fmt, data, 'policy_general', 'General Policy')
        self.display_field(fmt, data, 'policy_locations', 'Location Requirement')
        self.display_field(fmt, data, 'policy_ratio', 'Ratio Requirement')
        self.display_field(fmt, data, 'policy_contracts', 'Contract Requirement')
        print("\n", file=self.fobj)

        self.check_set(data, 'poc_set')
        self.check_set(data, 'netixlan_set')
        self.check_set(data, 'netfac_set')

    def print_poc_set(self, data):
        self.display_section("Contact Information")
        fmt = self.mk_fmt(6, 20, 15, 20, 14)
        hdr = ('Role', 'Name', 'Email', 'URL', 'Phone')
        self.display_headers(fmt, hdr)
        for poc in data:
            print(fmt % (poc.get('role', ''), poc.get('name', ''), poc.get('email', ''), poc.get('url', ''), poc.get('phone', '')), file=self.fobj)

        for poc in data:
            print(fmt % (poc.get('role', ''), poc.get('name', ''), poc.get('email', ''), poc.get('url', ''), poc.get('phone', '')), file=self.fobj)

        print("\n", file=self.fobj)

    def print_netfac_set(self, data):
        self.display_section("Private Peering Facilities (%d)" % len(data))
        fmt = self.mk_fmt(51, 8, 15, 2)
        hdr = ('Facility Name', 'ASN', 'City', 'CO')
        self.display_headers(fmt, hdr)
        for each in data:
            print(fmt % (each.get('name', each.get('id')), each.get('local_asn', ''), each.get('city', ''), each.get('country', '')), file=self.fobj)
        print("\n", file=self.fobj)

    def print_netixlan_set(self, data):
        self.display_section("Public Peering Points (%d)" % len(data))
        fmt = self.mk_fmt(36, 8, 27, 5)
        hdr = ('Exchange Point', 'ASN', 'IP Address', 'Speed')
        self.display_headers(fmt, hdr)
        for ix in data:
            if ix.get('ipaddr4', None):
                print(fmt % (ix.get('name', ix.get('ixlan_id')), ix['asn'], ix['ipaddr4'], pretty_speed(ix['speed'])), file=self.fobj)
            if ix.get('ipaddr6', None):
                if ix.get('ipaddr4', None):
                    print(fmt % ('', '', ix['ipaddr6'], ''), file=self.fobj)
                else:
                    print(fmt % (ix['name'], ix['asn'], ix['ipaddr6'], pretty_speed(ix['speed'])), file=self.fobj)
        print("\n", file=self.fobj)

    def print(self, typ, data):
        if hasattr(self, 'print_' + typ):
            getattr(self, 'print_' + typ)(data)

        elif not data:
            print("%s: %s" % (typ, data), file=self.fobj)

        elif isinstance(data, collections.Mapping):
            print("\n", typ, file=self.fobj)
            for k,v in data.items():
                self.print(k, v)

        elif isinstance(data, (list, tuple)):
            # tabular data layout for lists of dicts
            if isinstance(data[0], collections.Mapping):
                self.display_set(typ, data, self._get_columns(data[0]))
            else:
                for each in data:
                    self.print(typ, each)

        else:
            print("%s: %s" % (typ, data), file=self.fobj)
<commit_msg>Fix field name for info_prefixe4 in whois format.<commit_after>
from __future__ import print_function

import collections
import sys
from peeringdb.util import pretty_speed

class WhoisFormat(object):
    def __init__(self, fobj=sys.stdout):
        self.fobj = fobj

        self.display_names = {
            'fac_set': 'Facilities',
            }

    def mk_fmt(self, *widths):
        return '%-' + 's %-'.join(map(str, widths)) + 's'

    def mk_set_headers(self, data, columns):
        """ figure out sizes and create header fmt """
        columns = tuple(columns)
        lens = []

        for key in columns:
            value_len = max(len(str(each.get(key, ''))) for each in data)
            # account for header lengths
            lens.append(max(value_len, len(self._get_name(key))))

        fmt = self.mk_fmt(*lens)
        return fmt

    def _get_name(self, key):
        """ get display name for a key, or mangle for display """
        if key in self.display_names:
            return self.display_names[key]

        return key.capitalize()

    def _get_val(self, data, key):
        """ get value from a dict, format if necessary """
        return data.get(key, '')

    def _get_columns(self, data):
        """ get columns from a dict """
        return data.keys()

    def display_section(self, name):
        print(name, file=self.fobj)
        print('=' * len(name), file=self.fobj)
        print("", file=self.fobj)

    def display_headers(self, fmt, headers):
        print(fmt % headers, file=self.fobj)
        print(fmt % tuple('-' * len(x) for x in headers), file=self.fobj)

    def display_set(self, typ, data, columns):
        """ display a list of dicts """
        self.display_section("%s (%d)" % (self._get_name(typ), len(data)))
        headers = tuple(map(self._get_name, columns))
        fmt = self.mk_set_headers(data, columns)
        self.display_headers(fmt, headers)

        for each in data:
            row = tuple(self._get_val(each, k) for k,v in each.items())
            print(fmt % row, file=self.fobj)

        print("\n", file=self.fobj)

    def display_field(self, fmt, obj, field, display=None):
        if not display:
            display = self._get_name(field)
        print(fmt % (display, obj[field]), file=self.fobj)

    def check_set(self, data, name):
        if data.get(name, None):
            if hasattr(self, 'print_' + name):
                getattr(self, 'print_' + name)(data[name])

    def print_net(self, data):
        self.display_section("Network Information")
        fmt = "%-21s: %s"
        self.display_field(fmt, data, 'name', 'Name')
        self.display_field(fmt, data, 'asn', 'Primary ASN')
        self.display_field(fmt, data, 'aka', 'Also Known As')
        self.display_field(fmt, data, 'website', 'Website')
        self.display_field(fmt, data, 'irr_as_set', 'IRR AS-SET')
        self.display_field(fmt, data, 'info_type', 'Network Type')
        self.display_field(fmt, data, 'info_prefixes6', 'Approx IPv6 Prefixes')
        self.display_field(fmt, data, 'info_prefixes4', 'Approx IPv4 Prefixes')
        self.display_field(fmt, data, 'looking_glass', 'Looking Glass')
        self.display_field(fmt, data, 'route_server', 'Route Server')
        self.display_field(fmt, data, 'created', 'Created at')
        self.display_field(fmt, data, 'updated', 'Updated at')
        print("\n", file=self.fobj)

        self.display_section("Peering Policy Information")
        self.display_field(fmt, data, 'policy_url', 'URL')
        self.display_field(fmt, data, 'policy_general', 'General Policy')
        self.display_field(fmt, data, 'policy_locations', 'Location Requirement')
        self.display_field(fmt, data, 'policy_ratio', 'Ratio Requirement')
        self.display_field(fmt, data, 'policy_contracts', 'Contract Requirement')
        print("\n", file=self.fobj)

        self.check_set(data, 'poc_set')
        self.check_set(data, 'netixlan_set')
        self.check_set(data, 'netfac_set')

    def print_poc_set(self, data):
        self.display_section("Contact Information")
        fmt = self.mk_fmt(6, 20, 15, 20, 14)
        hdr = ('Role', 'Name', 'Email', 'URL', 'Phone')
        self.display_headers(fmt, hdr)
        for poc in data:
            print(fmt % (poc.get('role', ''), poc.get('name', ''), poc.get('email', ''), poc.get('url', ''), poc.get('phone', '')), file=self.fobj)

        for poc in data:
            print(fmt % (poc.get('role', ''), poc.get('name', ''), poc.get('email', ''), poc.get('url', ''), poc.get('phone', '')), file=self.fobj)

        print("\n", file=self.fobj)

    def print_netfac_set(self, data):
        self.display_section("Private Peering Facilities (%d)" % len(data))
        fmt = self.mk_fmt(51, 8, 15, 2)
        hdr = ('Facility Name', 'ASN', 'City', 'CO')
        self.display_headers(fmt, hdr)
        for each in data:
            print(fmt % (each.get('name', each.get('id')), each.get('local_asn', ''), each.get('city', ''), each.get('country', '')), file=self.fobj)
        print("\n", file=self.fobj)

    def print_netixlan_set(self, data):
        self.display_section("Public Peering Points (%d)" % len(data))
        fmt = self.mk_fmt(36, 8, 27, 5)
        hdr = ('Exchange Point', 'ASN', 'IP Address', 'Speed')
        self.display_headers(fmt, hdr)
        for ix in data:
            if ix.get('ipaddr4', None):
                print(fmt % (ix.get('name', ix.get('ixlan_id')), ix['asn'], ix['ipaddr4'], pretty_speed(ix['speed'])), file=self.fobj)
            if ix.get('ipaddr6', None):
                if ix.get('ipaddr4', None):
                    print(fmt % ('', '', ix['ipaddr6'], ''), file=self.fobj)
                else:
                    print(fmt % (ix['name'], ix['asn'], ix['ipaddr6'], pretty_speed(ix['speed'])), file=self.fobj)
        print("\n", file=self.fobj)

    def print(self, typ, data):
        if hasattr(self, 'print_' + typ):
            getattr(self, 'print_' + typ)(data)

        elif not data:
            print("%s: %s" % (typ, data), file=self.fobj)

        elif isinstance(data, collections.Mapping):
            print("\n", typ, file=self.fobj)
            for k,v in data.items():
                self.print(k, v)

        elif isinstance(data, (list, tuple)):
            # tabular data layout for lists of dicts
            if isinstance(data[0], collections.Mapping):
                self.display_set(typ, data, self._get_columns(data[0]))
            else:
                for each in data:
                    self.print(typ, each)

        else:
            print("%s: %s" % (typ, data), file=self.fobj)
<|endoftext|>


<commit_before>import base64
import functools
import inspect
import json
import logging
import re
import subprocess
import sys
import time

from jinja2 import Template

def retry(times=3, sleep_time=3):
    def wrap(f):
        @functools.wraps(f)
        def inner(*args, **kwargs):
            for i in range(0, times):
                try:
                    return f(*args, **kwargs)
                except:
                    time.sleep(sleep_time)
            raise

        return inner

    return wrap

def in_bash(func):
    @functools.wraps(func)
    def wrap(*args, **kwargs):
        __BASH_DEBUG_INFO__ = []

        start_time = time.time()
        try:
            return func(*args, **kwargs)
        finally:
            end_time = time.time()
            logger.debug('BASH COMMAND DETAILS IN %s [cost %s ms]: %s' % (
            func.__name__, (end_time - start_time) * 1000, json.dumps(__BASH_DEBUG_INFO__)))
            del __BASH_DEBUG_INFO__

    return wrap

def __collect_locals_on_stack():
    frames = []
    frame = inspect.currentframe()
    while frame:
        frames.insert(0, frame)
        frame = frame.f_back

    ctx = {}
    for f in frames:
        ctx.update(f.f_locals)

    return ctx

def get_process(cmd, shell=None, workdir=None, pipe=None, executable=None):
    if pipe:
        return subprocess.Popen(cmd, shell=shell, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                close_fds=True, executable=executable, cwd=workdir)
    else:
        return subprocess.Popen(cmd, shell=shell, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                close_fds=True, executable=executable, cwd=workdir)

def bash_eval(raw_str, ctx=None):
    if not ctx:
        ctx = __collect_locals_on_stack()

    while True:
        unresolved = re.findall('{{(.+?)}}', raw_str)
        if not unresolved:
            break

        for u in unresolved:
            if u not in ctx:
                raise Exception('unresolved symbol {{%s}}' % u)

        tmpt = Template(raw_str)
        raw_str = tmpt.render(ctx)

    return raw_str

# @return: return code, stdout, stderr
def bash_roe(cmd, errorout=False, ret_code=0, pipe_fail=False):
    # type: (str, bool, int, bool) -> (int, str, str)
    ctx = __collect_locals_on_stack()

    cmd = bash_eval(cmd, ctx)
    p = get_process("/bin/bash", pipe=True)
    if pipe_fail:
        cmd = 'set -o pipefail; %s' % cmd
    o, e = p.communicate(cmd)
    r = p.returncode

    __BASH_DEBUG_INFO__ = ctx.get('__BASH_DEBUG_INFO__')
    if __BASH_DEBUG_INFO__ is not None:
        __BASH_DEBUG_INFO__.append({
            'cmd': cmd,
            'return_code': p.returncode,
            'stdout': o,
            'stderr': e
        })

    if r != ret_code and errorout:
        raise Exception('failed to execute bash[%s], return code: %s, stdout: %s, stderr: %s' % (cmd, r, o, e))
    if r == ret_code:
        e = None

    return r, o, e

# @return: return code, stdout
def bash_ro(cmd, pipe_fail=False):
    # type: (str, bool) -> (int, str)
    ret, o, _ = bash_roe(cmd, pipe_fail=pipe_fail)
    return ret, o

# @return: stdout
def bash_o(cmd, pipe_fail=False):
    # type: (str, bool) -> str
    _, o, _ = bash_roe(cmd, pipe_fail=pipe_fail)
    return o

# @return: return code
def bash_r(cmd, pipe_fail=False):
    ret, _, _ = bash_roe(cmd, pipe_fail=pipe_fail)
    return ret

def get_disk_holders(disk_names):
    holders = []
    for disk_name in disk_names:
        h = bash_o("ls /sys/class/block/%s/holders/" % disk_name).strip().splitlines()
        if len(h) == 0:
            continue
        holders.extend(h)
        holders.extend(get_disk_holders(h))
    holders.reverse()
    return holders

def drop_vg_lock(vgUuid):
    bash_roe("lvmlockctl --gl-disable %s" % vgUuid)
    bash_roe("lvmlockctl --drop %s" % vgUuid)

def remove_device_map_for_vg(vgUuid):
    o = bash_o("dmsetup ls | grep %s | awk '{print $1}'" % vgUuid).strip().splitlines()
    if len(o) == 0:
        return
    for dm in o:
        bash_roe("dmsetup remove %s" % dm.strip())

@in_bash
def wipe_fs(disks, expected_vg=None, with_lock=True):
    for disk in disks:
        exists_vg = None
        r = bash_r("pvdisplay %s | grep %s" % (disk, expected_vg))
        if r == 0:
            continue

        r, o = bash_ro("pvs --nolocking --noheading -o vg_name %s" % disk)
        if r == 0 and o.strip() != "":
            exists_vg = o.strip()

        need_flush_mpath = False

        bash_roe("partprobe -s %s" % disk)

        cmd_type = bash_o("lsblk %s -oTYPE | grep mpath" % disk)
        if cmd_type.strip() != "":
            need_flush_mpath = True

        bash_roe("wipefs -af %s" % disk)

        if need_flush_mpath:
            bash_roe("multipath -f %s && systemctl restart multipathd.service && sleep 1" % disk)

        for holder in get_disk_holders([disk.split("/")[-1]]):
            if not holder.startswith("dm-"):
                continue
            bash_roe("dmsetup remove /dev/%s" % holder)

        if exists_vg is not None:
            bash_r("grep %s /etc/drbd.d/* | awk '{print $1}' | sort | uniq | tr -d ':' | xargs rm" % exists_vg)
            logger.debug("found vg %s exists on this pv %s, start wipe" %
                         (exists_vg, disk))
            try:
                if with_lock:
                    drop_vg_lock(exists_vg)
                remove_device_map_for_vg(exists_vg)
            finally:
                pass

# from zstacklib.utils import bash
# from zstacklib.utils import linux
# from zstacklib.utils import lvm

logging.basicConfig(filename='/tmp/reset_mini_host.log', level=logging.DEBUG,
                    format='%(asctime)s %(process)d %(levelname)s %(funcName)s %(message)s')
logger = logging.getLogger(__name__)

@in_bash
@retry(times=3, sleep_time=1)
def stop_server():
    r, o, e = bash_roe("zsha2 stop-node || zstack-ctl stop")
    r1, o1 = bash_ro("pgrep -af -- '-DappName=zstack start'")
    if r1 == 0:
        raise Exception(
            "stop zstack failed, return code: %s, stdout: %s, stderr: %s, pgrep zstack return code: %s, stdout: %s" %
            (r, o, e, r1, o1))

@in_bash
def stop_kvmagent():
    bash_roe("/etc/init.d/zstack-kvmagent stop")
    bash_roe("pkill -9 -f 'c from kvmagent import kdaemon'")

@in_bash
def stop_vms():
    bash_roe("pkill -f -9 '/usr/libexec/qemu-kvm -name guest'")
    time.sleep(1)

@in_bash
def cleanup_storage():
    def get_live_drbd_minor():
        return bash_o(
            "cat /proc/drbd | grep -E '^[[:space:]]*[0-9]+\: ' | awk '{print $1}' | cut -d ':' -f1").strip().splitlines()

    def kill_drbd_holder(minor):
        if minor == "" or minor is None:
            return
        lines = bash_o("lsof /dev/drbd%s | grep -v '^COMMAND'" % minor).splitlines()
        if len(lines) == 0:
            return
        for line in lines:
            bash_r("kill -9 %s" % line.split()[1])

    def get_mini_pv():
        vg_name = bash_o("vgs --nolocking -oname,tags | grep zs::ministorage | awk '{print $1}'").strip()
        pv_names = bash_o("pvs --nolocking -oname -Svg_name=%s | grep -v PV" % vg_name).strip().splitlines()
        return [p.strip() for p in pv_names]

    bash_r("rm -rf /zstack_bs")
    bash_roe("drbdadm down all")
    if len(get_live_drbd_minor()) != 0:
        for m in get_live_drbd_minor():
            kill_drbd_holder(m)
        bash_roe("drbdadm down all")
    bash_r("/bin/rm /etc/drbd.d/*.res")
    wipe_fs(get_mini_pv())

@in_bash
def cleanup_zstack():
    bash_r("mysql -uroot -pzstack.mysql.password -e 'STOP SLAVE; RESET SLAVE ALL;'")
    bash_r("rm -rf /usr/local/zstack/apache-tomcat/logs/")
    bash_r("rm -rf /var/log/zstack/")
    bash_r("rm -rf /usr/local/zstack")
    bash_r("bash /opt/zstack-dvd/x86_64/c76/zstack-installer.bin -D -I 127.0.0.1")
    bash_r("bash /opt/zstack-dvd/x86_64/c76/zstack_mini_server.bin -a")
    bash_r("systemctl disable zstack.service")

@in_bash
def clear_network():
    bash_r("iptables -F; iptables -t nat -F; iptables -t mangle -F; iptables -t raw -F")
    bash_r("ebtables -F; ebtables -t nat -F")
    all_links = [x.strip().strip(":") for x in bash_o("ip -o link | awk '{print $2}'").strip().splitlines()]
    for i in all_links:
        if i not in ["lo", "eno1", "eno2", "ens2f0", "ens2f1"]:
            bash_r("ip link del %s" % i.split("@")[0])

@in_bash
def reset_network():
    bash_r(
        "ip link set dev eno1 nomaster; ip link set dev eno2 nomaster; ip link set dev ens2f0 nomaster; ip link set dev ens2f1 nomaster")
    bash_r("ip -4 a flush eno1; ip -4 a flush eno2; ip -4 a flush ens2f0; ip -4 a flush ens2f1")
    bash_r("ip r del default")
    bash_r("ls /etc/sysconfig/network-scripts/ifcfg-* | grep -E 'ifcfg-eno1|ifcfg-eno2|ifcfg-ens2f0|ifcfg-ens2f1' | xargs sed -i -E 's/(IPADDR.*|NETMASK.*|GATEWAY.*)//g'")
    bash_r("ls /etc/sysconfig/network-scripts/ifcfg-* | grep -Ev 'ifcfg-lo|ifcfg-eno1|ifcfg-eno2|ifcfg-ens2f0|ifcfg-ens2f1' | xargs /bin/rm")
    bash_r("systemctl restart network")
    clear_network()
    sn = bash_o("dmidecode -s system-serial-number").strip()
    if sn[-1] == "B":
        bash_r(
            "export PATH=$PATH:/usr/local/bin/; /usr/local/bin/zs-network-setting -i eno1 100.66.66.67 255.255.255.0")
        bash_r("ip link set dev eno1 up")
    elif sn[-1] == "A":
        bash_r(
            "export PATH=$PATH:/usr/local/bin/; /usr/local/bin/zs-network-setting -i eno1 100.66.66.66 255.255.255.0")
        bash_r("ip link set dev eno1 up")
    else:
        raise Exception("serial number not last with A or B!")
    # bash.bash_r("sed -i 's/ui_mode = .*/ui_mode = mini/g' /usr/local/zstack/apache-tomcat/webapps/zstack/WEB-INF/classes/zstack.properties")
    bash_r("systemctl restart zstack-network-agent.service")

@in_bash
def reset_system():
    reset_network()
    p = "enN0YWNrLm9yZ0A2MzdF"
    bash_r("echo 'root:%s' | chpasswd" % base64.decodestring(p))
    logger.info("reset system done")

def main(args):
    stop_server()
    stop_kvmagent()
    stop_vms()
    cleanup_storage()
    cleanup_zstack()
    reset_system()

if __name__ == "__main__":
    main(sys.argv[1:])
<commit_msg>Fix reset mini<commit_after>import base64
import functools
import inspect
import json
import logging
import re
import subprocess
import sys
import time

from jinja2 import Template

def retry(times=3, sleep_time=3):
    def wrap(f):
        @functools.wraps(f)
        def inner(*args, **kwargs):
            for i in range(0, times):
                try:
                    return f(*args, **kwargs)
                except:
                    time.sleep(sleep_time)
            raise

        return inner

    return wrap

def in_bash(func):
    @functools.wraps(func)
    def wrap(*args, **kwargs):
        __BASH_DEBUG_INFO__ = []

        start_time = time.time()
        try:
            return func(*args, **kwargs)
        finally:
            end_time = time.time()
            logger.debug('BASH COMMAND DETAILS IN %s [cost %s ms]: %s' % (
            func.__name__, (end_time - start_time) * 1000, json.dumps(__BASH_DEBUG_INFO__)))
            del __BASH_DEBUG_INFO__

    return wrap

def __collect_locals_on_stack():
    frames = []
    frame = inspect.currentframe()
    while frame:
        frames.insert(0, frame)
        frame = frame.f_back

    ctx = {}
    for f in frames:
        ctx.update(f.f_locals)

    return ctx

def get_process(cmd, shell=None, workdir=None, pipe=None, executable=None):
    if pipe:
        return subprocess.Popen(cmd, shell=shell, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                close_fds=True, executable=executable, cwd=workdir)
    else:
        return subprocess.Popen(cmd, shell=shell, stdout=subprocess.PIPE, stderr=subprocess.PIPE,
                                close_fds=True, executable=executable, cwd=workdir)

def bash_eval(raw_str, ctx=None):
    if not ctx:
        ctx = __collect_locals_on_stack()

    while True:
        unresolved = re.findall('{{(.+?)}}', raw_str)
        if not unresolved:
            break

        for u in unresolved:
            if u not in ctx:
                raise Exception('unresolved symbol {{%s}}' % u)

        tmpt = Template(raw_str)
        raw_str = tmpt.render(ctx)

    return raw_str

# @return: return code, stdout, stderr
def bash_roe(cmd, errorout=False, ret_code=0, pipe_fail=False):
    # type: (str, bool, int, bool) -> (int, str, str)
    ctx = __collect_locals_on_stack()

    cmd = bash_eval(cmd, ctx)
    p = get_process("/bin/bash", pipe=True)
    if pipe_fail:
        cmd = 'set -o pipefail; %s' % cmd
    o, e = p.communicate(cmd)
    r = p.returncode

    __BASH_DEBUG_INFO__ = ctx.get('__BASH_DEBUG_INFO__')
    if __BASH_DEBUG_INFO__ is not None:
        __BASH_DEBUG_INFO__.append({
            'cmd': cmd,
            'return_code': p.returncode,
            'stdout': o,
            'stderr': e
        })

    if r != ret_code and errorout:
        raise Exception('failed to execute bash[%s], return code: %s, stdout: %s, stderr: %s' % (cmd, r, o, e))
    if r == ret_code:
        e = None

    return r, o, e

# @return: return code, stdout
def bash_ro(cmd, pipe_fail=False):
    # type: (str, bool) -> (int, str)
    ret, o, _ = bash_roe(cmd, pipe_fail=pipe_fail)
    return ret, o

# @return: stdout
def bash_o(cmd, pipe_fail=False):
    # type: (str, bool) -> str
    _, o, _ = bash_roe(cmd, pipe_fail=pipe_fail)
    return o

# @return: return code
def bash_r(cmd, pipe_fail=False):
    ret, _, _ = bash_roe(cmd, pipe_fail=pipe_fail)
    return ret

def get_disk_holders(disk_names):
    holders = []
    for disk_name in disk_names:
        h = bash_o("ls /sys/class/block/%s/holders/" % disk_name).strip().splitlines()
        if len(h) == 0:
            continue
        holders.extend(h)
        holders.extend(get_disk_holders(h))
    holders.reverse()
    return holders

def drop_vg_lock(vgUuid):
    bash_roe("lvmlockctl --gl-disable %s" % vgUuid)
    bash_roe("lvmlockctl --drop %s" % vgUuid)

def remove_device_map_for_vg(vgUuid):
    o = bash_o("dmsetup ls | grep %s | awk '{print $1}'" % vgUuid).strip().splitlines()
    if len(o) == 0:
        return
    for dm in o:
        bash_roe("dmsetup remove %s" % dm.strip())

@in_bash
def wipe_fs(disks, expected_vg=None, with_lock=True):
    for disk in disks:
        exists_vg = None
        r = bash_r("pvdisplay %s | grep %s" % (disk, expected_vg))
        if r == 0:
            continue

        r, o = bash_ro("pvs --nolocking --noheading -o vg_name %s" % disk)
        if r == 0 and o.strip() != "":
            exists_vg = o.strip()

        need_flush_mpath = False

        bash_roe("partprobe -s %s" % disk)

        cmd_type = bash_o("lsblk %s -oTYPE | grep mpath" % disk)
        if cmd_type.strip() != "":
            need_flush_mpath = True

        bash_roe("wipefs -af %s" % disk)

        if need_flush_mpath:
            bash_roe("multipath -f %s && systemctl restart multipathd.service && sleep 1" % disk)

        for holder in get_disk_holders([disk.split("/")[-1]]):
            if not holder.startswith("dm-"):
                continue
            bash_roe("dmsetup remove /dev/%s" % holder)

        if exists_vg is not None:
            bash_r("grep %s /etc/drbd.d/* | awk '{print $1}' | sort | uniq | tr -d ':' | xargs rm" % exists_vg)
            logger.debug("found vg %s exists on this pv %s, start wipe" %
                         (exists_vg, disk))
            try:
                if with_lock:
                    drop_vg_lock(exists_vg)
                remove_device_map_for_vg(exists_vg)
            finally:
                pass

# from zstacklib.utils import bash
# from zstacklib.utils import linux
# from zstacklib.utils import lvm

logging.basicConfig(filename='/tmp/reset_mini_host.log', level=logging.DEBUG,
                    format='%(asctime)s %(process)d %(levelname)s %(funcName)s %(message)s')
logger = logging.getLogger(__name__)

@in_bash
@retry(times=3, sleep_time=1)
def stop_server():
    r, o, e = bash_roe("zsha2 stop-node || zstack-ctl stop")
    r1, o1 = bash_ro("pgrep -af -- '-DappName=zstack start'")
    if r1 == 0:
        raise Exception(
            "stop zstack failed, return code: %s, stdout: %s, stderr: %s, pgrep zstack return code: %s, stdout: %s" %
            (r, o, e, r1, o1))

@in_bash
def stop_kvmagent():
    bash_roe("/etc/init.d/zstack-kvmagent stop")
    bash_roe("pkill -9 -f 'c from kvmagent import kdaemon'")

@in_bash
def stop_vms():
    bash_roe("pkill -f -9 '/usr/libexec/qemu-kvm -name guest'")
    time.sleep(1)

@in_bash
def cleanup_storage():
    def get_live_drbd_minor():
        return bash_o(
            "cat /proc/drbd | grep -E '^[[:space:]]*[0-9]+\: ' | awk '{print $1}' | cut -d ':' -f1").strip().splitlines()

    def kill_drbd_holder(minor):
        if minor == "" or minor is None:
            return
        lines = bash_o("lsof /dev/drbd%s | grep -v '^COMMAND'" % minor).splitlines()
        if len(lines) == 0:
            return
        for line in lines:
            bash_r("kill -9 %s" % line.split()[1])

    def get_mini_pv():
        vg_name = bash_o("vgs --nolocking -oname,tags | grep zs::ministorage | awk '{print $1}'").strip()
        pv_names = bash_o("pvs --nolocking -oname -Svg_name=%s | grep -v PV" % vg_name).strip().splitlines()
        return [p.strip() for p in pv_names]

    bash_r("rm -rf /zstack_bs")
    bash_roe("drbdadm down all")
    if len(get_live_drbd_minor()) != 0:
        for m in get_live_drbd_minor():
            kill_drbd_holder(m)
        bash_roe("drbdadm down all")
    bash_r("/bin/rm /etc/drbd.d/*.res")
    wipe_fs(get_mini_pv())

@in_bash
def cleanup_zstack():
    bash_r("mysql -uroot -pzstack.mysql.password -e 'STOP SLAVE; RESET SLAVE ALL;'")
    bash_r("rm -rf /usr/local/zstack/apache-tomcat/logs/")
    bash_r("rm -rf /var/log/zstack/")
    bash_r("rm -rf /usr/local/zstack")
    bash_r("bash /opt/zstack-dvd/x86_64/c76/zstack-installer.bin -D -I 127.0.0.1")
    bash_r("bash /opt/zstack-dvd/x86_64/c76/zstack_mini_server.bin -a")
    bash_r("systemctl disable zstack.service")

@in_bash
def clear_network():
    bash_r("iptables -F; iptables -t nat -F; iptables -t mangle -F; iptables -t raw -F")
    bash_r("ebtables -F; ebtables -t nat -F")
    all_links = [x.strip().strip(":") for x in bash_o("ip -o link | awk '{print $2}'").strip().splitlines()]
    for i in all_links:
        if i not in ["lo", "eno1", "eno2", "ens2f0", "ens2f1"]:
            bash_r("ip link del %s" % i.split("@")[0])

@in_bash
def reset_network():
    bash_r(
        "ip link set dev eno1 nomaster; ip link set dev eno2 nomaster; ip link set dev ens2f0 nomaster; ip link set dev ens2f1 nomaster")
    bash_r("ip -4 a flush eno1; ip -4 a flush eno2; ip -4 a flush ens2f0; ip -4 a flush ens2f1")
    bash_r("ip r del default")
    bash_r("ls /etc/sysconfig/network-scripts/ifcfg-* | grep -E 'ifcfg-eno1|ifcfg-eno2|ifcfg-ens2f0|ifcfg-ens2f1' | xargs sed -i -E 's/(IPADDR.*|NETMASK.*|GATEWAY.*|SLAVE.*|MASTER.*)//g'")
    bash_r("ls /etc/sysconfig/network-scripts/ifcfg-* | grep -Ev 'ifcfg-lo|ifcfg-eno1|ifcfg-eno2|ifcfg-ens2f0|ifcfg-ens2f1' | xargs /bin/rm")
    bash_r("systemctl restart network")
    clear_network()
    sn = bash_o("dmidecode -s system-serial-number").strip()
    if sn[-1] == "B":
        bash_r(
            "export PATH=$PATH:/usr/local/bin/; /usr/local/bin/zs-network-setting -i eno1 100.66.66.67 255.255.255.0")
        bash_r("ip link set dev eno1 up")
    elif sn[-1] == "A":
        bash_r(
            "export PATH=$PATH:/usr/local/bin/; /usr/local/bin/zs-network-setting -i eno1 100.66.66.66 255.255.255.0")
        bash_r("ip link set dev eno1 up")
    else:
        raise Exception("serial number not last with A or B!")
    # bash.bash_r("sed -i 's/ui_mode = .*/ui_mode = mini/g' /usr/local/zstack/apache-tomcat/webapps/zstack/WEB-INF/classes/zstack.properties")
    bash_r("systemctl restart zstack-network-agent.service")

@in_bash
def reset_system():
    reset_network()
    p = "enN0YWNrLm9yZ0A2MzdF"
    bash_r("echo 'root:%s' | chpasswd" % base64.decodestring(p))
    logger.info("reset system done")

def main(args):
    stop_server()
    stop_kvmagent()
    stop_vms()
    cleanup_storage()
    cleanup_zstack()
    reset_system()

if __name__ == "__main__":
    main(sys.argv[1:])
<|endoftext|>


<commit_before>#
# Log-likelihood functions
#
# This file is part of PINTS.
#  Copyright (c) 2017-2018, University of Oxford.
#  For licensing information, see the LICENSE file distributed with the PINTS
#  software package.
#
from __future__ import absolute_import, division
from __future__ import print_function, unicode_literals
import pints
import numpy as np
import scipy.special

class KnownNoiseLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood assuming independent normally-distributed noise
    at each time point, using a known value for the standard deviation (sigma)
    of that noise.

    Arguments:

    ``problem``
        A :class:`SingleOutputProblem` or :class:`MultiOutputProblem`.
    ``sigma``
        The standard devation(s) of the noise. Can be a single value or a
        sequence of sigma's for each output.

    """
    def __init__(self, problem, sigma):
        super(KnownNoiseLogLikelihood, self).__init__(problem)

        # Check sigma
        no = problem.n_outputs()
        if np.isscalar(sigma):
            sigma = np.ones(no) * float(sigma)
        else:
            sigma = pints.vector(sigma)
            if len(sigma) != no:
                raise ValueError(
                    'Sigma must be a scalar or a vector of length n_outputs.')
        if np.any(sigma <= 0):
            raise ValueError('Standard deviation must be greater than zero.')

        # Pre-calculate parts
        self._offset = -0.5 * len(self._times) * np.log(2 * np.pi)
        self._offset -= len(self._times) * np.log(sigma)
        self._multip = -1 / (2.0 * sigma**2)

    def __call__(self, x):
        error = self._values - self._problem.evaluate(x)
        return np.sum(self._offset + self._multip * np.sum(error**2, axis=0))

class UnknownNoiseLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood assuming independent normally-distributed noise
    at each time point, and adds a parameter representing the standard
    deviation (sigma) of the noise on each output.

    For a noise level of ``sigma``, the likelihood becomes:

    .. math::
        L(\\theta, \sigma) = p(data | \\theta, \sigma) =
            \prod_{i=1}^N \\frac{1}{2\pi\sigma^2}\exp\left(
            -\\frac{(x_i - f_i(\\theta))^2}{2\sigma^2}\\right)

    leading to a log likelihood of:

    .. math::
        \log{L(\\theta, \sigma)} =
            -\\frac{N}{2}\log{2\pi}
            -N\log{\sigma}
            -\\frac{1}{2\sigma^2}\sum_{i=1}^N{(x_i - f_i(\\theta))^2}

    Arguments:

    ``problem``
        A :class:`SingleOutputProblem` or :class:`MultiOutputProblem`. For a
        single-output problem a single parameter is added, for a multi-output
        problem ``n_outputs`` parameters are added.

    """
    def __init__(self, problem):
        super(UnknownNoiseLogLikelihood, self).__init__(problem)

        # Get number of times, number of outputs
        self._nt = len(self._times)
        self._no = problem.n_outputs()

        # Add parameters to problem
        self._dimension = problem.n_parameters() + self._no

        # Pre-calculate parts
        self._logn = 0.5 * len(self._times) * np.log(2 * np.pi)

    def __call__(self, x):
        sigma = np.asarray(x[-self._no:])
        error = self._values - self._problem.evaluate(x[:-self._no])
        return np.sum(- self._logn - self._nt * np.log(sigma)
                      - np.sum(error**2, axis=0) / (2 * sigma**2))

class StudentTLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood assuming independent Student-t-distributed
    noise at each time point, and adds two parameters: one representing the
    degrees of freedom (``nu``), the other representing the scale (``sigma``).

    For a noise characterised by ``nu'' and ``sigma``, the log likelihood is of
    the form:

    .. math::
        \log{L(\\theta, \\nu, \sigma)} =
            N\\frac{\\nu}{2}\log(\\nu) - N\log(\sigma) -
            N\log B(\\nu/2, 1/2)
            -\\frac{1+\\nu}{2}\sum_{i=1}^N\log(\\nu +
            \\frac{x_i - f(\\theta)}{\sigma}^2)

    where ``B(.,.)`` is a beta function.

    Arguments:

    ``problem``
        A :class:`SingleOutputProblem` or :class:`MultiOutputProblem`. For a
        single-output problem two parameters are added ``(nu, sigma)``, where
        ``nu`` is the degrees of freedom and ``sigma`` is scale, for a
        multi-output problem ``2 * n_outputs`` parameters are added.

    """
    def __init__(self, problem):
        super(StudentTLogLikelihood, self).__init__(problem)

        # Get number of times, number of outputs
        self._nt = len(self._times)
        self._no = problem.n_outputs()

        # Add parameters to problem (two for each output)
        self._dimension = problem.n_parameters() + 2 * self._no

        # Pre-calculate
        self._n = len(self._times)

    def __call__(self, x):
        # For multiparameter problems the parameters are stored as
        # (model_params_1, model_params_2, ..., model_params_k,
        # nu_1, sigma_1, nu_2, sigma_2,...)
        n = self._n
        m = 2 * self._no

        # problem parameters
        problem_parameters = x[:-m]
        error = self._values - self._problem.evaluate(problem_parameters)

        # Distribution parameters
        parameters = x[-m:]
        nu = np.asarray(parameters[0::2])
        sigma = np.asarray(parameters[1::2])

        # Calculate
        return np.sum(
            + 0.5 * n * nu * np.log(nu)
            - n * np.log(sigma)
            - n * np.log(scipy.special.beta(0.5 * nu, 0.5))
            - 0.5 * (1 + nu) * np.sum(np.log(nu + (error / sigma)**2), axis=0)
        )

class ScaledLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood based on a (conditional)
    :class:`ProblemLogLikelihood` divided by the number of time samples.

    The returned value will be ``(1 / n) * log_likelihood(x|problem)``, where
    ``n`` is the number of time samples multiplied by the number of outputs.

    Arguments:

    ``log_likelihood``
        A :class:`ProblemLogLikelihood`.

    This log-likelihood operates on both single and multi-output problems.
    """
    def __init__(self, log_likelihood):
        # Check arguments
        if not isinstance(log_likelihood, pints.ProblemLogLikelihood):
            raise ValueError(
                'Given log_likelihood must extend pints.ProblemLogLikelihood')

        # Call parent constructor
        super(ScaledLogLikelihood, self).__init__(log_likelihood._problem)

        # Store log-likelihood
        self._log_likelihood = log_likelihood

        # Pre-calculate parts
        self._f = 1.0 / np.product(self._values.shape)

    def __call__(self, x):
        return self._log_likelihood(x) * self._f

class SumOfIndependentLogLikelihoods(pints.LogLikelihood):
    """
    *Extends:* :class:`LogLikelihood`

    Calculates a sum of :class:`LogLikelihood` objects, all defined on the same
    parameter space.

    This is useful for e.g. Bayesian inference using a single model evaluated
    on two **independent** data sets ``D`` and ``E``. In this case,

    .. math::
        f(\\theta|D,E) &= \\frac{f(D, E|\\theta)f(\\theta)}{f(D, E)} \\\\
                       &= \\frac{f(D|\\theta)f(E|\\theta)f(\\theta)}{f(D, E)}

    Arguments:

    ``log_likelihoods``
        A sequence of :class:`LogLikelihood` objects.

    Example::

        log_likelihood = pints.SumOfIndependentLogLikelihoods([
            pints.UnknownNoiseLogLikelihood(problem1),
            pints.UnknownNoiseLogLikelihood(problem2),
        ])

    """
    def __init__(self, log_likelihoods):
        super(SumOfIndependentLogLikelihoods, self).__init__()

        # Check input arguments
        if len(log_likelihoods) < 2:
            raise ValueError(
                'SumOfIndependentLogLikelihoods requires at least two log'
                ' likelihoods.')
        for i, e in enumerate(log_likelihoods):
            if not isinstance(e, pints.LogLikelihood):
                raise ValueError(
                    'All objects passed to SumOfIndependentLogLikelihoods must'
                    ' be instances of pints.LogLikelihood (failed on argument '
                    + str(i) + ').')
        self._log_likelihoods = list(log_likelihoods)

        # Get and check dimension
        i = iter(self._log_likelihoods)
        self._dimension = next(i).n_parameters()
        for e in i:
            if e.n_parameters() != self._dimension:
                raise ValueError(
                    'All log-likelihoods passed to'
                    ' SumOfIndependentLogLikelihoods must have same'
                    ' dimension.')

    def n_parameters(self):
        """ See :meth:`LogPDF.n_parameters()`. """
        return self._dimension

    def __call__(self, x):
        total = 0
        for e in self._log_likelihoods:
            total += e(x)
        return total
<commit_msg>Added equation to KnownNoiseLogLikelihood docs.<commit_after>#
# Log-likelihood functions
#
# This file is part of PINTS.
#  Copyright (c) 2017-2018, University of Oxford.
#  For licensing information, see the LICENSE file distributed with the PINTS
#  software package.
#
from __future__ import absolute_import, division
from __future__ import print_function, unicode_literals
import pints
import numpy as np
import scipy.special

class KnownNoiseLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood assuming independent normally-distributed noise
    at each time point, using a known value for the standard deviation (sigma)
    of that noise:

    .. math::
        \log{L(\\theta, \sigma)} =
            -\\frac{N}{2}\log{2\pi}
            -N\log{\sigma}
            -\\frac{1}{2\sigma^2}\sum_{i=1}^N{(x_i - f_i(\\theta))^2}

    Arguments:

    ``problem``
        A :class:`SingleOutputProblem` or :class:`MultiOutputProblem`.
    ``sigma``
        The standard devation(s) of the noise. Can be a single value or a
        sequence of sigma's for each output.

    """
    def __init__(self, problem, sigma):
        super(KnownNoiseLogLikelihood, self).__init__(problem)

        # Check sigma
        no = problem.n_outputs()
        if np.isscalar(sigma):
            sigma = np.ones(no) * float(sigma)
        else:
            sigma = pints.vector(sigma)
            if len(sigma) != no:
                raise ValueError(
                    'Sigma must be a scalar or a vector of length n_outputs.')
        if np.any(sigma <= 0):
            raise ValueError('Standard deviation must be greater than zero.')

        # Pre-calculate parts
        self._offset = -0.5 * len(self._times) * np.log(2 * np.pi)
        self._offset -= len(self._times) * np.log(sigma)
        self._multip = -1 / (2.0 * sigma**2)

    def __call__(self, x):
        error = self._values - self._problem.evaluate(x)
        return np.sum(self._offset + self._multip * np.sum(error**2, axis=0))

class UnknownNoiseLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood assuming independent normally-distributed noise
    at each time point, and adds a parameter representing the standard
    deviation (sigma) of the noise on each output.

    For a noise level of ``sigma``, the likelihood becomes:

    .. math::
        L(\\theta, \sigma) = p(data | \\theta, \sigma) =
            \prod_{i=1}^N \\frac{1}{2\pi\sigma^2}\exp\left(
            -\\frac{(x_i - f_i(\\theta))^2}{2\sigma^2}\\right)

    leading to a log likelihood of:

    .. math::
        \log{L(\\theta, \sigma)} =
            -\\frac{N}{2}\log{2\pi}
            -N\log{\sigma}
            -\\frac{1}{2\sigma^2}\sum_{i=1}^N{(x_i - f_i(\\theta))^2}

    Arguments:

    ``problem``
        A :class:`SingleOutputProblem` or :class:`MultiOutputProblem`. For a
        single-output problem a single parameter is added, for a multi-output
        problem ``n_outputs`` parameters are added.

    """
    def __init__(self, problem):
        super(UnknownNoiseLogLikelihood, self).__init__(problem)

        # Get number of times, number of outputs
        self._nt = len(self._times)
        self._no = problem.n_outputs()

        # Add parameters to problem
        self._dimension = problem.n_parameters() + self._no

        # Pre-calculate parts
        self._logn = 0.5 * len(self._times) * np.log(2 * np.pi)

    def __call__(self, x):
        sigma = np.asarray(x[-self._no:])
        error = self._values - self._problem.evaluate(x[:-self._no])
        return np.sum(- self._logn - self._nt * np.log(sigma)
                      - np.sum(error**2, axis=0) / (2 * sigma**2))

class StudentTLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood assuming independent Student-t-distributed
    noise at each time point, and adds two parameters: one representing the
    degrees of freedom (``nu``), the other representing the scale (``sigma``).

    For a noise characterised by ``nu'' and ``sigma``, the log likelihood is of
    the form:

    .. math::
        \log{L(\\theta, \\nu, \sigma)} =
            N\\frac{\\nu}{2}\log(\\nu) - N\log(\sigma) -
            N\log B(\\nu/2, 1/2)
            -\\frac{1+\\nu}{2}\sum_{i=1}^N\log(\\nu +
            \\frac{x_i - f(\\theta)}{\sigma}^2)

    where ``B(.,.)`` is a beta function.

    Arguments:

    ``problem``
        A :class:`SingleOutputProblem` or :class:`MultiOutputProblem`. For a
        single-output problem two parameters are added ``(nu, sigma)``, where
        ``nu`` is the degrees of freedom and ``sigma`` is scale, for a
        multi-output problem ``2 * n_outputs`` parameters are added.

    """
    def __init__(self, problem):
        super(StudentTLogLikelihood, self).__init__(problem)

        # Get number of times, number of outputs
        self._nt = len(self._times)
        self._no = problem.n_outputs()

        # Add parameters to problem (two for each output)
        self._dimension = problem.n_parameters() + 2 * self._no

        # Pre-calculate
        self._n = len(self._times)

    def __call__(self, x):
        # For multiparameter problems the parameters are stored as
        # (model_params_1, model_params_2, ..., model_params_k,
        # nu_1, sigma_1, nu_2, sigma_2,...)
        n = self._n
        m = 2 * self._no

        # problem parameters
        problem_parameters = x[:-m]
        error = self._values - self._problem.evaluate(problem_parameters)

        # Distribution parameters
        parameters = x[-m:]
        nu = np.asarray(parameters[0::2])
        sigma = np.asarray(parameters[1::2])

        # Calculate
        return np.sum(
            + 0.5 * n * nu * np.log(nu)
            - n * np.log(sigma)
            - n * np.log(scipy.special.beta(0.5 * nu, 0.5))
            - 0.5 * (1 + nu) * np.sum(np.log(nu + (error / sigma)**2), axis=0)
        )

class ScaledLogLikelihood(pints.ProblemLogLikelihood):
    """
    *Extends:* :class:`ProblemLogLikelihood`

    Calculates a log-likelihood based on a (conditional)
    :class:`ProblemLogLikelihood` divided by the number of time samples.

    The returned value will be ``(1 / n) * log_likelihood(x|problem)``, where
    ``n`` is the number of time samples multiplied by the number of outputs.

    Arguments:

    ``log_likelihood``
        A :class:`ProblemLogLikelihood`.

    This log-likelihood operates on both single and multi-output problems.
    """
    def __init__(self, log_likelihood):
        # Check arguments
        if not isinstance(log_likelihood, pints.ProblemLogLikelihood):
            raise ValueError(
                'Given log_likelihood must extend pints.ProblemLogLikelihood')

        # Call parent constructor
        super(ScaledLogLikelihood, self).__init__(log_likelihood._problem)

        # Store log-likelihood
        self._log_likelihood = log_likelihood

        # Pre-calculate parts
        self._f = 1.0 / np.product(self._values.shape)

    def __call__(self, x):
        return self._log_likelihood(x) * self._f

class SumOfIndependentLogLikelihoods(pints.LogLikelihood):
    """
    *Extends:* :class:`LogLikelihood`

    Calculates a sum of :class:`LogLikelihood` objects, all defined on the same
    parameter space.

    This is useful for e.g. Bayesian inference using a single model evaluated
    on two **independent** data sets ``D`` and ``E``. In this case,

    .. math::
        f(\\theta|D,E) &= \\frac{f(D, E|\\theta)f(\\theta)}{f(D, E)} \\\\
                       &= \\frac{f(D|\\theta)f(E|\\theta)f(\\theta)}{f(D, E)}

    Arguments:

    ``log_likelihoods``
        A sequence of :class:`LogLikelihood` objects.

    Example::

        log_likelihood = pints.SumOfIndependentLogLikelihoods([
            pints.UnknownNoiseLogLikelihood(problem1),
            pints.UnknownNoiseLogLikelihood(problem2),
        ])

    """
    def __init__(self, log_likelihoods):
        super(SumOfIndependentLogLikelihoods, self).__init__()

        # Check input arguments
        if len(log_likelihoods) < 2:
            raise ValueError(
                'SumOfIndependentLogLikelihoods requires at least two log'
                ' likelihoods.')
        for i, e in enumerate(log_likelihoods):
            if not isinstance(e, pints.LogLikelihood):
                raise ValueError(
                    'All objects passed to SumOfIndependentLogLikelihoods must'
                    ' be instances of pints.LogLikelihood (failed on argument '
                    + str(i) + ').')
        self._log_likelihoods = list(log_likelihoods)

        # Get and check dimension
        i = iter(self._log_likelihoods)
        self._dimension = next(i).n_parameters()
        for e in i:
            if e.n_parameters() != self._dimension:
                raise ValueError(
                    'All log-likelihoods passed to'
                    ' SumOfIndependentLogLikelihoods must have same'
                    ' dimension.')

    def n_parameters(self):
        """ See :meth:`LogPDF.n_parameters()`. """
        return self._dimension

    def __call__(self, x):
        total = 0
        for e in self._log_likelihoods:
            total += e(x)
        return total
<|endoftext|>


<commit_before>"""

"""

import os
import glob
import importlib
import mimetypes
from bson import ObjectId
from mako.lookup import TemplateLookup
from modularodm import fields

from framework.mongo import StoredObject
from framework.routing import process_rules
from framework.guid.model import GuidStoredObject

from website import settings

lookup = TemplateLookup(
    directories=[
        settings.TEMPLATES_PATH
    ]
)

def _is_image(filename):
    mtype, _ = mimetypes.guess_type(filename)
    return mtype and mtype.startswith('image')

class AddonConfig(object):

    def __init__(self, short_name, full_name, owners, categories,
                 added_default=None, added_mandatory=None,
                 node_settings_model=None, user_settings_model=None, include_js=None, include_css=None,
                 widget_help=None, views=None, configs=None, models=None,
                 has_hgrid_files=False, get_hgrid_data=None, max_file_size=None, high_max_file_size=None,
                 accept_extensions=True,
                 **kwargs):

        self.models = models
        self.settings_models = {}

        if node_settings_model:
            node_settings_model.config = self
            self.settings_models['node'] = node_settings_model

        if user_settings_model:
            user_settings_model.config = self
            self.settings_models['user'] = user_settings_model

        self.short_name = short_name
        self.full_name = full_name
        self.owners = owners
        self.categories = categories

        self.added_default = added_default or []
        self.added_mandatory = added_mandatory or []
        if set(self.added_mandatory).difference(self.added_default):
            raise ValueError('All mandatory targets must also be defaults.')

        self.include_js = self._include_to_static(include_js or {})
        self.include_css = self._include_to_static(include_css or {})

        self.widget_help = widget_help

        self.views = views or []
        self.configs = configs or []

        self.has_hgrid_files = has_hgrid_files
        # WARNING: get_hgrid_data can return None if the addon is added but has no credentials.
        self.get_hgrid_data = get_hgrid_data  # if has_hgrid_files and not get_hgrid_data rubeus.make_dummy()
        self.max_file_size = max_file_size
        self.high_max_file_size = high_max_file_size
        self.accept_extensions = accept_extensions

        # Build template lookup
        template_path = os.path.join('website', 'addons', short_name, 'templates')
        if os.path.exists(template_path):
            self.template_lookup = TemplateLookup(
                directories=[
                    template_path,
                    settings.TEMPLATES_PATH,
                ]
            )
        else:
            self.template_lookup = None

    def _static_url(self, filename):
        """Build static URL for file; use the current addon if relative path,
        else the global static directory.

        :param str filename: Local path to file
        :return str: Static URL for file

        """
        if filename.startswith('/'):
            return filename
        return '/static/addons/{addon}/{filename}'.format(
            addon=self.short_name,
            filename=filename,
        )

    def _include_to_static(self, include):
        """

        """
        # TODO: minify static assets
        return {
            key: [
                self._static_url(item)
                for item in value
            ]
            for key, value in include.iteritems()
        }

    # TODO: Make INCLUDE_JS and INCLUDE_CSS one option

    @property
    def icon(self):

        try:
            return self._icon
        except:
            static_path = os.path.join('website', 'addons', self.short_name, 'static')
            static_files = glob.glob(os.path.join(static_path, 'comicon.*'))
            image_files = [
                os.path.split(filename)[1]
                for filename in static_files
                if _is_image(filename)
            ]
            if len(image_files) == 1:
                self._icon = image_files[0]
            else:
                self._icon = None
            return self._icon

    @property
    def icon_url(self):
        return self._static_url(self.icon) if self.icon else None

    def to_json(self):
        return {
            'short_name': self.short_name,
            'full_name': self.full_name,
            'capabilities': self.short_name in settings.ADDON_CAPABILITIES,
            'addon_capabilities': settings.ADDON_CAPABILITIES.get(self.short_name),
            'icon': self.icon_url,
            'has_page': 'page' in self.views,
            'has_widget': 'widget' in self.views,
        }

    @property
    def path(self):
        return os.path.join(settings.BASE_PATH, self.short_name)

class GuidFile(GuidStoredObject):

    redirect_mode = 'proxy'

    _id = fields.StringField(primary=True)
    node = fields.ForeignField('node', required=True, index=True)

    _meta = {
        'abstract': True,
    }

    @property
    def file_url(self):
        raise NotImplementedError

    @property
    def deep_url(self):
        if self.node is None:
            raise ValueError('Node field must be defined.')
        return os.path.join(
            self.node.deep_url, self.file_url,
        )

class AddonSettingsBase(StoredObject):

    _id = fields.StringField(default=lambda: str(ObjectId()))
    deleted = fields.BooleanField(default=False)

    _meta = {
        'abstract': True,
    }

    def delete(self, save=True):
        self.deleted = True
        self.on_delete()
        if save:
            self.save()

    def undelete(self, save=True):
        self.deleted = False
        self.on_add()
        if save:
            self.save()

    def to_json(self, user):
        return {
            'addon_short_name': self.config.short_name,
            'addon_full_name': self.config.full_name,
        }

    #############
    # Callbacks #
    #############

    def on_add(self):
        """Called when the addon is added (or re-added) to the owner (User or Node)."""
        pass

    def on_delete(self):
        """Called when the addon is deleted from the owner (User or Node)."""
        pass

class AddonUserSettingsBase(AddonSettingsBase):

    owner = fields.ForeignField('user', backref='addons')

    _meta = {
        'abstract': True,
    }

    def __repr__(self):
        if self.owner:
            return '<{cls} owned by user {uid}>'.format(cls=self.__class__.__name__, uid=self.owner._id)
        else:
            return '<{cls} with no owner>'.format(cls=self.__class__.__name__)

    @property
    def public_id(self):
        return None

    @property
    def has_auth(self):
        """Whether the user has added credentials for this addon."""
        return False

    def get_backref_key(self, schema, backref_name):
        return schema._name + '__' + backref_name

    # TODO: Test me @asmacdo
    @property
    def nodes_authorized(self):
        """Get authorized, non-deleted nodes. Returns an empty list if the
        attached add-on does not include a node model.

        """
        try:
            schema = self.config.settings_models['node']
        except KeyError:
            return []
        nodes_backref = self.get_backref_key(schema, 'authorized')
        return [
            node_addon.owner
            for node_addon in getattr(self, nodes_backref)
            if not node_addon.owner.is_deleted
        ]

    def to_json(self, user):
        ret = super(AddonUserSettingsBase, self).to_json(user)
        ret['has_auth'] = self.has_auth
        ret.update({
            'nodes': [
                {
                    '_id': node._id,
                    'url': node.url,
                    'title': node.title,
                    'registered': node.is_registration,
                }
                for node in self.nodes_authorized
            ]
        })
        return ret

class AddonNodeSettingsBase(AddonSettingsBase):

    owner = fields.ForeignField('node', backref='addons')

    _meta = {
        'abstract': True,
    }

    def to_json(self, user):
        ret = super(AddonNodeSettingsBase, self).to_json(user)
        ret.update({
            'user': {
                'permissions': self.owner.get_permissions(user)
            },
            'node': {
                'id': self.owner._id,
                'api_url': self.owner.api_url,
                'url': self.owner.url,
                'is_registration': self.owner.is_registration,
            }
        })
        return ret

    def render_config_error(self, data):
        """

        """
        # Note: `config` is added to `self` in `AddonConfig::__init__`.
        template = lookup.get_template('project/addon/config_error.mako')
        return template.get_def('config_error').render(
            title=self.config.full_name,
            name=self.config.short_name,
            **data
        )

    #############
    # Callbacks #
    #############

    def before_page_load(self, node, user):
        """

        :param User user:
        :param Node node:

        """
        pass

    def before_remove_contributor(self, node, removed):
        """

        :param Node node:
        :param User removed:

        """
        pass

    def after_remove_contributor(self, node, removed):
        """

        :param Node node:
        :param User removed:

        """
        pass

    def before_make_public(self, node):

        """

        :param Node node:
        :returns: Alert message or None

        """
        pass

    def before_make_private(self, node):
        """

        :param Node node:
        :returns: Alert message or None

        """
        pass

    def after_set_privacy(self, node, permissions):
        """

        :param Node node:
        :param str permissions:

        """
        pass

    def before_fork(self, node, user):
        """

        :param Node node:
        :param User user:
        :returns: Alert message

        """
        pass

    def after_fork(self, node, fork, user, save=True):
        """

        :param Node node:
        :param Node fork:
        :param User user:
        :param bool save:
        :returns: Tuple of cloned settings and alert message

        """
        clone = self.clone()
        clone.owner = fork

        if save:
            clone.save()

        return clone, None

    def before_register(self, node, user):
        """

        :param Node node:
        :param User user:
        :returns: Alert message

        """
        pass

    def after_register(self, node, registration, user, save=True):
        """

        :param Node node:
        :param Node registration:
        :param User user:
        :param bool save:
        :returns: Tuple of cloned settings and alert message

        """
        clone = self.clone()
        clone.owner = registration

        if save:
            clone.save()

        return clone, None

    def after_delete(self, node, user):
        """

        :param Node node:
        :param User user:

        """
        pass

# TODO: No more magicks
def init_addon(app, addon_name, routes=True):
    """Load addon module return its create configuration object.

    If `log_fp` is provided, the addon's log templates will be appended
    to the file.

    :param app: Flask app object
    :param addon_name: Name of addon directory
    :param file log_fp: File pointer for the built logs file.
    :param bool routes: Add routes
    :return AddonConfig: AddonConfig configuration object if module found,
        else None

    """
    import_path = 'website.addons.{0}'.format(addon_name)

    # Import addon module
    addon_module = importlib.import_module(import_path)

    data = vars(addon_module)

    # Add routes
    if routes:
        for route_group in getattr(addon_module, 'ROUTES', []):
            process_rules(app, **route_group)

    # Build AddonConfig object
    return AddonConfig(
        **{
            key.lower(): value
            for key, value in data.iteritems()
        }
    )
<commit_msg>Rewrite the base fileguid<commit_after>"""

"""

import os
import glob
import importlib
import mimetypes
from bson import ObjectId
from flask import request
from modularodm import fields
from mako.lookup import TemplateLookup

import furl
import requests

from framework.mongo import StoredObject
from framework.routing import process_rules
from framework.guid.model import GuidStoredObject

from website import settings
from website.addons.base import exceptions

lookup = TemplateLookup(
    directories=[
        settings.TEMPLATES_PATH
    ]
)

def _is_image(filename):
    mtype, _ = mimetypes.guess_type(filename)
    return mtype and mtype.startswith('image')

class AddonConfig(object):

    def __init__(self, short_name, full_name, owners, categories,
                 added_default=None, added_mandatory=None,
                 node_settings_model=None, user_settings_model=None, include_js=None, include_css=None,
                 widget_help=None, views=None, configs=None, models=None,
                 has_hgrid_files=False, get_hgrid_data=None, max_file_size=None, high_max_file_size=None,
                 accept_extensions=True,
                 **kwargs):

        self.models = models
        self.settings_models = {}

        if node_settings_model:
            node_settings_model.config = self
            self.settings_models['node'] = node_settings_model

        if user_settings_model:
            user_settings_model.config = self
            self.settings_models['user'] = user_settings_model

        self.short_name = short_name
        self.full_name = full_name
        self.owners = owners
        self.categories = categories

        self.added_default = added_default or []
        self.added_mandatory = added_mandatory or []
        if set(self.added_mandatory).difference(self.added_default):
            raise ValueError('All mandatory targets must also be defaults.')

        self.include_js = self._include_to_static(include_js or {})
        self.include_css = self._include_to_static(include_css or {})

        self.widget_help = widget_help

        self.views = views or []
        self.configs = configs or []

        self.has_hgrid_files = has_hgrid_files
        # WARNING: get_hgrid_data can return None if the addon is added but has no credentials.
        self.get_hgrid_data = get_hgrid_data  # if has_hgrid_files and not get_hgrid_data rubeus.make_dummy()
        self.max_file_size = max_file_size
        self.high_max_file_size = high_max_file_size
        self.accept_extensions = accept_extensions

        # Build template lookup
        template_path = os.path.join('website', 'addons', short_name, 'templates')
        if os.path.exists(template_path):
            self.template_lookup = TemplateLookup(
                directories=[
                    template_path,
                    settings.TEMPLATES_PATH,
                ]
            )
        else:
            self.template_lookup = None

    def _static_url(self, filename):
        """Build static URL for file; use the current addon if relative path,
        else the global static directory.

        :param str filename: Local path to file
        :return str: Static URL for file

        """
        if filename.startswith('/'):
            return filename
        return '/static/addons/{addon}/{filename}'.format(
            addon=self.short_name,
            filename=filename,
        )

    def _include_to_static(self, include):
        """

        """
        # TODO: minify static assets
        return {
            key: [
                self._static_url(item)
                for item in value
            ]
            for key, value in include.iteritems()
        }

    # TODO: Make INCLUDE_JS and INCLUDE_CSS one option

    @property
    def icon(self):

        try:
            return self._icon
        except:
            static_path = os.path.join('website', 'addons', self.short_name, 'static')
            static_files = glob.glob(os.path.join(static_path, 'comicon.*'))
            image_files = [
                os.path.split(filename)[1]
                for filename in static_files
                if _is_image(filename)
            ]
            if len(image_files) == 1:
                self._icon = image_files[0]
            else:
                self._icon = None
            return self._icon

    @property
    def icon_url(self):
        return self._static_url(self.icon) if self.icon else None

    def to_json(self):
        return {
            'short_name': self.short_name,
            'full_name': self.full_name,
            'capabilities': self.short_name in settings.ADDON_CAPABILITIES,
            'addon_capabilities': settings.ADDON_CAPABILITIES.get(self.short_name),
            'icon': self.icon_url,
            'has_page': 'page' in self.views,
            'has_widget': 'widget' in self.views,
        }

    @property
    def path(self):
        return os.path.join(settings.BASE_PATH, self.short_name)

class GuidFile(GuidStoredObject):

    redirect_mode = 'proxy'

    _id = fields.StringField(primary=True)

    name = fields.StringField()
    enriched = fields.BooleanField(default=False)
    node = fields.ForeignField('node', required=True, index=True)

    _meta = {
        'abstract': True,
    }

    @property
    def provider(self):
        raise NotImplementedError

    @property
    def waterbutler_extras(self):
        return {}

    @property
    def guid_url(self):
        return '/{0}/'.format(self._id)

    @property
    def file_name(self):
        return '{0}.html'.format(self._id)

    @property
    def joinable_path(self):
        if self.path.startswith('/'):
            return self.path[1:]
        return self.path

    @property
    def _base_butler_url(self):
        url = furl.furl(settings.WATERBUTLER_URL)
        url.args.update(self.waterbutler_extras)

        url.args.update({
            'nid': self.node._id,
            'path': self.path,
            'provider': self.provider,
            'cookie': request.cookies.get(settings.COOKIE_NAME)
        })

        return url

    @property
    def download_url(self):
        url = self._base_butler_url
        url.path.add('file')
        return url.url

    @property
    def metadata_url(self):
        url = self._base_butler_url
        url.path.add('data')

        return url.url

    @property
    def cache_path(self):
        return os.path.join(
            settings.MFR_CACHE_PATH,
            self.provider,
            self.file_name
        )

    @property
    def temp_path(self):
        return os.path.join(
            settings.MFR_TEMP_PATH,
            self.provider,
            self.file_name + os.path.splitext(self.path)[1]
        )

    @property
    def deep_url(self):
        if self.node is None:
            raise ValueError('Node field must be defined.')

        url = os.path.join(
            self.node.deep_url,
            'files',
            self.provider,
            self.joinable_path
        )

        if url.endswith('/'):
            return url
        else:
            return url + '/'

    def _fetch_metadata(self, should_raise=False):
        resp = requests.get(self.metadata_url)

        if should_raise:
            if resp.status_code != 200:
                raise exceptions.AddonEnrichmentError(resp.status_code)

        return resp

    def enrich(self, save=True):
        resp = self._fetch_metadata(should_raise=True)

        if resp.json().get('code', 200) != 200:
            raise exceptions.AddonEnrichmentError(resp.json()['code'])

        metadata = resp.json()['data']
        self.name = metadata['name']
        self.enriched = True
        self.save()

class AddonSettingsBase(StoredObject):

    _id = fields.StringField(default=lambda: str(ObjectId()))
    deleted = fields.BooleanField(default=False)

    _meta = {
        'abstract': True,
    }

    def delete(self, save=True):
        self.deleted = True
        self.on_delete()
        if save:
            self.save()

    def undelete(self, save=True):
        self.deleted = False
        self.on_add()
        if save:
            self.save()

    def to_json(self, user):
        return {
            'addon_short_name': self.config.short_name,
            'addon_full_name': self.config.full_name,
        }

    #############
    # Callbacks #
    #############

    def on_add(self):
        """Called when the addon is added (or re-added) to the owner (User or Node)."""
        pass

    def on_delete(self):
        """Called when the addon is deleted from the owner (User or Node)."""
        pass

class AddonUserSettingsBase(AddonSettingsBase):

    owner = fields.ForeignField('user', backref='addons')

    _meta = {
        'abstract': True,
    }

    def __repr__(self):
        if self.owner:
            return '<{cls} owned by user {uid}>'.format(cls=self.__class__.__name__, uid=self.owner._id)
        else:
            return '<{cls} with no owner>'.format(cls=self.__class__.__name__)

    @property
    def public_id(self):
        return None

    @property
    def has_auth(self):
        """Whether the user has added credentials for this addon."""
        return False

    def get_backref_key(self, schema, backref_name):
        return schema._name + '__' + backref_name

    # TODO: Test me @asmacdo
    @property
    def nodes_authorized(self):
        """Get authorized, non-deleted nodes. Returns an empty list if the
        attached add-on does not include a node model.

        """
        try:
            schema = self.config.settings_models['node']
        except KeyError:
            return []
        nodes_backref = self.get_backref_key(schema, 'authorized')
        return [
            node_addon.owner
            for node_addon in getattr(self, nodes_backref)
            if not node_addon.owner.is_deleted
        ]

    def to_json(self, user):
        ret = super(AddonUserSettingsBase, self).to_json(user)
        ret['has_auth'] = self.has_auth
        ret.update({
            'nodes': [
                {
                    '_id': node._id,
                    'url': node.url,
                    'title': node.title,
                    'registered': node.is_registration,
                }
                for node in self.nodes_authorized
            ]
        })
        return ret

class AddonNodeSettingsBase(AddonSettingsBase):

    owner = fields.ForeignField('node', backref='addons')

    _meta = {
        'abstract': True,
    }

    def to_json(self, user):
        ret = super(AddonNodeSettingsBase, self).to_json(user)
        ret.update({
            'user': {
                'permissions': self.owner.get_permissions(user)
            },
            'node': {
                'id': self.owner._id,
                'api_url': self.owner.api_url,
                'url': self.owner.url,
                'is_registration': self.owner.is_registration,
            }
        })
        return ret

    def render_config_error(self, data):
        """

        """
        # Note: `config` is added to `self` in `AddonConfig::__init__`.
        template = lookup.get_template('project/addon/config_error.mako')
        return template.get_def('config_error').render(
            title=self.config.full_name,
            name=self.config.short_name,
            **data
        )

    #############
    # Callbacks #
    #############

    def before_page_load(self, node, user):
        """

        :param User user:
        :param Node node:

        """
        pass

    def before_remove_contributor(self, node, removed):
        """

        :param Node node:
        :param User removed:

        """
        pass

    def after_remove_contributor(self, node, removed):
        """

        :param Node node:
        :param User removed:

        """
        pass

    def before_make_public(self, node):

        """

        :param Node node:
        :returns: Alert message or None

        """
        pass

    def before_make_private(self, node):
        """

        :param Node node:
        :returns: Alert message or None

        """
        pass

    def after_set_privacy(self, node, permissions):
        """

        :param Node node:
        :param str permissions:

        """
        pass

    def before_fork(self, node, user):
        """

        :param Node node:
        :param User user:
        :returns: Alert message

        """
        pass

    def after_fork(self, node, fork, user, save=True):
        """

        :param Node node:
        :param Node fork:
        :param User user:
        :param bool save:
        :returns: Tuple of cloned settings and alert message

        """
        clone = self.clone()
        clone.owner = fork

        if save:
            clone.save()

        return clone, None

    def before_register(self, node, user):
        """

        :param Node node:
        :param User user:
        :returns: Alert message

        """
        pass

    def after_register(self, node, registration, user, save=True):
        """

        :param Node node:
        :param Node registration:
        :param User user:
        :param bool save:
        :returns: Tuple of cloned settings and alert message

        """
        clone = self.clone()
        clone.owner = registration

        if save:
            clone.save()

        return clone, None

    def after_delete(self, node, user):
        """

        :param Node node:
        :param User user:

        """
        pass

# TODO: No more magicks
def init_addon(app, addon_name, routes=True):
    """Load addon module return its create configuration object.

    If `log_fp` is provided, the addon's log templates will be appended
    to the file.

    :param app: Flask app object
    :param addon_name: Name of addon directory
    :param file log_fp: File pointer for the built logs file.
    :param bool routes: Add routes
    :return AddonConfig: AddonConfig configuration object if module found,
        else None

    """
    import_path = 'website.addons.{0}'.format(addon_name)

    # Import addon module
    addon_module = importlib.import_module(import_path)

    data = vars(addon_module)

    # Add routes
    if routes:
        for route_group in getattr(addon_module, 'ROUTES', []):
            process_rules(app, **route_group)

    # Build AddonConfig object
    return AddonConfig(
        **{
            key.lower(): value
            for key, value in data.iteritems()
        }
    )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright 2016 Yelp Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import absolute_import

from kafka.structs import PartitionMetadata

from kafka_utils.util.client import KafkaToolClient
from kafka_utils.util.zookeeper import ZK

LEADER_NOT_AVAILABLE_ERROR = 5
REPLICA_NOT_AVAILABLE_ERROR = 9

def get_topic_partition_metadata(hosts):
    """Returns topic-partition metadata from Kafka broker.

    kafka-python 1.3+ doesn't include partition metadata information in
    topic_partitions so we extract it from metadata ourselves.
    """
    kafka_client = KafkaToolClient(hosts, timeout=10)
    kafka_client.load_metadata_for_topics()
    topic_partitions = kafka_client.topic_partitions
    resp = kafka_client.send_metadata_request()

    for _, topic, partitions in resp.topics:
        for partition_error, partition, leader, replicas, isr in partitions:
            if topic_partitions.get(topic, {}).get(partition) is not None:
                topic_partitions[topic][partition] = PartitionMetadata(topic, partition, leader,
                                                                       replicas, isr, partition_error)
    return topic_partitions

def get_unavailable_brokers(zk, partition_metadata):
    """Returns the set of unavailable brokers from the difference of replica
    set of given partition to the set of available replicas.
    """
    topic_data = zk.get_topics(partition_metadata.topic)
    topic = partition_metadata.topic
    partition = partition_metadata.partition
    expected_replicas = set(topic_data[topic]['partitions'][str(partition)]['replicas'])
    available_replicas = set(partition_metadata.replicas)
    return expected_replicas - available_replicas

def get_topic_partition_with_error(cluster_config, error, fetch_unavailable_brokers=False):
    """Fetches the metadata from the cluster and returns the set of
    (topic, partition) tuples containing all the topic-partitions
    currently affected by the specified error. It also fetches unavailable-broker list
    if required."""

    metadata = get_topic_partition_metadata(cluster_config.broker_list)
    affected_partitions = set()
    if fetch_unavailable_brokers:
        unavailable_brokers = set()
    with ZK(cluster_config) as zk:
        for partitions in metadata.values():
            for partition_metadata in partitions.values():
                if int(partition_metadata.error) == error:
                    if fetch_unavailable_brokers:
                        unavailable_brokers |= get_unavailable_brokers(zk, partition_metadata)
                    affected_partitions.add((partition_metadata.topic, partition_metadata.partition))

    if fetch_unavailable_brokers:
        return affected_partitions, unavailable_brokers
    else:
        return affected_partitions
<commit_msg>Make kafka_check/offline resilient to topic deletion<commit_after># -*- coding: utf-8 -*-
# Copyright 2016 Yelp Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from __future__ import absolute_import

from collections import defaultdict

from kafka.structs import PartitionMetadata

from kafka_utils.util.client import KafkaToolClient
from kafka_utils.util.zookeeper import ZK

LEADER_NOT_AVAILABLE_ERROR = 5
REPLICA_NOT_AVAILABLE_ERROR = 9

def get_topic_partition_metadata(hosts):
    """Returns topic-partition metadata from Kafka broker.

    kafka-python 1.3+ doesn't include partition metadata information in
    topic_partitions so we extract it from metadata ourselves.
    """
    topic_partitions = defaultdict(dict)

    kafka_client = KafkaToolClient(hosts, timeout=10)
    resp = kafka_client.send_metadata_request()

    for _, topic, partitions in resp.topics:
        for partition_error, partition, leader, replicas, isr in partitions:
            topic_partitions[topic][partition] = PartitionMetadata(
                topic,
                partition,
                leader,
                replicas,
                isr,
                partition_error,
            )
    return topic_partitions

def get_unavailable_brokers(zk, partition_metadata):
    """Returns the set of unavailable brokers from the difference of replica
    set of given partition to the set of available replicas.
    """
    topic_data = zk.get_topics(partition_metadata.topic)
    topic = partition_metadata.topic
    partition = partition_metadata.partition
    expected_replicas = set(topic_data[topic]['partitions'][str(partition)]['replicas'])
    available_replicas = set(partition_metadata.replicas)
    return expected_replicas - available_replicas

def get_topic_partition_with_error(cluster_config, error, fetch_unavailable_brokers=False):
    """Fetches the metadata from the cluster and returns the set of
    (topic, partition) tuples containing all the topic-partitions
    currently affected by the specified error. It also fetches unavailable-broker list
    if required."""

    metadata = get_topic_partition_metadata(cluster_config.broker_list)
    affected_partitions = set()
    if fetch_unavailable_brokers:
        unavailable_brokers = set()
    with ZK(cluster_config) as zk:
        for partitions in metadata.values():
            for partition_metadata in partitions.values():
                if int(partition_metadata.error) == error:
                    if fetch_unavailable_brokers:
                        unavailable_brokers |= get_unavailable_brokers(zk, partition_metadata)
                    affected_partitions.add((partition_metadata.topic, partition_metadata.partition))

    if fetch_unavailable_brokers:
        return affected_partitions, unavailable_brokers
    else:
        return affected_partitions
<|endoftext|>