<commit_before>import csv
import inspect
import sys
from collections import namedtuple
from datetime import datetime

from django.core.management.base import BaseCommand, CommandError

import dateutil

from dimagi.utils.chunked import chunked
from dimagi.utils.retry import retry_on

from corehq.apps.es import CaseES, CaseSearchES, FormES
from corehq.elastic import ES_EXPORT_INSTANCE
from corehq.form_processor.backends.sql.dbaccessors import (
    CaseReindexAccessor,
    FormReindexAccessor,
)
from corehq.pillows.case_search import domains_needing_search_index
from corehq.util.dates import iso_string_to_datetime
from corehq.util.doc_processor.progress import (
    ProcessorProgressLogger,
    ProgressManager,
)
from corehq.util.doc_processor.sql import resumable_sql_model_iterator
from corehq.elastic import ESError

CHUNK_SIZE = 1000

RunConfig = namedtuple('RunConfig', ['iteration_key', 'domain', 'start_date', 'end_date', 'case_type'])
DataRow = namedtuple('DataRow', ['doc_id', 'doc_type', 'doc_subtype', 'domain', 'es_date', 'primary_date'])

ALL_DOMAINS = object()
HEADER_ROW = DataRow(
    doc_id='Doc ID',
    doc_type='Doc Type',
    doc_subtype='Doc Subtype',
    domain='Domain',
    es_date='ES Date',
    primary_date='Correct Date',
)

def get_csv_args(delimiter):
    return {
        'delimiter': delimiter,
        'lineterminator': '\n',
    }

retry_on_es_timeout = retry_on(ESError, delays=[2**x for x in range(10)])

class Command(BaseCommand):
    """
    Returns list of (doc_id, doc_type, doc_subtype, es_server_modified_on, primary_modified_on)
    tuples that are not updated in ES. Works for cases and forms.

    Can be used in conjunction with republish_doc_changes

        1. Generate tuples not updated in ES with extra debug columns
        $ ./manage.py stale_data_in_es case form --domain DOMAIN > stale_ids.tsv

        (Can call with just "case" or "form" if only want to use one data model type)

        2. Republish doc changes
        $ ./manage.py republish_doc_changes stale_ids.tsv
    """
    help = inspect.cleandoc(__doc__).split('\n')[0]

    def add_arguments(self, parser):
        parser.add_argument('data_models', nargs='+',
                            help='A list of data models to check. Valid options are "case" and "form".')
        parser.add_argument('--domain', default=ALL_DOMAINS)
        parser.add_argument('--iteration_key', help='Unique slug to identify this run. Used to allow resuming.')
        parser.add_argument(
            '--start',
            action='store',
            help='Only include data modified after this date',
        )
        parser.add_argument(
            '--end',
            action='store',
            help='Only include data modified before this date',
        )
        parser.add_argument(
            '--case_type',
            action='store',
            help='Only run for the specified case type',
        )
        parser.add_argument('--delimiter', default='\t', choices=('\t', ','))

    def handle(self, domain, data_models, delimiter, **options):
        data_models = set(data_models)

        default_start, default_end = datetime(2010, 1, 1), datetime.utcnow()
        start = dateutil.parser.parse(options['start']) if options['start'] else default_start
        end = dateutil.parser.parse(options['end']) if options['end'] else default_end
        case_type = options['case_type']
        iteration_key = options['iteration_key']
        if iteration_key:
            print(f'\nResuming previous run. Iteration Key:\n\t"{iteration_key}"\n', file=self.stderr)
        else:
            iteration_key = (
                f'stale_data_{datetime.utcnow().isoformat()}'
                f'-{"" if domain == ALL_DOMAINS else domain}'
                f'-{start.isoformat() if start != default_start else ""}'
                f'-{end.isoformat() if end != default_end else ""}'
                f'-{case_type or ""}'
            )
            print(f'\nStarting new run. Iteration key:\n\t"{iteration_key}"\n', file=self.stderr)

        run_config = RunConfig(iteration_key, domain, start, end, case_type)

        if run_config.domain is ALL_DOMAINS:
            print('Running for all domains', file=self.stderr)

        csv_writer = csv.writer(self.stdout, **get_csv_args(delimiter))

        def print_data_row(data_row):
            # Casting as str print `None` as 'None'
            csv_writer.writerow(map(str, data_row))

        print_data_row(HEADER_ROW)

        try:
            for data_model in data_models:
                try:
                    process_data_model_fn = DATA_MODEL_HELPERS[data_model.lower()]()
                except KeyError:
                    raise CommandError('Only valid options for data model are "{}"'.format(
                        '", "'.join(DATA_MODEL_HELPERS.keys())
                    ))
                data_rows = process_data_model_fn(run_config)

                for data_row in data_rows:
                    print_data_row(data_row)
        except:  # noqa: E722
            print(f'\nERROR: To resume the previous run add the "iteration_key" parameter to the command:\n'
                  f"\t--iteration_key '{iteration_key}'\n", file=self.stderr)
            raise

DATA_MODEL_HELPERS = {
    'case': lambda: CaseHelper.run,
    'form': lambda: FormHelper.run,
}

class CaseHelper:

    @classmethod
    def run(cls, run_config):
        for chunk in cls.get_sql_chunks(run_config):
            yield from cls._yield_missing_in_es(chunk)

    @staticmethod
    def get_sql_chunks(run_config):
        domain = run_config.domain if run_config.domain is not ALL_DOMAINS else None

        accessor = CaseReindexAccessor(
            domain,
            start_date=run_config.start_date, end_date=run_config.end_date,
            case_type=run_config.case_type
        )
        iteration_key = f'sql_cases-{run_config.iteration_key}'
        for chunk in _get_resumable_chunked_iterator(accessor, iteration_key, '[SQL cases] '):
            matching_records = [
                (case.case_id, case.type, case.server_modified_on, case.domain)
                for case in chunk
            ]
            yield matching_records

    @staticmethod
    def _yield_missing_in_es(chunk):
        case_ids, case_search_ids = CaseHelper._get_ids_from_chunk(chunk)
        es_modified_on_by_ids = CaseHelper._get_es_modified_dates(case_ids)
        case_search_es_modified_on_by_ids = CaseHelper._get_case_search_es_modified_dates(case_search_ids)
        for case_id, case_type, modified_on, domain in chunk:
            es_modified_on, es_domain = es_modified_on_by_ids.get(case_id, (None, None))
            if (es_modified_on, es_domain) != (modified_on, domain):
                yield DataRow(doc_id=case_id, doc_type='CommCareCase', doc_subtype=case_type, domain=domain,
                              es_date=es_modified_on, primary_date=modified_on)
            elif domain in domains_needing_search_index():
                es_modified_on, es_domain = case_search_es_modified_on_by_ids.get(case_id, (None, None))
                if (es_modified_on, es_domain) != (modified_on, domain):
                    yield DataRow(doc_id=case_id, doc_type='CommCareCase', doc_subtype=case_type, domain=domain,
                                  es_date=es_modified_on, primary_date=modified_on)

    @staticmethod
    @retry_on_es_timeout
    def _get_es_modified_dates(case_ids):
        results = (
            CaseES(es_instance_alias=ES_EXPORT_INSTANCE)
            .case_ids(case_ids)
            .values_list('_id', 'server_modified_on', 'domain')
        )
        return {_id: (iso_string_to_datetime(server_modified_on), domain)
                for _id, server_modified_on, domain in results}

    @staticmethod
    @retry_on_es_timeout
    def _get_case_search_es_modified_dates(case_ids):
        results = (
            CaseSearchES(es_instance_alias=ES_EXPORT_INSTANCE)
            .case_ids(case_ids)
            .values_list('_id', 'server_modified_on', 'domain')
        )
        return {_id: (iso_string_to_datetime(server_modified_on), domain)
                for _id, server_modified_on, domain in results}

    @staticmethod
    def _get_ids_from_chunk(chunk):
        case_ids = []
        case_search_ids = []
        for val in chunk:
            case_ids.append(val[0])
            if val[3] in domains_needing_search_index():
                case_search_ids.append(val[0])
        return case_ids, case_search_ids

class FormHelper:
    @classmethod
    def run(cls, run_config):
        for chunk in cls.get_sql_chunks(run_config):
            yield from cls._yield_missing_in_es(chunk)

    @staticmethod
    def get_sql_chunks(run_config):
        domain = run_config.domain if run_config.domain is not ALL_DOMAINS else None

        accessor = FormReindexAccessor(
            domain,
            start_date=run_config.start_date, end_date=run_config.end_date,
        )
        iteration_key = f'sql_forms-{run_config.iteration_key}'
        for chunk in _get_resumable_chunked_iterator(accessor, iteration_key, '[SQL forms] '):
            matching_records = [
                (form.form_id, form.doc_type, form.xmlns, form.received_on, form.domain)
                for form in chunk
                # Only check for "normal" and "archived" forms
                if form.is_normal or form.is_archived
            ]
            yield matching_records

    @staticmethod
    def _yield_missing_in_es(chunk):
        form_ids = [val[0] for val in chunk]
        es_modified_on_by_ids = FormHelper._get_es_modified_dates_for_forms(form_ids)
        for form_id, doc_type, xmlns, modified_on, domain in chunk:
            es_modified_on, es_doc_type, es_domain = es_modified_on_by_ids.get(form_id, (None, None, None))
            if (es_modified_on, es_doc_type, es_domain) != (modified_on, doc_type, domain):
                yield DataRow(doc_id=form_id, doc_type=doc_type, doc_subtype=xmlns, domain=domain,
                              es_date=es_modified_on, primary_date=modified_on)

    @staticmethod
    def _get_es_modified_dates_for_forms(form_ids):
        results = (
            FormES(es_instance_alias=ES_EXPORT_INSTANCE).remove_default_filters()
            .form_ids(form_ids)
            .values_list('_id', 'received_on', 'doc_type', 'domain')
        )
        return {_id: (iso_string_to_datetime(received_on), doc_type, domain)
                for _id, received_on, doc_type, domain in results}

def _get_resumable_chunked_iterator(dbaccessor, iteration_key, log_prefix):
    total_docs = 0
    for db in dbaccessor.sql_db_aliases:
        total_docs += dbaccessor.get_approximate_doc_count(db)

    iterable = resumable_sql_model_iterator(
        iteration_key,
        dbaccessor,
        chunk_size=CHUNK_SIZE,
        transform=lambda x: x
    )
    progress = ProgressManager(
        iterable,
        total=total_docs,
        reset=False,
        chunk_size=CHUNK_SIZE,
        logger=ProcessorProgressLogger(log_prefix, sys.stderr)
    )
    with progress:
        for chunk in chunked(iterable, CHUNK_SIZE):
            yield chunk
            progress.add(len(chunk))
<commit_msg>refresh case if newer in es<commit_after>import csv
import inspect
import sys
from collections import namedtuple
from datetime import datetime

from django.core.management.base import BaseCommand, CommandError

import dateutil

from dimagi.utils.chunked import chunked
from dimagi.utils.retry import retry_on

from corehq.apps.es import CaseES, CaseSearchES, FormES
from corehq.elastic import ES_EXPORT_INSTANCE
from corehq.form_processor.backends.sql.dbaccessors import (
    CaseReindexAccessor,
    FormReindexAccessor,
)
from corehq.form_processor.interfaces.dbaccessors import CaseAccessors
from corehq.pillows.case_search import domains_needing_search_index
from corehq.util.dates import iso_string_to_datetime
from corehq.util.doc_processor.progress import (
    ProcessorProgressLogger,
    ProgressManager,
)
from corehq.util.doc_processor.sql import resumable_sql_model_iterator
from corehq.elastic import ESError

CHUNK_SIZE = 1000

RunConfig = namedtuple('RunConfig', ['iteration_key', 'domain', 'start_date', 'end_date', 'case_type'])
DataRow = namedtuple('DataRow', ['doc_id', 'doc_type', 'doc_subtype', 'domain', 'es_date', 'primary_date'])

ALL_DOMAINS = object()
HEADER_ROW = DataRow(
    doc_id='Doc ID',
    doc_type='Doc Type',
    doc_subtype='Doc Subtype',
    domain='Domain',
    es_date='ES Date',
    primary_date='Correct Date',
)

def get_csv_args(delimiter):
    return {
        'delimiter': delimiter,
        'lineterminator': '\n',
    }

retry_on_es_timeout = retry_on(ESError, delays=[2**x for x in range(10)])

class Command(BaseCommand):
    """
    Returns list of (doc_id, doc_type, doc_subtype, es_server_modified_on, primary_modified_on)
    tuples that are not updated in ES. Works for cases and forms.

    Can be used in conjunction with republish_doc_changes

        1. Generate tuples not updated in ES with extra debug columns
        $ ./manage.py stale_data_in_es case form --domain DOMAIN > stale_ids.tsv

        (Can call with just "case" or "form" if only want to use one data model type)

        2. Republish doc changes
        $ ./manage.py republish_doc_changes stale_ids.tsv
    """
    help = inspect.cleandoc(__doc__).split('\n')[0]

    def add_arguments(self, parser):
        parser.add_argument('data_models', nargs='+',
                            help='A list of data models to check. Valid options are "case" and "form".')
        parser.add_argument('--domain', default=ALL_DOMAINS)
        parser.add_argument('--iteration_key', help='Unique slug to identify this run. Used to allow resuming.')
        parser.add_argument(
            '--start',
            action='store',
            help='Only include data modified after this date',
        )
        parser.add_argument(
            '--end',
            action='store',
            help='Only include data modified before this date',
        )
        parser.add_argument(
            '--case_type',
            action='store',
            help='Only run for the specified case type',
        )
        parser.add_argument('--delimiter', default='\t', choices=('\t', ','))

    def handle(self, domain, data_models, delimiter, **options):
        data_models = set(data_models)

        default_start, default_end = datetime(2010, 1, 1), datetime.utcnow()
        start = dateutil.parser.parse(options['start']) if options['start'] else default_start
        end = dateutil.parser.parse(options['end']) if options['end'] else default_end
        case_type = options['case_type']
        iteration_key = options['iteration_key']
        if iteration_key:
            print(f'\nResuming previous run. Iteration Key:\n\t"{iteration_key}"\n', file=self.stderr)
        else:
            iteration_key = (
                f'stale_data_{datetime.utcnow().isoformat()}'
                f'-{"" if domain == ALL_DOMAINS else domain}'
                f'-{start.isoformat() if start != default_start else ""}'
                f'-{end.isoformat() if end != default_end else ""}'
                f'-{case_type or ""}'
            )
            print(f'\nStarting new run. Iteration key:\n\t"{iteration_key}"\n', file=self.stderr)

        run_config = RunConfig(iteration_key, domain, start, end, case_type)

        if run_config.domain is ALL_DOMAINS:
            print('Running for all domains', file=self.stderr)

        csv_writer = csv.writer(self.stdout, **get_csv_args(delimiter))

        def print_data_row(data_row):
            # Casting as str print `None` as 'None'
            csv_writer.writerow(map(str, data_row))

        print_data_row(HEADER_ROW)

        try:
            for data_model in data_models:
                try:
                    process_data_model_fn = DATA_MODEL_HELPERS[data_model.lower()]()
                except KeyError:
                    raise CommandError('Only valid options for data model are "{}"'.format(
                        '", "'.join(DATA_MODEL_HELPERS.keys())
                    ))
                data_rows = process_data_model_fn(run_config)

                for data_row in data_rows:
                    print_data_row(data_row)
        except:  # noqa: E722
            print(f'\nERROR: To resume the previous run add the "iteration_key" parameter to the command:\n'
                  f"\t--iteration_key '{iteration_key}'\n", file=self.stderr)
            raise

DATA_MODEL_HELPERS = {
    'case': lambda: CaseHelper.run,
    'form': lambda: FormHelper.run,
}

class CaseHelper:

    @classmethod
    def run(cls, run_config):
        for chunk in cls.get_sql_chunks(run_config):
            yield from cls._yield_missing_in_es(chunk)

    @staticmethod
    def get_sql_chunks(run_config):
        domain = run_config.domain if run_config.domain is not ALL_DOMAINS else None

        accessor = CaseReindexAccessor(
            domain,
            start_date=run_config.start_date, end_date=run_config.end_date,
            case_type=run_config.case_type
        )
        iteration_key = f'sql_cases-{run_config.iteration_key}'
        for chunk in _get_resumable_chunked_iterator(accessor, iteration_key, '[SQL cases] '):
            matching_records = [
                (case.case_id, case.type, case.server_modified_on, case.domain)
                for case in chunk
            ]
            yield matching_records

    @staticmethod
    def _yield_missing_in_es(chunk):
        case_ids, case_search_ids = CaseHelper._get_ids_from_chunk(chunk)
        es_modified_on_by_ids = CaseHelper._get_es_modified_dates(case_ids)
        case_search_es_modified_on_by_ids = CaseHelper._get_case_search_es_modified_dates(case_search_ids)
        for case_id, case_type, modified_on, domain in chunk:
            stale, data_row = CaseHelper._check_stale(case_id, case_type, modified_on, domain, es_modified_on_by_ids, case_search_es_modified_on_by_ids)
            if stale:
                yield data_row

    @staticmethod
    def _check_stale(case_id, case_type, modified_on, domain,
                     es_modified_on_by_ids, case_search_es_modified_on_by_ids):
        es_modified_on, es_domain = es_modified_on_by_ids.get(case_id, (None, None))
        if (es_modified_on, es_domain) != (modified_on, domain):
            # if the doc is newer in ES than sql, refetch from sql to get newest
            if es_modified_on is not None and es_modified_on > modified_on:
                refreshed = CaseAccessors(domain).get_case(case_id)
                return _check_stale(case_id, case_type, refreshed.server_modified_on,
                                    refreshed.domain, es_modified_on_by_ids, case_search_es_modified_on_by_ids)
            else:
                return True, DataRow(doc_id=case_id, doc_type='CommCareCase', doc_subtype=case_type, domain=domain,
                                     es_date=es_modified_on, primary_date=modified_on)
        elif domain in domains_needing_search_index():
            es_modified_on, es_domain = case_search_es_modified_on_by_ids.get(case_id, (None, None))
            if (es_modified_on, es_domain) != (modified_on, domain):
                # if the doc is newer in ES than sql, refetch from sql to get newest
                if es_modified_on is not None and es_modified_on > modified_on:
                    refreshed = CaseAccessors(domain).get_case(case_id)
                    return _check_stale(case_id, case_type, refreshed.server_modified_on,
                                        refreshed.domain, es_modified_on_by_ids, case_search_es_modified_on_by_ids)
                else:
                    return True, DataRow(doc_id=case_id, doc_type='CommCareCase', doc_subtype=case_type, domain=domain,
                                        es_date=es_modified_on, primary_date=modified_on)
        return False, None

    @staticmethod
    @retry_on_es_timeout
    def _get_es_modified_dates(case_ids):
        results = (
            CaseES(es_instance_alias=ES_EXPORT_INSTANCE)
            .case_ids(case_ids)
            .values_list('_id', 'server_modified_on', 'domain')
        )
        return {_id: (iso_string_to_datetime(server_modified_on), domain)
                for _id, server_modified_on, domain in results}

    @staticmethod
    @retry_on_es_timeout
    def _get_case_search_es_modified_dates(case_ids):
        results = (
            CaseSearchES(es_instance_alias=ES_EXPORT_INSTANCE)
            .case_ids(case_ids)
            .values_list('_id', 'server_modified_on', 'domain')
        )
        return {_id: (iso_string_to_datetime(server_modified_on), domain)
                for _id, server_modified_on, domain in results}

    @staticmethod
    def _get_ids_from_chunk(chunk):
        case_ids = []
        case_search_ids = []
        for val in chunk:
            case_ids.append(val[0])
            if val[3] in domains_needing_search_index():
                case_search_ids.append(val[0])
        return case_ids, case_search_ids

class FormHelper:
    @classmethod
    def run(cls, run_config):
        for chunk in cls.get_sql_chunks(run_config):
            yield from cls._yield_missing_in_es(chunk)

    @staticmethod
    def get_sql_chunks(run_config):
        domain = run_config.domain if run_config.domain is not ALL_DOMAINS else None

        accessor = FormReindexAccessor(
            domain,
            start_date=run_config.start_date, end_date=run_config.end_date,
        )
        iteration_key = f'sql_forms-{run_config.iteration_key}'
        for chunk in _get_resumable_chunked_iterator(accessor, iteration_key, '[SQL forms] '):
            matching_records = [
                (form.form_id, form.doc_type, form.xmlns, form.received_on, form.domain)
                for form in chunk
                # Only check for "normal" and "archived" forms
                if form.is_normal or form.is_archived
            ]
            yield matching_records

    @staticmethod
    def _yield_missing_in_es(chunk):
        form_ids = [val[0] for val in chunk]
        es_modified_on_by_ids = FormHelper._get_es_modified_dates_for_forms(form_ids)
        for form_id, doc_type, xmlns, modified_on, domain in chunk:
            es_modified_on, es_doc_type, es_domain = es_modified_on_by_ids.get(form_id, (None, None, None))
            if (es_modified_on, es_doc_type, es_domain) != (modified_on, doc_type, domain):
                yield DataRow(doc_id=form_id, doc_type=doc_type, doc_subtype=xmlns, domain=domain,
                              es_date=es_modified_on, primary_date=modified_on)

    @staticmethod
    def _get_es_modified_dates_for_forms(form_ids):
        results = (
            FormES(es_instance_alias=ES_EXPORT_INSTANCE).remove_default_filters()
            .form_ids(form_ids)
            .values_list('_id', 'received_on', 'doc_type', 'domain')
        )
        return {_id: (iso_string_to_datetime(received_on), doc_type, domain)
                for _id, received_on, doc_type, domain in results}

def _get_resumable_chunked_iterator(dbaccessor, iteration_key, log_prefix):
    total_docs = 0
    for db in dbaccessor.sql_db_aliases:
        total_docs += dbaccessor.get_approximate_doc_count(db)

    iterable = resumable_sql_model_iterator(
        iteration_key,
        dbaccessor,
        chunk_size=CHUNK_SIZE,
        transform=lambda x: x
    )
    progress = ProgressManager(
        iterable,
        total=total_docs,
        reset=False,
        chunk_size=CHUNK_SIZE,
        logger=ProcessorProgressLogger(log_prefix, sys.stderr)
    )
    with progress:
        for chunk in chunked(iterable, CHUNK_SIZE):
            yield chunk
            progress.add(len(chunk))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
#
# Electrum - lightweight Bitcoin client
# Copyright (C) 2018 The Electrum developers
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import json

from .util import inv_dict

def read_json(filename, default):
    path = os.path.join(os.path.dirname(__file__), filename)
    try:
        with open(path, 'r') as f:
            r = json.loads(f.read())
    except:
        r = default
    return r

GIT_REPO_URL = "https://github.com/spesmilo/electrum"
GIT_REPO_ISSUES_URL = "https://github.com/spesmilo/electrum/issues"

class AbstractNet:

    @classmethod
    def max_checkpoint(cls) -> int:
        return max(0, len(cls.CHECKPOINTS) * 2016 - 1)

class BitcoinMainnet(AbstractNet):

    TESTNET = False
    WIF_PREFIX = 0x80
    ADDRTYPE_P2PKH = 0
    ADDRTYPE_P2SH = 5
    SEGWIT_HRP = "bc"
    GENESIS = "000000000019d6689c085ae165831e934ff763ae46a2a6c172b3f1b60a8ce26f"
    DEFAULT_PORTS = {'t': '50001', 's': '50002'}
    DEFAULT_SERVERS = read_json('servers.json', {})
    CHECKPOINTS = read_json('checkpoints.json', [])

    XPRV_HEADERS = {
        'standard':    0x0488ade4,  # xprv
        'p2wpkh-p2sh': 0x049d7878,  # yprv
        'p2wsh-p2sh':  0x0295b005,  # Yprv
        'p2wpkh':      0x04b2430c,  # zprv
        'p2wsh':       0x02aa7a99,  # Zprv
    }
    XPRV_HEADERS_INV = inv_dict(XPRV_HEADERS)
    XPUB_HEADERS = {
        'standard':    0x0488b21e,  # xpub
        'p2wpkh-p2sh': 0x049d7cb2,  # ypub
        'p2wsh-p2sh':  0x0295b43f,  # Ypub
        'p2wpkh':      0x04b24746,  # zpub
        'p2wsh':       0x02aa7ed3,  # Zpub
    }
    XPUB_HEADERS_INV = inv_dict(XPUB_HEADERS)
    BIP44_COIN_TYPE = 0

class BitcoinTestnet(AbstractNet):

    TESTNET = True
    WIF_PREFIX = 0xef
    ADDRTYPE_P2PKH = 111
    ADDRTYPE_P2SH = 196
    SEGWIT_HRP = "tb"
    GENESIS = "000000000933ea01ad0ee984209779baaec3ced90fa3f408719526f8d77f4943"
    DEFAULT_PORTS = {'t': '51001', 's': '51002'}
    DEFAULT_SERVERS = read_json('servers_testnet.json', {})
    CHECKPOINTS = read_json('checkpoints_testnet.json', [])

    XPRV_HEADERS = {
        'standard':    0x04358394,  # tprv
        'p2wpkh-p2sh': 0x044a4e28,  # uprv
        'p2wsh-p2sh':  0x024285b5,  # Uprv
        'p2wpkh':      0x045f18bc,  # vprv
        'p2wsh':       0x02575048,  # Vprv
    }
    XPRV_HEADERS_INV = inv_dict(XPRV_HEADERS)
    XPUB_HEADERS = {
        'standard':    0x043587cf,  # tpub
        'p2wpkh-p2sh': 0x044a5262,  # upub
        'p2wsh-p2sh':  0x024289ef,  # Upub
        'p2wpkh':      0x045f1cf6,  # vpub
        'p2wsh':       0x02575483,  # Vpub
    }
    XPUB_HEADERS_INV = inv_dict(XPUB_HEADERS)
    BIP44_COIN_TYPE = 1

class BitcoinSimnet(BitcoinTestnet):
    ADDRTYPE_P2PKH = 0x3f
    ADDRTYPE_P2SH = 0x7b
    SEGWIT_HRP = "sb"
    GENESIS = "683e86bd5c6d110d91b94b97137ba6bfe02dbbdb8e3dff722a669b5d69d77af6"
    WIF_PREFIX = 0x00
    DEFAULT_SERVERS = read_json('servers_regtest.json', {})  # Note: regtest!
    CHECKPOINTS = []

class BitcoinRegtest(BitcoinTestnet):

    SEGWIT_HRP = "bcrt"
    GENESIS = "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
    DEFAULT_SERVERS = read_json('servers_regtest.json', {})
    CHECKPOINTS = []

class BitcoinSimnet(BitcoinTestnet):

    SEGWIT_HRP = "sb"
    GENESIS = "683e86bd5c6d110d91b94b97137ba6bfe02dbbdb8e3dff722a669b5d69d77af6"
    DEFAULT_SERVERS = read_json('servers_regtest.json', {})
    CHECKPOINTS = []

# don't import net directly, import the module instead (so that net is singleton)
net = BitcoinMainnet

def set_simnet():
    global net
    net = BitcoinSimnet

def set_mainnet():
    global net
    net = BitcoinMainnet

def set_testnet():
    global net
    net = BitcoinTestnet

def set_regtest():
    global net
    net = BitcoinRegtest
<commit_msg>constants.py: remove duplicate Simnet declaration<commit_after># -*- coding: utf-8 -*-
#
# Electrum - lightweight Bitcoin client
# Copyright (C) 2018 The Electrum developers
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation files
# (the "Software"), to deal in the Software without restriction,
# including without limitation the rights to use, copy, modify, merge,
# publish, distribute, sublicense, and/or sell copies of the Software,
# and to permit persons to whom the Software is furnished to do so,
# subject to the following conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF
# MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT HOLDERS
# BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY, WHETHER IN AN
# ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, OUT OF OR IN
# CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

import os
import json

from .util import inv_dict

def read_json(filename, default):
    path = os.path.join(os.path.dirname(__file__), filename)
    try:
        with open(path, 'r') as f:
            r = json.loads(f.read())
    except:
        r = default
    return r

GIT_REPO_URL = "https://github.com/spesmilo/electrum"
GIT_REPO_ISSUES_URL = "https://github.com/spesmilo/electrum/issues"

class AbstractNet:

    @classmethod
    def max_checkpoint(cls) -> int:
        return max(0, len(cls.CHECKPOINTS) * 2016 - 1)

class BitcoinMainnet(AbstractNet):

    TESTNET = False
    WIF_PREFIX = 0x80
    ADDRTYPE_P2PKH = 0
    ADDRTYPE_P2SH = 5
    SEGWIT_HRP = "bc"
    GENESIS = "000000000019d6689c085ae165831e934ff763ae46a2a6c172b3f1b60a8ce26f"
    DEFAULT_PORTS = {'t': '50001', 's': '50002'}
    DEFAULT_SERVERS = read_json('servers.json', {})
    CHECKPOINTS = read_json('checkpoints.json', [])

    XPRV_HEADERS = {
        'standard':    0x0488ade4,  # xprv
        'p2wpkh-p2sh': 0x049d7878,  # yprv
        'p2wsh-p2sh':  0x0295b005,  # Yprv
        'p2wpkh':      0x04b2430c,  # zprv
        'p2wsh':       0x02aa7a99,  # Zprv
    }
    XPRV_HEADERS_INV = inv_dict(XPRV_HEADERS)
    XPUB_HEADERS = {
        'standard':    0x0488b21e,  # xpub
        'p2wpkh-p2sh': 0x049d7cb2,  # ypub
        'p2wsh-p2sh':  0x0295b43f,  # Ypub
        'p2wpkh':      0x04b24746,  # zpub
        'p2wsh':       0x02aa7ed3,  # Zpub
    }
    XPUB_HEADERS_INV = inv_dict(XPUB_HEADERS)
    BIP44_COIN_TYPE = 0

class BitcoinTestnet(AbstractNet):

    TESTNET = True
    WIF_PREFIX = 0xef
    ADDRTYPE_P2PKH = 111
    ADDRTYPE_P2SH = 196
    SEGWIT_HRP = "tb"
    GENESIS = "000000000933ea01ad0ee984209779baaec3ced90fa3f408719526f8d77f4943"
    DEFAULT_PORTS = {'t': '51001', 's': '51002'}
    DEFAULT_SERVERS = read_json('servers_testnet.json', {})
    CHECKPOINTS = read_json('checkpoints_testnet.json', [])

    XPRV_HEADERS = {
        'standard':    0x04358394,  # tprv
        'p2wpkh-p2sh': 0x044a4e28,  # uprv
        'p2wsh-p2sh':  0x024285b5,  # Uprv
        'p2wpkh':      0x045f18bc,  # vprv
        'p2wsh':       0x02575048,  # Vprv
    }
    XPRV_HEADERS_INV = inv_dict(XPRV_HEADERS)
    XPUB_HEADERS = {
        'standard':    0x043587cf,  # tpub
        'p2wpkh-p2sh': 0x044a5262,  # upub
        'p2wsh-p2sh':  0x024289ef,  # Upub
        'p2wpkh':      0x045f1cf6,  # vpub
        'p2wsh':       0x02575483,  # Vpub
    }
    XPUB_HEADERS_INV = inv_dict(XPUB_HEADERS)
    BIP44_COIN_TYPE = 1

class BitcoinRegtest(BitcoinTestnet):

    SEGWIT_HRP = "bcrt"
    GENESIS = "0f9188f13cb7b2c71f2a335e3a4fc328bf5beb436012afca590b1a11466e2206"
    DEFAULT_SERVERS = read_json('servers_regtest.json', {})
    CHECKPOINTS = []

class BitcoinSimnet(BitcoinTestnet):

    WIF_PREFIX = 0x64
    ADDRTYPE_P2PKH = 0x3f
    ADDRTYPE_P2SH = 0x7b
    SEGWIT_HRP = "sb"
    GENESIS = "683e86bd5c6d110d91b94b97137ba6bfe02dbbdb8e3dff722a669b5d69d77af6"
    DEFAULT_SERVERS = read_json('servers_regtest.json', {})
    CHECKPOINTS = []

# don't import net directly, import the module instead (so that net is singleton)
net = BitcoinMainnet

def set_simnet():
    global net
    net = BitcoinSimnet

def set_mainnet():
    global net
    net = BitcoinMainnet

def set_testnet():
    global net
    net = BitcoinTestnet

def set_regtest():
    global net
    net = BitcoinRegtest
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from django.contrib import admin
from django.contrib import messages
from django.core.urlresolvers import reverse
from django.db.models import Count
from django.utils.translation import ugettext as _
from django.utils.translation import ugettext_lazy
from lizard_auth_server import forms
from lizard_auth_server import models
from tls import request as tls_request

class InvitationAdmin(admin.ModelAdmin):
    model = models.Invitation
    list_display = [
        "email",
        "name",
        "is_activated",
        "user_profile_link",
        "created_at",
        "activated_on",
    ]
    search_fields = ["name", "email"]
    list_filter = ["is_activated"]
    list_select_related = ["user", "user__user_profile"]

    readonly_fields = [
        "created_at",
        "activation_key",
        "activation_key_date",
        "activated_on",
        "shortcut_urls",
    ]
    actions = ["send_new_activation_email"]
    filter_horizontal = ["portals"]

    def send_new_activation_email(self, request, queryset):
        for invitation in queryset:
            if invitation.is_activated:
                messages.error(
                    request,
                    "Invitation {} is already activated!".format(invitation),
                    fail_silently=False,
                )
            else:
                invitation.send_new_activation_email()

    send_new_activation_email.short_description = ugettext_lazy(
        "Resend the activation email, with a new key"
    )

    def shortcut_urls(self, obj):
        if obj.is_activated:
            return ""
        else:
            url = reverse(
                "lizard_auth_server.activate",
                kwargs={"activation_key": obj.activation_key},
            )
            return '{}'.format(url, _("Activate manually"))

    shortcut_urls.allow_tags = True
    shortcut_urls.short_description = ugettext_lazy("Shortcut URLs")

    def user_profile_link(self, obj):
        if not obj.user or not obj.user.user_profile:
            return
        url = reverse(
            "admin:lizard_auth_server_userprofile_change",
            args=[obj.user.user_profile.id],
        )
        return '→ {}'.format(url, obj.user)

    user_profile_link.allow_tags = True
    user_profile_link.short_description = ugettext_lazy("user profile")

class UserProfileAdmin(admin.ModelAdmin):
    model = models.UserProfile
    form = forms.UserProfileForm
    list_display = ["username", "full_name", "email", "created_at"]
    search_fields = ["user__first_name", "user__last_name", "user__email"]
    list_filter = ["portals", "organisations"]
    readonly_fields = ["updated_at", "created_at"]
    list_select_related = ["user"]

    filter_horizontal = ("portals", "organisations", "roles")
    readonly_fields = ["created_at", "updated_at", "first_name", "last_name", "email"]
    fieldsets = (
        (
            None,
            {
                "fields": [
                    "user",
                    "first_name",
                    "last_name",
                ]
            },
        ),
        (
            ugettext_lazy("V1 authorization"),
            {
                "fields": [
                    "portals",
                    "organisations",
                    "roles",
                ]
            },
        ),
        (
            ugettext_lazy("Dates"),
            {
                "fields": [
                    "created_at",
                    "updated_at",
                ]
            },
        ),
    )

class RelevantPortalFilter(admin.SimpleListFilter):
    title = _("portal")
    parameter_name = "portal__id__exact"

    def lookups(self, request, model_admin):
        return (
            models.Portal.objects.filter(roles__isnull=False)
            .distinct()
            .values_list("id", "name")
        )

    def queryset(self, request, queryset):
        if not self.value():
            return queryset
        return queryset.filter(portal=self.value())

class RoleInline(admin.TabularInline):
    model = models.Role
    fields = [
        "code",
        "name",
        "internal_description",
        "external_description",
        "num_inheriting_roles",
    ]
    readonly_fields = [
        "internal_description",
        "external_description",
        "num_inheriting_roles",
    ]
    # TODO: add show_change_link when we move to django 1.8.
    extra = 1

    def get_queryset(self, request):
        # Direct copy/paste from RoleAdmin (apart from the different 'super').
        queryset = super(RoleInline, self).get_queryset(request)
        return queryset.annotate(
            inheriting_roles_count=Count("inheriting_roles", distinct=True)
        )

    def num_inheriting_roles(self, obj):
        # Direct copy/paste from RoleAdmin
        count = obj.inheriting_roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_role_changelist")
        url += "?base_role={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_inheriting_roles.short_description = ugettext_lazy("number of inheriting roles")
    num_inheriting_roles.admin_order_field = "inheriting_roles_count"
    num_inheriting_roles.allow_tags = True

class OrganisationRoleInline(admin.TabularInline):
    model = models.OrganisationRole
    extra = 1

class RelevantBaseRoleFilter(admin.SimpleListFilter):
    title = _("base role")
    parameter_name = "base_role"

    def lookups(self, request, model_admin):
        return models.Role.objects.exclude(inheriting_roles__isnull=True).values_list(
            "id", "name"
        )

    def queryset(self, request, queryset):
        if not self.value():
            return queryset
        return queryset.filter(base_roles=self.value())

class RoleAdmin(admin.ModelAdmin):
    model = models.Role
    search_fields = [
        "code",
        "name",
        "portal__name",
        "portal__visit_url",
        "portal__allowed_domain",
        "external_description",
        "internal_description",
    ]
    list_display = [
        "code",
        "portal",
        "name",
        "internal_description",
        "num_organisation_roles",
        "num_inheriting_roles",
    ]
    list_filter = [RelevantPortalFilter, RelevantBaseRoleFilter]
    readonly_fields = ["unique_id"]
    filter_horizontal = ["inheriting_roles"]

    # inlines = [OrganisationRoleInline]
    # ^^^ This is easy to enable, but I [reinout] found it unclear how to use
    # it. Better to only have this inline on Organisation only.

    def get_queryset(self, request):
        queryset = super(RoleAdmin, self).get_queryset(request)
        return queryset.annotate(
            organisation_roles_count=Count("organisation_roles", distinct=True),
            inheriting_roles_count=Count("inheriting_roles", distinct=True),
        )

    def num_organisation_roles(self, obj):
        count = obj.organisation_roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_organisationrole_changelist")
        url += "?role__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_organisation_roles.short_description = ugettext_lazy(
        "number of organisation roles"
    )
    num_organisation_roles.admin_order_field = "organisation_roles_count"
    num_organisation_roles.allow_tags = True

    def num_inheriting_roles(self, obj):
        count = obj.inheriting_roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_role_changelist")
        url += "?base_role={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_inheriting_roles.short_description = ugettext_lazy("number of inheriting roles")
    num_inheriting_roles.admin_order_field = "inheriting_roles_count"
    num_inheriting_roles.allow_tags = True

class PortalAdmin(admin.ModelAdmin):
    model = models.Portal
    search_fields = ["name", "visit_url", "allowed_domain"]
    list_display = [
        "name",
        "visit_url",
        "allowed_domain",
        "num_user_profiles",
        "num_roles",
    ]
    readonly_fields = ["sso_secret", "sso_key", "v2_config"]
    inlines = [RoleInline]

    def get_queryset(self, request):
        queryset = super(PortalAdmin, self).get_queryset(request)
        return queryset.annotate(
            user_profiles_count=Count("user_profiles", distinct=True),
            roles_count=Count("roles", distinct=True),
        )

    def num_user_profiles(self, obj):
        count = obj.user_profiles_count
        url = reverse("admin:lizard_auth_server_userprofile_changelist")
        url += "?portals__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_user_profiles.short_description = ugettext_lazy("number of user profiles")
    num_user_profiles.admin_order_field = "user_profiles_count"
    num_user_profiles.allow_tags = True

    def num_roles(self, obj):
        count = obj.roles_count
        url = reverse("admin:lizard_auth_server_role_changelist")
        url += "?portal__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_roles.short_description = ugettext_lazy("number of roles")
    num_roles.admin_order_field = "roles_count"
    num_roles.allow_tags = True

    def v2_config(self, obj):
        config_lines = [
            "SSO_ENABLED = True",
            "SSO_USE_V2_LOGIN = True",
            "SSO_SERVER_API_START_URL = '{}'".format(
                tls_request.build_absolute_uri(
                    reverse("lizard_auth_server.api_v2.start")
                )
            ),
            "SSO_KEY = '{}'".format(obj.sso_key),
            "SSO_SECRET = '{}'".format(obj.sso_secret),
        ]
        return "
{}
".format("\n".join(config_lines))

    v2_config.short_description = ugettext_lazy("settings for the v2 API")
    v2_config.allow_tags = True

class OrganisationAdmin(admin.ModelAdmin):
    model = models.Organisation
    search_fields = ["name", "unique_id"]
    list_display = ["name", "num_user_profiles", "num_roles"]
    readonly_fields = ["unique_id"]
    inlines = [OrganisationRoleInline]

    def get_queryset(self, request):
        queryset = super(OrganisationAdmin, self).get_queryset(request)
        return queryset.annotate(
            user_profiles_count=Count("user_profiles", distinct=True),
            roles_count=Count("roles", distinct=True),
        )

    def num_user_profiles(self, obj):
        count = obj.user_profiles_count
        url = reverse("admin:lizard_auth_server_userprofile_changelist")
        url += "?organisations__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_user_profiles.short_description = ugettext_lazy("number of user profiles")
    num_user_profiles.admin_order_field = "user_profiles_count"
    num_user_profiles.allow_tags = True

    def num_roles(self, obj):
        count = obj.roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_organisationrole_changelist")
        url += "?organisation__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_roles.short_description = ugettext_lazy("number of roles")
    num_roles.admin_order_field = "roles_count"
    num_roles.allow_tags = True

class TokenAdmin(admin.ModelAdmin):
    model = models.Token
    search_fields = ["portal__name", "portal__visit_url", "portal__allowed_domain"]
    list_display = ["created", "portal", "user"]
    readonly_fields = ["created", "request_token", "auth_token"]

class RelevantOrganisationFilter(admin.SimpleListFilter):
    title = _("organisation")
    parameter_name = "organisation__id__exact"

    def lookups(self, request, model_admin):
        return (
            models.Organisation.objects.filter(organisation_roles__isnull=False)
            .distinct()
            .values_list("id", "name")
        )

    def queryset(self, request, queryset):
        if not self.value():
            return queryset
        return queryset.filter(organisation=self.value())

class OrganisationRoleAdmin(admin.ModelAdmin):
    model = models.OrganisationRole
    ordering = ("role__portal", "organisation", "role")
    list_display = ["__str__", "role", "organisation"]
    list_filter = ["role", RelevantOrganisationFilter]
    search_fields = ["organisation__name", "role__name", "role__portal__name"]

admin.site.register(models.Portal, PortalAdmin)
admin.site.register(models.Token, TokenAdmin)
admin.site.register(models.Invitation, InvitationAdmin)
admin.site.register(models.UserProfile, UserProfileAdmin)
admin.site.register(models.Role, RoleAdmin)
admin.site.register(models.Organisation, OrganisationAdmin)
admin.site.register(models.OrganisationRole, OrganisationRoleAdmin)
<commit_msg>Fixed django import<commit_after># -*- coding: utf-8 -*-
from django.contrib import admin
from django.contrib import messages
from django.db.models import Count
from django.urls import reverse
from django.utils.translation import ugettext as _
from django.utils.translation import ugettext_lazy
from lizard_auth_server import forms
from lizard_auth_server import models
from tls import request as tls_request

class InvitationAdmin(admin.ModelAdmin):
    model = models.Invitation
    list_display = [
        "email",
        "name",
        "is_activated",
        "user_profile_link",
        "created_at",
        "activated_on",
    ]
    search_fields = ["name", "email"]
    list_filter = ["is_activated"]
    list_select_related = ["user", "user__user_profile"]

    readonly_fields = [
        "created_at",
        "activation_key",
        "activation_key_date",
        "activated_on",
        "shortcut_urls",
    ]
    actions = ["send_new_activation_email"]
    filter_horizontal = ["portals"]

    def send_new_activation_email(self, request, queryset):
        for invitation in queryset:
            if invitation.is_activated:
                messages.error(
                    request,
                    "Invitation {} is already activated!".format(invitation),
                    fail_silently=False,
                )
            else:
                invitation.send_new_activation_email()

    send_new_activation_email.short_description = ugettext_lazy(
        "Resend the activation email, with a new key"
    )

    def shortcut_urls(self, obj):
        if obj.is_activated:
            return ""
        else:
            url = reverse(
                "lizard_auth_server.activate",
                kwargs={"activation_key": obj.activation_key},
            )
            return '{}'.format(url, _("Activate manually"))

    shortcut_urls.allow_tags = True
    shortcut_urls.short_description = ugettext_lazy("Shortcut URLs")

    def user_profile_link(self, obj):
        if not obj.user or not obj.user.user_profile:
            return
        url = reverse(
            "admin:lizard_auth_server_userprofile_change",
            args=[obj.user.user_profile.id],
        )
        return '→ {}'.format(url, obj.user)

    user_profile_link.allow_tags = True
    user_profile_link.short_description = ugettext_lazy("user profile")

class UserProfileAdmin(admin.ModelAdmin):
    model = models.UserProfile
    form = forms.UserProfileForm
    list_display = ["username", "full_name", "email", "created_at"]
    search_fields = ["user__first_name", "user__last_name", "user__email"]
    list_filter = ["portals", "organisations"]
    readonly_fields = ["updated_at", "created_at"]
    list_select_related = ["user"]

    filter_horizontal = ("portals", "organisations", "roles")
    readonly_fields = ["created_at", "updated_at", "first_name", "last_name", "email"]
    fieldsets = (
        (
            None,
            {
                "fields": [
                    "user",
                    "first_name",
                    "last_name",
                ]
            },
        ),
        (
            ugettext_lazy("V1 authorization"),
            {
                "fields": [
                    "portals",
                    "organisations",
                    "roles",
                ]
            },
        ),
        (
            ugettext_lazy("Dates"),
            {
                "fields": [
                    "created_at",
                    "updated_at",
                ]
            },
        ),
    )

class RelevantPortalFilter(admin.SimpleListFilter):
    title = _("portal")
    parameter_name = "portal__id__exact"

    def lookups(self, request, model_admin):
        return (
            models.Portal.objects.filter(roles__isnull=False)
            .distinct()
            .values_list("id", "name")
        )

    def queryset(self, request, queryset):
        if not self.value():
            return queryset
        return queryset.filter(portal=self.value())

class RoleInline(admin.TabularInline):
    model = models.Role
    fields = [
        "code",
        "name",
        "internal_description",
        "external_description",
        "num_inheriting_roles",
    ]
    readonly_fields = [
        "internal_description",
        "external_description",
        "num_inheriting_roles",
    ]
    # TODO: add show_change_link when we move to django 1.8.
    extra = 1

    def get_queryset(self, request):
        # Direct copy/paste from RoleAdmin (apart from the different 'super').
        queryset = super(RoleInline, self).get_queryset(request)
        return queryset.annotate(
            inheriting_roles_count=Count("inheriting_roles", distinct=True)
        )

    def num_inheriting_roles(self, obj):
        # Direct copy/paste from RoleAdmin
        count = obj.inheriting_roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_role_changelist")
        url += "?base_role={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_inheriting_roles.short_description = ugettext_lazy("number of inheriting roles")
    num_inheriting_roles.admin_order_field = "inheriting_roles_count"
    num_inheriting_roles.allow_tags = True

class OrganisationRoleInline(admin.TabularInline):
    model = models.OrganisationRole
    extra = 1

class RelevantBaseRoleFilter(admin.SimpleListFilter):
    title = _("base role")
    parameter_name = "base_role"

    def lookups(self, request, model_admin):
        return models.Role.objects.exclude(inheriting_roles__isnull=True).values_list(
            "id", "name"
        )

    def queryset(self, request, queryset):
        if not self.value():
            return queryset
        return queryset.filter(base_roles=self.value())

class RoleAdmin(admin.ModelAdmin):
    model = models.Role
    search_fields = [
        "code",
        "name",
        "portal__name",
        "portal__visit_url",
        "portal__allowed_domain",
        "external_description",
        "internal_description",
    ]
    list_display = [
        "code",
        "portal",
        "name",
        "internal_description",
        "num_organisation_roles",
        "num_inheriting_roles",
    ]
    list_filter = [RelevantPortalFilter, RelevantBaseRoleFilter]
    readonly_fields = ["unique_id"]
    filter_horizontal = ["inheriting_roles"]

    # inlines = [OrganisationRoleInline]
    # ^^^ This is easy to enable, but I [reinout] found it unclear how to use
    # it. Better to only have this inline on Organisation only.

    def get_queryset(self, request):
        queryset = super(RoleAdmin, self).get_queryset(request)
        return queryset.annotate(
            organisation_roles_count=Count("organisation_roles", distinct=True),
            inheriting_roles_count=Count("inheriting_roles", distinct=True),
        )

    def num_organisation_roles(self, obj):
        count = obj.organisation_roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_organisationrole_changelist")
        url += "?role__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_organisation_roles.short_description = ugettext_lazy(
        "number of organisation roles"
    )
    num_organisation_roles.admin_order_field = "organisation_roles_count"
    num_organisation_roles.allow_tags = True

    def num_inheriting_roles(self, obj):
        count = obj.inheriting_roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_role_changelist")
        url += "?base_role={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_inheriting_roles.short_description = ugettext_lazy("number of inheriting roles")
    num_inheriting_roles.admin_order_field = "inheriting_roles_count"
    num_inheriting_roles.allow_tags = True

class PortalAdmin(admin.ModelAdmin):
    model = models.Portal
    search_fields = ["name", "visit_url", "allowed_domain"]
    list_display = [
        "name",
        "visit_url",
        "allowed_domain",
        "num_user_profiles",
        "num_roles",
    ]
    readonly_fields = ["sso_secret", "sso_key", "v2_config"]
    inlines = [RoleInline]

    def get_queryset(self, request):
        queryset = super(PortalAdmin, self).get_queryset(request)
        return queryset.annotate(
            user_profiles_count=Count("user_profiles", distinct=True),
            roles_count=Count("roles", distinct=True),
        )

    def num_user_profiles(self, obj):
        count = obj.user_profiles_count
        url = reverse("admin:lizard_auth_server_userprofile_changelist")
        url += "?portals__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_user_profiles.short_description = ugettext_lazy("number of user profiles")
    num_user_profiles.admin_order_field = "user_profiles_count"
    num_user_profiles.allow_tags = True

    def num_roles(self, obj):
        count = obj.roles_count
        url = reverse("admin:lizard_auth_server_role_changelist")
        url += "?portal__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_roles.short_description = ugettext_lazy("number of roles")
    num_roles.admin_order_field = "roles_count"
    num_roles.allow_tags = True

    def v2_config(self, obj):
        config_lines = [
            "SSO_ENABLED = True",
            "SSO_USE_V2_LOGIN = True",
            "SSO_SERVER_API_START_URL = '{}'".format(
                tls_request.build_absolute_uri(
                    reverse("lizard_auth_server.api_v2.start")
                )
            ),
            "SSO_KEY = '{}'".format(obj.sso_key),
            "SSO_SECRET = '{}'".format(obj.sso_secret),
        ]
        return "
{}
".format("\n".join(config_lines))

    v2_config.short_description = ugettext_lazy("settings for the v2 API")
    v2_config.allow_tags = True

class OrganisationAdmin(admin.ModelAdmin):
    model = models.Organisation
    search_fields = ["name", "unique_id"]
    list_display = ["name", "num_user_profiles", "num_roles"]
    readonly_fields = ["unique_id"]
    inlines = [OrganisationRoleInline]

    def get_queryset(self, request):
        queryset = super(OrganisationAdmin, self).get_queryset(request)
        return queryset.annotate(
            user_profiles_count=Count("user_profiles", distinct=True),
            roles_count=Count("roles", distinct=True),
        )

    def num_user_profiles(self, obj):
        count = obj.user_profiles_count
        url = reverse("admin:lizard_auth_server_userprofile_changelist")
        url += "?organisations__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_user_profiles.short_description = ugettext_lazy("number of user profiles")
    num_user_profiles.admin_order_field = "user_profiles_count"
    num_user_profiles.allow_tags = True

    def num_roles(self, obj):
        count = obj.roles_count
        if not count:
            return ""
        url = reverse("admin:lizard_auth_server_organisationrole_changelist")
        url += "?organisation__id__exact={}".format(obj.id)
        return '→ {}'.format(url, count)

    num_roles.short_description = ugettext_lazy("number of roles")
    num_roles.admin_order_field = "roles_count"
    num_roles.allow_tags = True

class TokenAdmin(admin.ModelAdmin):
    model = models.Token
    search_fields = ["portal__name", "portal__visit_url", "portal__allowed_domain"]
    list_display = ["created", "portal", "user"]
    readonly_fields = ["created", "request_token", "auth_token"]

class RelevantOrganisationFilter(admin.SimpleListFilter):
    title = _("organisation")
    parameter_name = "organisation__id__exact"

    def lookups(self, request, model_admin):
        return (
            models.Organisation.objects.filter(organisation_roles__isnull=False)
            .distinct()
            .values_list("id", "name")
        )

    def queryset(self, request, queryset):
        if not self.value():
            return queryset
        return queryset.filter(organisation=self.value())

class OrganisationRoleAdmin(admin.ModelAdmin):
    model = models.OrganisationRole
    ordering = ("role__portal", "organisation", "role")
    list_display = ["__str__", "role", "organisation"]
    list_filter = ["role", RelevantOrganisationFilter]
    search_fields = ["organisation__name", "role__name", "role__portal__name"]

admin.site.register(models.Portal, PortalAdmin)
admin.site.register(models.Token, TokenAdmin)
admin.site.register(models.Invitation, InvitationAdmin)
admin.site.register(models.UserProfile, UserProfileAdmin)
admin.site.register(models.Role, RoleAdmin)
admin.site.register(models.Organisation, OrganisationAdmin)
admin.site.register(models.OrganisationRole, OrganisationRoleAdmin)
<|endoftext|>


<commit_before># encoding: utf-8

from __future__ import absolute_import
from __future__ import with_statement

import multiprocessing
from StringIO import StringIO
import Queue
import threading
import time

import twill
import twill.commands
import twill.parse

from twillmanager.db import get_db_connection, close_db_connection
from twillmanager.mail import create_mailer
from twillmanager.log import logger
import twillmanager.async

__all__ = ['STATUS_FAILED', 'STATUS_OK', 'STATUS_UNKNOWN', 'Watch', 'WorkerSet']

# Consts for statuses
STATUS_OK = 'OK'
STATUS_FAILED = 'FAILED'
STATUS_UNKNOWN = 'UNKNOWN'

class Watch(object):
    """ A simple data transfer object for describing watches (with data access methods
        for loading/storing watches into the database)
    """
    # Rows in database table
    COLUMNS = ['id',
               'name',
               'interval',
               'script',
               'emails',
               'status',
               'time',
               'reminder_interval',
               'last_alert',
              ]

    def __init__(self, name, interval, script,
            emails=None, status=STATUS_UNKNOWN, time=None,
            reminder_interval=600, last_alert=None, id=None):
        """ Create a new Watch

            :param name: Name of the watch
            :param interval: Interval (in seconds) between watch runs
            :param script: Twill script that should be executed by the watch
            :param emails: String containing all the e-mail addresses (comma-separated)
                of people who should receive failure notifications.
            :param status: Current status of the watch
            :param time: Last update time (as number of seconds since epoch - obtained by call to `time.time()`)
            :param reminder_interval: Interval (seconds) after which reminder that the watch is still down is sent
            :param last_alert: Time of last alert sent (as number of seconds since epoch - obtained by call to `time.time()`)
            :param id: Key in the database of the watch
        """
        self.name = name
        self.interval = interval
        self.script = script
        self.emails = emails
        self.status = status
        self.time = time
        self.reminder_interval = reminder_interval
        self.last_alert = last_alert
        self.id = id

    def formatted_time(self):
        if self.time:
            return time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime(self.time));
        else:
            return None

    def dict(self):
        """ Watch status data as dictionary (not a complete watch) """
        data = {}
        data['id'] = self.id
        data['name'] = self.name
        data['status'] = self.status
        data['time'] = self.formatted_time();
        return data

    @classmethod
    def columns(cls):
        return ','.join(cls.COLUMNS)

    @classmethod
    def construct_from_row(cls, row):
        kwargs = dict()
        for name in cls.COLUMNS:
            try:
                kwargs[name] = row[name]
            except IndexError:
                pass
        return cls(**kwargs)

    def save(self, connection):
        """ Save the watch into the database (using given connection) """
        if self.id is None:
            self.insert(connection)
        else:
            self.update(connection)

    def insert(self, connection):
        """ Insert the watch into the database (using given connection) """
        c = connection.cursor()
        c.execute("INSERT INTO twills (name, interval, script, emails, status, time, reminder_interval, last_alert) VALUES (?,?,?,?,?,?,?,?)",
            (self.name, self.interval, self.script, self.emails, self.status, self.time, self.reminder_interval, self.last_alert))
        self.id = c.lastrowid
        c.close()
        connection.commit()

    def update(self, connection):
        """ Update the watch into the database (using given connection) """
        assert self.id is not None
        c = connection.cursor()
        c.execute("UPDATE twills SET name=?, interval=?, script=?, emails=?, status=?, time=?, reminder_interval=?, last_alert=? WHERE id = ?",
            (self.name, self.interval, self.script, self.emails, self.status, self.time, self.reminder_interval, self.last_alert, self.id))
        c.close()
        connection.commit()

    def update_status(self, connection):
        """ Updates only information related to status check
            (status, check time, messages).

            This avoids overwriting script definition by a worker.
        """
        assert self.id is not None
        c = connection.cursor()
        c.execute("UPDATE twills SET status=?, time=?, last_alert=? WHERE id = ?",
            (self.status, self.time, self.last_alert, self.id))
        c.close()
        connection.commit()

    def delete(self, connection):
        """ Deletes the watch from the database (using given connection) """
        c = connection.cursor()
        c.execute("DELETE FROM twills WHERE id = ?", (self.id,))
        c.close()
        connection.commit()

    @classmethod
    def load(cls, id, connection):
        """ Loads the watch with given id """
        watch = None
        c = connection.cursor()
        c.execute("SELECT " + cls.columns() + " FROM twills WHERE id = ?", (id,))
        for row in c:
            watch = cls.construct_from_row(row)
        c.close()
        return watch

    @classmethod
    def load_by_name(cls, name, connection):
        """ Loads the watch with given name """
        watch = None
        c = connection.cursor()
        c.execute("SELECT " + cls.columns() + " FROM twills WHERE name = ?", (name,))
        for row in c:
            watch = cls.construct_from_row(row)
        c.close()
        return watch

    @classmethod
    def load_all(cls, connection):
        """ Loads all watches """
        watches = []
        c = connection.cursor()
        c.execute("SELECT " + cls.columns() + " FROM twills ORDER BY name")
        for row in c:
            watches.append(cls.construct_from_row(row))
        c.close()
        return watches

class WorkerProxy(twillmanager.async.WorkerProxy):
    """ A proxy for controlling `Worker` processes """
    def __init__(self, id, config):
        twillmanager.async.WorkerProxy.__init__(self)
        self.id = id
        self.config = config
        self.on_twill_start = None # executed when twill check starts
        self.on_twill_end = None # executed when twill check starts

    def set_twill_callbacks(self, on_start, on_end):
        """ Sets callbacks to be executed when twill check starts and ends.
            Must be set before starting the worker

            :param on_start: Executed when twill starts. A zero-argument callable
            :param on_start: Executed when twill end. A zero-argument callable
        """
        assert not self.already_started(), "Can't call set_twill_callbacks on worker that is already started"
        self.on_twill_start = on_start
        self.on_twill_end = on_end

    def make_worker(self, queue):
        return Worker(queue, self.id, self.config, self.on_twill_start, self.on_twill_end)

    def quit(self):
        """ Send 'quit' signal to the worker """
        self.queue_command('quit')

    def execute(self):
        """ Send 'execute the script now' signal to the worker """
        self.queue_command('execute')

class Worker(twillmanager.async.Worker):
    """ Worker - a process that monitors if given twill script executes properly"""
    def __init__(self, queue, id, config, on_start=None, on_end=None):
        """ Creates a new `Worker`
            :param queue: The command queue, as needed by `twillmanager.async.Worker`
            :param id: Id (database primary key) of the watch to use
            :param config: Configuration dict to be used (needed for e-mail addresses etc)
            :param on_start: Callable (zero-argument) to be invoked when
                twill script starts execution
            :param on_end: Callable (zero-argument) to be invoked when twill
                script finishes execution (disregarding status)
        """
        twillmanager.async.Worker.__init__(self, queue)
        self.id = id
        self.config = config
        self.watch = None
        self.connection = None
        self.on_start = on_start
        self.on_end = on_end

    def main(self):
        """ Process main function """
        # to make sure we do not use inherited descriptor
        # from the parent process
        close_db_connection()
        self.connection = get_db_connection(self.config)
        self.watch = Watch.load(self.id, self.connection)

        if self.watch:
            logger.info("Starting worker for watch `%s` (id: %s)" % (self.watch.name, self.id))
            self.tick_interval = self.watch.interval
            twillmanager.async.Worker.main(self)
        else:
            logger.warn("Failed to start worker for watch (id: %s) - no such watch" % self.id)

    def tick(self):
        """ Executed every self.watch.interval seconds """
        self.execute()

    def quit(self):
        self.running = False

    def execute(self):
        """ Called by `tick` and when `execute` schedules immediate script execution."""
        try: # large try block to ensure on_end is called
            if self.on_start:
                self.on_start()

            new_status, output = self.execute_script()

            old_status = self.watch.status
            self.watch.status = new_status
            self.watch.time = time.time()
            self.watch.update_status(self.connection)

            msg = "Status for watch `%s` (id: %s): %s" % (self.watch.name, self.id, new_status)

            if new_status != STATUS_OK:
                logger.warn(msg)
            else:
                logger.info(msg)

            # when was last e-mail alert sent
            if self.watch.last_alert is None:
                time_since_last_alert = None
            else:
                time_since_last_alert = self.watch.time - self.watch.last_alert

            status_has_changed = (old_status != new_status)

            # whether last alert was sent long ago enough to send a failure reminder
            # (normally e-mails are sent only on change, but a reminder is sent
            # if the watch keeps failing)
            if self.watch.reminder_interval:
                last_alert_was_long_ago = time_since_last_alert is None or time_since_last_alert > self.watch.reminder_interval
            else:
                last_alert_was_long_ago = False

            if status_has_changed or (last_alert_was_long_ago and new_status == STATUS_FAILED):
                logger.info("Sending notification for watch `%s` (id: %s)" % (self.watch.name, self.id))
                self.status_notify(old_status, new_status, out.getvalue())
                self.watch.last_alert = time.time()
                self.watch.update_status(self.connection)
        except Exception, e:
            logger.error("Worker `%s` (id: %s) failed with exception: %s" % (self.watch.name, self.id, e.message))
            raise
        finally:
            if self.on_end:
                self.on_end()

    def execute_script(self):
        """ Executes twill script. Returns a tuple status, output """
        out = StringIO()
        # execute the twill, catching any exceptions
        try:
            twill.set_errout(out)
            twill.set_output(out)
            twill.parse._execute_script(self.watch.script.split("\n"))
            status = STATUS_OK
        except Exception, e:
            status = STATUS_FAILED
        finally:
            twill.commands.reset_error()
            twill.commands.reset_output()
        return status, out.getvalue()

    def status_notify(self, old_status, new_status, message):
        """ Sends out notifications about watch status change """
        recipients = self.watch.emails

        if not recipients:
            return

        # strip
        recipients = [r.strip() for r in recipients.split(',')]
        # remove empty addresses
        recipients = [r for r in recipients if r]

        if len(recipients) == 0:
            return

        sender = self.config['mail.from']

        if old_status != new_status:
            subject = "Watch %s status change %s -> %s" % (self.watch.name, old_status, new_status)
        else:
            subject = "Watch %s status is still %s" % (self.watch.name, new_status)

        body = "Script:\n%s\n\nResult:\n%s" % (self.watch.script, message)

        mailer = create_mailer(self.config)
        mailer.send_mail(sender, recipients, subject, body)

class WorkerSet(object):
    """ Object managing spawning of workers."""
    def __init__(self, config):
        """ Create a new `WorkerSet`.

            :param config: Configuration dict to be passed to spawned workers
        """
        # synchronization between threads (WorkerSet is used from CherryPy)
        self._lock = threading.RLock()
        self.workers = {}
        self.now_building = {}
        self.config = config

        # thread that checks for workers that died unexpectedly
        # and listens to their status update messages
        # only tuples (command, argument) should be put into that queue.
        self.manager_thread_queue = multiprocessing.Queue(0)
        self.manager_thread = threading.Thread(target=self.manager_thread_main)
        self.manager_thread.daemon = True
        self.manager_thread.start()

    def worker_status_dict(self, id):
        with self._lock:
            return {'alive': self.is_alive(id), 'building': self.is_building(id)}

    def get(self, id):
        return self.workers.get(id, None)

    def finish(self):
        """ Call this to clean up when the application is shut down """
        self.manager_thread_queue.put(('quit', None))
        self.manager_thread.join()

    def is_alive(self, id):
        """ Check if worker with given id is alive """
        with self._lock:
            return id in self.workers and self.workers[id].is_alive()

    def is_building(self, id):
        """ Check if worker with given id is currently building"""
        with self._lock:
            return self.now_building.get(id, False)

    def check_now(self, id):
        """ Tell the worker with given id to check immediately.
            This also restarts the worker if it died.
        """
        with self._lock:
            self.restart(id)
            self.now_building[id] = True
            self.workers[id].execute()

    def restart(self, id):
        """ Restarts given worker. """
        with self._lock:
            self.remove(id)
            self.add(id)

    def add(self, id):
        """ Adds worker for watch with given id.

            It is safe to call this method even if the worker is already running.
        """
        with self._lock:
            if id in self.workers:
                return

            def on_start():
                self.manager_thread_queue.put(('start', id))

            def on_end():
                self.manager_thread_queue.put(('end', id))

            worker = WorkerProxy(id, self.config)
            worker.set_twill_callbacks(on_start, on_end)

            self.workers[id] = worker
            self.now_building[id] = False
            worker.start(True)

    def remove(self, id):
        """ Removes worker for watch with given id.

            It is safe to call this method even if the worker is not running.
        """
        with self._lock:
            if id in self.workers:
                worker = self.workers[id]
                worker.quit()
                del self.workers[id]
                del self.now_building[id]
                if worker.process:
                    worker.process.join()

    def manager_thread_main(self):
        """ Checks for workers that died unexpectedly and listens to their
            status update messages.
        """
        while True:
            # wait up to 60 seconds
            try:
                command, argument = self.manager_thread_queue.get(True, 60)

                if command == 'quit':
                    break
                elif command == 'start':
                    with self._lock:
                        self.now_building[argument] = True
                elif command == 'end':
                    with self._lock:
                        self.now_building[argument] = False
                else:
                    logger.warn("Unknown command to manager thread: %s" % command)

            except Queue.Empty:
                pass

            # this one is to remove zombie processes
            multiprocessing.active_children()

            with self._lock:
                ids_to_restart = []
                for id, worker in self.workers.items():
                    if not worker.is_alive():
                        ids_to_restart.append(id)

                for id in ids_to_restart:
                    self.restart(id)
<commit_msg>Old variable reference removed.<commit_after># encoding: utf-8

from __future__ import absolute_import
from __future__ import with_statement

import multiprocessing
from StringIO import StringIO
import Queue
import threading
import time

import twill
import twill.commands
import twill.parse

from twillmanager.db import get_db_connection, close_db_connection
from twillmanager.mail import create_mailer
from twillmanager.log import logger
import twillmanager.async

__all__ = ['STATUS_FAILED', 'STATUS_OK', 'STATUS_UNKNOWN', 'Watch', 'WorkerSet']

# Consts for statuses
STATUS_OK = 'OK'
STATUS_FAILED = 'FAILED'
STATUS_UNKNOWN = 'UNKNOWN'

class Watch(object):
    """ A simple data transfer object for describing watches (with data access methods
        for loading/storing watches into the database)
    """
    # Rows in database table
    COLUMNS = ['id',
               'name',
               'interval',
               'script',
               'emails',
               'status',
               'time',
               'reminder_interval',
               'last_alert',
              ]

    def __init__(self, name, interval, script,
            emails=None, status=STATUS_UNKNOWN, time=None,
            reminder_interval=600, last_alert=None, id=None):
        """ Create a new Watch

            :param name: Name of the watch
            :param interval: Interval (in seconds) between watch runs
            :param script: Twill script that should be executed by the watch
            :param emails: String containing all the e-mail addresses (comma-separated)
                of people who should receive failure notifications.
            :param status: Current status of the watch
            :param time: Last update time (as number of seconds since epoch - obtained by call to `time.time()`)
            :param reminder_interval: Interval (seconds) after which reminder that the watch is still down is sent
            :param last_alert: Time of last alert sent (as number of seconds since epoch - obtained by call to `time.time()`)
            :param id: Key in the database of the watch
        """
        self.name = name
        self.interval = interval
        self.script = script
        self.emails = emails
        self.status = status
        self.time = time
        self.reminder_interval = reminder_interval
        self.last_alert = last_alert
        self.id = id

    def formatted_time(self):
        if self.time:
            return time.strftime('%Y-%m-%d %H:%M:%S UTC', time.gmtime(self.time));
        else:
            return None

    def dict(self):
        """ Watch status data as dictionary (not a complete watch) """
        data = {}
        data['id'] = self.id
        data['name'] = self.name
        data['status'] = self.status
        data['time'] = self.formatted_time();
        return data

    @classmethod
    def columns(cls):
        return ','.join(cls.COLUMNS)

    @classmethod
    def construct_from_row(cls, row):
        kwargs = dict()
        for name in cls.COLUMNS:
            try:
                kwargs[name] = row[name]
            except IndexError:
                pass
        return cls(**kwargs)

    def save(self, connection):
        """ Save the watch into the database (using given connection) """
        if self.id is None:
            self.insert(connection)
        else:
            self.update(connection)

    def insert(self, connection):
        """ Insert the watch into the database (using given connection) """
        c = connection.cursor()
        c.execute("INSERT INTO twills (name, interval, script, emails, status, time, reminder_interval, last_alert) VALUES (?,?,?,?,?,?,?,?)",
            (self.name, self.interval, self.script, self.emails, self.status, self.time, self.reminder_interval, self.last_alert))
        self.id = c.lastrowid
        c.close()
        connection.commit()

    def update(self, connection):
        """ Update the watch into the database (using given connection) """
        assert self.id is not None
        c = connection.cursor()
        c.execute("UPDATE twills SET name=?, interval=?, script=?, emails=?, status=?, time=?, reminder_interval=?, last_alert=? WHERE id = ?",
            (self.name, self.interval, self.script, self.emails, self.status, self.time, self.reminder_interval, self.last_alert, self.id))
        c.close()
        connection.commit()

    def update_status(self, connection):
        """ Updates only information related to status check
            (status, check time, messages).

            This avoids overwriting script definition by a worker.
        """
        assert self.id is not None
        c = connection.cursor()
        c.execute("UPDATE twills SET status=?, time=?, last_alert=? WHERE id = ?",
            (self.status, self.time, self.last_alert, self.id))
        c.close()
        connection.commit()

    def delete(self, connection):
        """ Deletes the watch from the database (using given connection) """
        c = connection.cursor()
        c.execute("DELETE FROM twills WHERE id = ?", (self.id,))
        c.close()
        connection.commit()

    @classmethod
    def load(cls, id, connection):
        """ Loads the watch with given id """
        watch = None
        c = connection.cursor()
        c.execute("SELECT " + cls.columns() + " FROM twills WHERE id = ?", (id,))
        for row in c:
            watch = cls.construct_from_row(row)
        c.close()
        return watch

    @classmethod
    def load_by_name(cls, name, connection):
        """ Loads the watch with given name """
        watch = None
        c = connection.cursor()
        c.execute("SELECT " + cls.columns() + " FROM twills WHERE name = ?", (name,))
        for row in c:
            watch = cls.construct_from_row(row)
        c.close()
        return watch

    @classmethod
    def load_all(cls, connection):
        """ Loads all watches """
        watches = []
        c = connection.cursor()
        c.execute("SELECT " + cls.columns() + " FROM twills ORDER BY name")
        for row in c:
            watches.append(cls.construct_from_row(row))
        c.close()
        return watches

class WorkerProxy(twillmanager.async.WorkerProxy):
    """ A proxy for controlling `Worker` processes """
    def __init__(self, id, config):
        twillmanager.async.WorkerProxy.__init__(self)
        self.id = id
        self.config = config
        self.on_twill_start = None # executed when twill check starts
        self.on_twill_end = None # executed when twill check starts

    def set_twill_callbacks(self, on_start, on_end):
        """ Sets callbacks to be executed when twill check starts and ends.
            Must be set before starting the worker

            :param on_start: Executed when twill starts. A zero-argument callable
            :param on_start: Executed when twill end. A zero-argument callable
        """
        assert not self.already_started(), "Can't call set_twill_callbacks on worker that is already started"
        self.on_twill_start = on_start
        self.on_twill_end = on_end

    def make_worker(self, queue):
        return Worker(queue, self.id, self.config, self.on_twill_start, self.on_twill_end)

    def quit(self):
        """ Send 'quit' signal to the worker """
        self.queue_command('quit')

    def execute(self):
        """ Send 'execute the script now' signal to the worker """
        self.queue_command('execute')

class Worker(twillmanager.async.Worker):
    """ Worker - a process that monitors if given twill script executes properly"""
    def __init__(self, queue, id, config, on_start=None, on_end=None):
        """ Creates a new `Worker`
            :param queue: The command queue, as needed by `twillmanager.async.Worker`
            :param id: Id (database primary key) of the watch to use
            :param config: Configuration dict to be used (needed for e-mail addresses etc)
            :param on_start: Callable (zero-argument) to be invoked when
                twill script starts execution
            :param on_end: Callable (zero-argument) to be invoked when twill
                script finishes execution (disregarding status)
        """
        twillmanager.async.Worker.__init__(self, queue)
        self.id = id
        self.config = config
        self.watch = None
        self.connection = None
        self.on_start = on_start
        self.on_end = on_end

    def main(self):
        """ Process main function """
        # to make sure we do not use inherited descriptor
        # from the parent process
        close_db_connection()
        self.connection = get_db_connection(self.config)
        self.watch = Watch.load(self.id, self.connection)

        if self.watch:
            logger.info("Starting worker for watch `%s` (id: %s)" % (self.watch.name, self.id))
            self.tick_interval = self.watch.interval
            twillmanager.async.Worker.main(self)
        else:
            logger.warn("Failed to start worker for watch (id: %s) - no such watch" % self.id)

    def tick(self):
        """ Executed every self.watch.interval seconds """
        self.execute()

    def quit(self):
        self.running = False

    def execute(self):
        """ Called by `tick` and when `execute` schedules immediate script execution."""
        try: # large try block to ensure on_end is called
            if self.on_start:
                self.on_start()

            new_status, output = self.execute_script()

            old_status = self.watch.status
            self.watch.status = new_status
            self.watch.time = time.time()
            self.watch.update_status(self.connection)

            msg = "Status for watch `%s` (id: %s): %s" % (self.watch.name, self.id, new_status)

            if new_status != STATUS_OK:
                logger.warn(msg)
            else:
                logger.info(msg)

            # when was last e-mail alert sent
            if self.watch.last_alert is None:
                time_since_last_alert = None
            else:
                time_since_last_alert = self.watch.time - self.watch.last_alert

            status_has_changed = (old_status != new_status)

            # whether last alert was sent long ago enough to send a failure reminder
            # (normally e-mails are sent only on change, but a reminder is sent
            # if the watch keeps failing)
            if self.watch.reminder_interval:
                last_alert_was_long_ago = time_since_last_alert is None or time_since_last_alert > self.watch.reminder_interval
            else:
                last_alert_was_long_ago = False

            if status_has_changed or (last_alert_was_long_ago and new_status == STATUS_FAILED):
                logger.info("Sending notification for watch `%s` (id: %s)" % (self.watch.name, self.id))
                self.status_notify(old_status, new_status, output)
                self.watch.last_alert = time.time()
                self.watch.update_status(self.connection)
        except Exception, e:
            logger.error("Worker `%s` (id: %s) failed with exception: %s" % (self.watch.name, self.id, e.message))
            raise
        finally:
            if self.on_end:
                self.on_end()

    def execute_script(self):
        """ Executes twill script. Returns a tuple status, output """
        out = StringIO()
        # execute the twill, catching any exceptions
        try:
            twill.set_errout(out)
            twill.set_output(out)
            twill.parse._execute_script(self.watch.script.split("\n"))
            status = STATUS_OK
        except Exception, e:
            status = STATUS_FAILED
        finally:
            twill.commands.reset_error()
            twill.commands.reset_output()
        return status, out.getvalue()

    def status_notify(self, old_status, new_status, message):
        """ Sends out notifications about watch status change """
        recipients = self.watch.emails

        if not recipients:
            return

        # strip
        recipients = [r.strip() for r in recipients.split(',')]
        # remove empty addresses
        recipients = [r for r in recipients if r]

        if len(recipients) == 0:
            return

        sender = self.config['mail.from']

        if old_status != new_status:
            subject = "Watch %s status change %s -> %s" % (self.watch.name, old_status, new_status)
        else:
            subject = "Watch %s status is still %s" % (self.watch.name, new_status)

        body = "Script:\n%s\n\nResult:\n%s" % (self.watch.script, message)

        mailer = create_mailer(self.config)
        mailer.send_mail(sender, recipients, subject, body)

class WorkerSet(object):
    """ Object managing spawning of workers."""
    def __init__(self, config):
        """ Create a new `WorkerSet`.

            :param config: Configuration dict to be passed to spawned workers
        """
        # synchronization between threads (WorkerSet is used from CherryPy)
        self._lock = threading.RLock()
        self.workers = {}
        self.now_building = {}
        self.config = config

        # thread that checks for workers that died unexpectedly
        # and listens to their status update messages
        # only tuples (command, argument) should be put into that queue.
        self.manager_thread_queue = multiprocessing.Queue(0)
        self.manager_thread = threading.Thread(target=self.manager_thread_main)
        self.manager_thread.daemon = True
        self.manager_thread.start()

    def worker_status_dict(self, id):
        with self._lock:
            return {'alive': self.is_alive(id), 'building': self.is_building(id)}

    def get(self, id):
        return self.workers.get(id, None)

    def finish(self):
        """ Call this to clean up when the application is shut down """
        self.manager_thread_queue.put(('quit', None))
        self.manager_thread.join()

    def is_alive(self, id):
        """ Check if worker with given id is alive """
        with self._lock:
            return id in self.workers and self.workers[id].is_alive()

    def is_building(self, id):
        """ Check if worker with given id is currently building"""
        with self._lock:
            return self.now_building.get(id, False)

    def check_now(self, id):
        """ Tell the worker with given id to check immediately.
            This also restarts the worker if it died.
        """
        with self._lock:
            self.restart(id)
            self.now_building[id] = True
            self.workers[id].execute()

    def restart(self, id):
        """ Restarts given worker. """
        with self._lock:
            self.remove(id)
            self.add(id)

    def add(self, id):
        """ Adds worker for watch with given id.

            It is safe to call this method even if the worker is already running.
        """
        with self._lock:
            if id in self.workers:
                return

            def on_start():
                self.manager_thread_queue.put(('start', id))

            def on_end():
                self.manager_thread_queue.put(('end', id))

            worker = WorkerProxy(id, self.config)
            worker.set_twill_callbacks(on_start, on_end)

            self.workers[id] = worker
            self.now_building[id] = False
            worker.start(True)

    def remove(self, id):
        """ Removes worker for watch with given id.

            It is safe to call this method even if the worker is not running.
        """
        with self._lock:
            if id in self.workers:
                worker = self.workers[id]
                worker.quit()
                del self.workers[id]
                del self.now_building[id]
                if worker.process:
                    worker.process.join()

    def manager_thread_main(self):
        """ Checks for workers that died unexpectedly and listens to their
            status update messages.
        """
        while True:
            # wait up to 60 seconds
            try:
                command, argument = self.manager_thread_queue.get(True, 60)

                if command == 'quit':
                    break
                elif command == 'start':
                    with self._lock:
                        self.now_building[argument] = True
                elif command == 'end':
                    with self._lock:
                        self.now_building[argument] = False
                else:
                    logger.warn("Unknown command to manager thread: %s" % command)

            except Queue.Empty:
                pass

            # this one is to remove zombie processes
            multiprocessing.active_children()

            with self._lock:
                ids_to_restart = []
                for id, worker in self.workers.items():
                    if not worker.is_alive():
                        ids_to_restart.append(id)

                for id in ids_to_restart:
                    self.restart(id)
<|endoftext|>