<commit_before>import logging
import logging.handlers
import os
import sys
import traceback

from acbs.find import Finder
from acbs import utils
from acbs.utils import ACBSGeneralError, ACBSLogFormatter, ACBSConfError
from acbs.parser import Parser, ACBSPackgeInfo
from acbs.src_fetch import SourceFetcher
from acbs.misc import Misc, Fortune
from acbs.src_process import SourceProcessor
from acbs.loader import LoaderHelper
from acbs.utils import ACBSVariables
from acbs.start_build import Autobuild
from acbs.deps import Dependencies

class BuildCore(object):

    def __init__(self, pkgs_name, debug_mode=False, tree='default', version='', init=True, syslog=False, download_only=False):
        '''
        '''
        self.pkgs_name = pkgs_name
        self.isdebug = debug_mode
        self.tree = tree
        self.pkgs_que = set()
        self.pkgs_done = []
        self.dump_loc = '/var/cache/acbs/tarballs/'
        self.tmp_loc = '/var/cache/acbs/build/'
        self.conf_loc = '/etc/acbs/'
        self.log_loc = '/var/log/acbs/'
        self.pkg_data = ACBSPackgeInfo()
        self.isgroup = False
        self.acbs_version = version
        self.tree_loc = None
        self.log_to_system = syslog
        self.shared_error = None
        self.download_only = download_only
        self.acbs_settings = {'debug_mode': self.isdebug, 'tree': self.tree,
                              'version': self.acbs_version}
        if init:
            self.init()

    def init(self):
        sys.excepthook = self.acbs_except_hdr
        print(utils.full_line_banner(
            'Welcome to ACBS - {}'.format(self.acbs_version)))
        if self.isdebug:
            str_verbosity = logging.DEBUG
        else:
            str_verbosity = logging.INFO
        try:
            for dir_loc in [self.dump_loc, self.tmp_loc, self.conf_loc,
                            self.log_loc]:
                if not os.path.isdir(dir_loc):
                    os.makedirs(dir_loc)
        except Exception:
            raise IOError('\033[93mFailed to make work directory\033[0m!')
        self.__install_logger(str_verbosity)
        Misc().dev_utilz_warn()
        if os.path.exists(os.path.join(self.conf_loc, 'forest.conf')):
            self.tree_loc = os.path.abspath(Parser(
                main_data=self).parse_acbs_conf(self.tree))
            if not self.tree_loc:
                raise ACBSConfError('Tree not found!')
        else:
            Parser(main_data=self).write_acbs_conf()
        if not ACBSVariables.get('pending'):
            ACBSVariables.set('pending', [])

        # @LoaderHelper.register('after_build_finish')
        # def fortune():
        #     Fortune().get_comment()

        LoaderHelper.callback('after_init')

    def __install_logger(self, str_verbosity=logging.INFO,
                         file_verbosity=logging.DEBUG):
        logger = logging.getLogger()
        logger.setLevel(0)  # Set to lowest to bypass the initial filter
        str_handler = logging.StreamHandler()
        str_handler.setLevel(str_verbosity)
        str_handler.setFormatter(ACBSLogFormatter())
        logger.addHandler(str_handler)
        if self.log_to_system:
            log_file_handler = logging.handlers.SysLogHandler(
                address='/dev/log')
        else:
            log_file_handler = logging.handlers.RotatingFileHandler(
                os.path.join(self.log_loc, 'acbs-build.log'), mode='a', maxBytes=2e5, backupCount=10)
        log_file_handler.setLevel(file_verbosity)
        log_file_handler.setFormatter(logging.Formatter(
            '%(asctime)s:%(levelname)s:%(message)s'))
        logger.addHandler(log_file_handler)

    def build(self, pkgs=None):
        LoaderHelper.callback('before_build_init')
        pkgs_to_build = pkgs or self.pkgs_name
        for pkg in pkgs_to_build:
            matched_pkg = Finder(
                pkg, search_path=self.tree_loc).acbs_pkg_match()
            if isinstance(matched_pkg, list):
                logging.info('Package build list found: \033[36m%s (%s)\033[0m' %
                             (os.path.basename(pkg), len(matched_pkg)))
                self.pkgs_name.remove(pkg)
                self.pkgs_name += matched_pkg
                return self.build()
            if not matched_pkg:
                raise ACBSGeneralError(
                    'No valid candidate package found for \033[36m%s\033[0m.' % pkg)
            else:
                self.pkgs_que.update(matched_pkg)
                self.build_single_pkg(matched_pkg)
        print(utils.full_line_banner('Build Summary:'))
        self.print_summary()
        LoaderHelper.callback('after_build_finish')
        return 0

    def build_pkg_group(self, pkgs_array, single_pkg):
        self.isgroup = single_pkg
        import collections
        pkgs_array = collections.OrderedDict(
            sorted(pkgs_array.items(), key=lambda x: x[0]))
        pkg_tuple = [(lambda x, y, z: (y, '%s/0%s-%s' % (z, x, y)) if x < 10 else (y, '%s/%s-%s' %
                                                                                (z, str(x), y)))(i, pkgs_array[i], single_pkg) for i in pkgs_array]
        logging.info('Package group detected\033[36m({})\033[0m: contains: \033[36m{}\033[0m'.format(
            len(pkg_tuple), ' '.join([i[0] for i in pkg_tuple])))
        logging.debug('Package group building order: {}'.format(pkgs_array))
        tmp_dir_loc = []
        self.build_main(single_pkg, tmp_dir_loc, skipbuild=True)
        abbs_data = self.pkg_data.buffer.get('abbs_data')
        for pkg_name, pkg_dir in pkg_tuple:
            print(utils.full_line_banner(''))
            logging.info('Start building \033[36m%s::%s\033[0m' % (
                single_pkg, pkg_name))
            self.pkg_data.buffer['abbs_data'] = abbs_data
            self.build_main(pkg_dir, tmp_dir_loc)

    def print_summary(self):
        i = 0
        group_name = None
        prev_group_name = None
        accum = 0.0

        def swap_vars(prev_group_name):
            ACBSVariables.get('timings').insert(i, accum)
            self.pkgs_done.insert(i, prev_group_name)
            prev_group_name = group_name

        for it in self.pkgs_done:
            if it.find('::') > -1:
                group_name, sub_name = it.split('::')
                if prev_group_name and (group_name != prev_group_name):
                    swap_vars(prev_group_name)
                if group_name == sub_name:
                    self.pkgs_done.remove(it)
                else:
                    accum += ACBSVariables.get('timings')[i]
                i += 1
        if group_name:
            swap_vars(group_name)
        if self.download_only:
            x = [[name, 'Downloaded'] for name in self.pkgs_done]
        else:
            x = [[name, utils.human_time(time)] for name, time in zip(self.pkgs_done, ACBSVariables.get('timings'))]
        print(utils.format_column(x))
        return

    def build_main(self, target, tmp_dir_loc=[], skipbuild=False):
        skipbuild = skipbuild or self.download_only
        if not self.isgroup:
            tmp_dir_loc.clear()
        try:
            pkg_slug = os.path.basename(target)
        except Exception:
            pkg_slug = target
            self.pkg_data.name = pkg_slug
        ACBSVariables.get('pending').append(pkg_slug)
        parser = Parser(
            pkg_name=pkg_slug, spec_file_loc=os.path.abspath(target))
        if not tmp_dir_loc:
            self.pkg_data.update(parser.parse_abbs_spec())
        repo_dir = os.path.abspath(os.path.join(self.tree_loc, target))
        if not skipbuild:
            defines_loc = 'defines' if self.isgroup else 'autobuild/defines'
            self.pkg_data.update(parser.parse_ab3_defines(
                os.path.join(self.tree_loc, target, defines_loc)))
            try_build = Dependencies().process_deps(
                self.pkg_data.build_deps, self.pkg_data.run_deps, pkg_slug)
            if try_build:
                if set(try_build).intersection(ACBSVariables.get('pending')):
                    # Suspect this is dependency loop
                    err_msg = 'Dependency loop: %s' % '<->'.join(
                        ACBSVariables.get('pending'))
                    utils.err_msg(err_msg)
                    raise ACBSGeneralError(err_msg)
                self.new_build_thread(try_build)
        if not tmp_dir_loc:
            src_fetcher = SourceFetcher(
                self.pkg_data.buffer['abbs_data'], self.dump_loc)
            self.pkg_data.src_name = src_fetcher.fetch_src()
            self.pkg_data.src_path = self.dump_loc
            tmp_dir_loc.append(SourceProcessor(self.pkg_data, self).process())
        if self.isgroup:
            repo_ab_dir = repo_dir
        else:
            repo_ab_dir = os.path.join(repo_dir, 'autobuild/')
        if not skipbuild:
            ab3 = Autobuild(tmp_dir_loc[0], repo_ab_dir, self.pkg_data)
            ab3.copy_abd()
            ab3.timed_start_ab3(rm_abdir=self.isgroup)
        if ACBSVariables.get('pending'):
            ACBSVariables.get('pending').pop()
        self.pkgs_done.append(
            target if not self.isgroup else '%s::%s' % (self.isgroup, target))

    def build_single_pkg(self, single_pkg):
        logging.info('Start building \033[36m{}\033[0m'.format(single_pkg))
        os.chdir(self.tree_loc)
        self.pkg_data.slug = single_pkg
        pkg_type_res = Finder.determine_pkg_type(single_pkg)
        if isinstance(pkg_type_res, dict):
            return self.build_pkg_group(pkg_type_res, single_pkg)  # FIXME
        self.isgroup = False
        self.build_main(single_pkg)
        return 0

    def new_build_thread(self, try_build):
        def slave_thread_build(pkg, shared_error):
            logging.debug(
                'New build thread started for \033[36m{}\033[0m'.format(pkg))
            try:
                new_build_instance = BuildCore(
                    **self.acbs_settings, pkgs_name=[pkg], init=False)
                new_build_instance.tree_loc = self.tree_loc
                new_build_instance.shared_error = shared_error
                new_build_instance.build()
            except Exception as ex:
                shared_error.set()
                raise
            return
        from multiprocessing import Process, Event, Lock
        self.shared_error = Event()
        for sub_pkg in list(try_build):
            dumb_mutex = Lock()
            dumb_mutex.acquire()
            sub_thread = Process(target=slave_thread_build, args=(sub_pkg, self.shared_error))
            sub_thread.start()
            sub_thread.join()
            dumb_mutex.release()
            if self.shared_error.is_set():
                raise ACBSGeneralError(
                    'Sub-build process building \033[36m{}\033[0m \033[93mfailed!\033[0m'.format(sub_pkg))

    def acbs_except_hdr(self, type, value, tb):
        logging.debug('Traceback:\n' + ''.join(traceback.format_tb(tb)))
        if self.isdebug:
            sys.__excepthook__(type, value, tb)
        else:
            print()
            logging.fatal('Oops! \033[93m%s\033[0m: \033[93m%s\033[0m' % (
                str(type.__name__), str(value)))
<commit_msg>mainflow: remove "goto" constructs<commit_after>import logging
import logging.handlers
import os
import sys
import traceback

from acbs.find import Finder
from acbs import utils
from acbs.utils import ACBSGeneralError, ACBSLogFormatter, ACBSConfError
from acbs.parser import Parser, ACBSPackgeInfo
from acbs.src_fetch import SourceFetcher
from acbs.misc import Misc, Fortune
from acbs.src_process import SourceProcessor
from acbs.loader import LoaderHelper
from acbs.utils import ACBSVariables
from acbs.start_build import Autobuild
from acbs.deps import Dependencies

class BuildCore(object):

    def __init__(self, pkgs_name, debug_mode=False, tree='default', version='', init=True, syslog=False, download_only=False):
        '''
        '''
        self.pkgs_name = pkgs_name
        self.isdebug = debug_mode
        self.tree = tree
        self.pkgs_que = set()
        self.pkgs_done = []
        self.dump_loc = '/var/cache/acbs/tarballs/'
        self.tmp_loc = '/var/cache/acbs/build/'
        self.conf_loc = '/etc/acbs/'
        self.log_loc = '/var/log/acbs/'
        self.pkg_data = ACBSPackgeInfo()
        self.isgroup = False
        self.acbs_version = version
        self.tree_loc = None
        self.log_to_system = syslog
        self.shared_error = None
        self.download_only = download_only
        self.acbs_settings = {'debug_mode': self.isdebug, 'tree': self.tree,
                              'version': self.acbs_version}
        if init:
            self.init()

    def init(self):
        sys.excepthook = self.acbs_except_hdr
        print(utils.full_line_banner(
            'Welcome to ACBS - {}'.format(self.acbs_version)))
        if self.isdebug:
            str_verbosity = logging.DEBUG
        else:
            str_verbosity = logging.INFO
        try:
            for dir_loc in [self.dump_loc, self.tmp_loc, self.conf_loc,
                            self.log_loc]:
                if not os.path.isdir(dir_loc):
                    os.makedirs(dir_loc)
        except Exception:
            raise IOError('\033[93mFailed to make work directory\033[0m!')
        self.__install_logger(str_verbosity)
        Misc().dev_utilz_warn()
        if os.path.exists(os.path.join(self.conf_loc, 'forest.conf')):
            self.tree_loc = os.path.abspath(Parser(
                main_data=self).parse_acbs_conf(self.tree))
            if not self.tree_loc:
                raise ACBSConfError('Tree not found!')
        else:
            Parser(main_data=self).write_acbs_conf()
        if not ACBSVariables.get('pending'):
            ACBSVariables.set('pending', [])

        # @LoaderHelper.register('after_build_finish')
        # def fortune():
        #     Fortune().get_comment()

        LoaderHelper.callback('after_init')

    def __install_logger(self, str_verbosity=logging.INFO,
                         file_verbosity=logging.DEBUG):
        logger = logging.getLogger()
        logger.setLevel(0)  # Set to lowest to bypass the initial filter
        str_handler = logging.StreamHandler()
        str_handler.setLevel(str_verbosity)
        str_handler.setFormatter(ACBSLogFormatter())
        logger.addHandler(str_handler)
        if self.log_to_system:
            log_file_handler = logging.handlers.SysLogHandler(
                address='/dev/log')
        else:
            log_file_handler = logging.handlers.RotatingFileHandler(
                os.path.join(self.log_loc, 'acbs-build.log'), mode='a', maxBytes=2e5, backupCount=10)
        log_file_handler.setLevel(file_verbosity)
        log_file_handler.setFormatter(logging.Formatter(
            '%(asctime)s:%(levelname)s:%(message)s'))
        logger.addHandler(log_file_handler)

    def build(self):
        LoaderHelper.callback('before_build_init')
        pkgs_to_build = []
        for pkg in self.pkgs_name:
            matched_pkg = Finder(
                pkg, search_path=self.tree_loc).acbs_pkg_match()
            if isinstance(matched_pkg, list):
                logging.info('Package build list found: \033[36m%s (%s)\033[0m' %
                             (os.path.basename(pkg), len(matched_pkg)))
                pkgs_to_build.extend(matched_pkg)
            elif not matched_pkg:
                raise ACBSGeneralError(
                    'No valid candidate package found for \033[36m%s\033[0m.' % pkg)
            else:
                pkgs_to_build.append(matched_pkg)
        for pkg in pkgs_to_build:
            self.pkgs_que.update(pkg)
            self.build_single_pkg(pkg)
        print(utils.full_line_banner('Build Summary:'))
        self.print_summary()
        LoaderHelper.callback('after_build_finish')
        return 0

    def build_pkg_group(self, pkgs_array, single_pkg):
        self.isgroup = single_pkg
        import collections
        pkgs_array = collections.OrderedDict(
            sorted(pkgs_array.items(), key=lambda x: x[0]))
        pkg_tuple = [(lambda x, y, z: (y, '%s/0%s-%s' % (z, x, y)) if x < 10 else (y, '%s/%s-%s' %
                                                                                (z, str(x), y)))(i, pkgs_array[i], single_pkg) for i in pkgs_array]
        logging.info('Package group detected\033[36m({})\033[0m: contains: \033[36m{}\033[0m'.format(
            len(pkg_tuple), ' '.join([i[0] for i in pkg_tuple])))
        logging.debug('Package group building order: {}'.format(pkgs_array))
        tmp_dir_loc = []
        self.build_main(single_pkg, tmp_dir_loc, skipbuild=True)
        abbs_data = self.pkg_data.buffer.get('abbs_data')
        for pkg_name, pkg_dir in pkg_tuple:
            print(utils.full_line_banner(''))
            logging.info('Start building \033[36m%s::%s\033[0m' % (
                single_pkg, pkg_name))
            self.pkg_data.buffer['abbs_data'] = abbs_data
            self.build_main(pkg_dir, tmp_dir_loc)

    def print_summary(self):
        i = 0
        group_name = None
        prev_group_name = None
        accum = 0.0

        def swap_vars(prev_group_name):
            ACBSVariables.get('timings').insert(i, accum)
            self.pkgs_done.insert(i, prev_group_name)
            prev_group_name = group_name

        for it in self.pkgs_done:
            if it.find('::') > -1:
                group_name, sub_name = it.split('::')
                if prev_group_name and (group_name != prev_group_name):
                    swap_vars(prev_group_name)
                if group_name == sub_name:
                    self.pkgs_done.remove(it)
                else:
                    accum += ACBSVariables.get('timings')[i]
                i += 1
        if group_name:
            swap_vars(group_name)
        if self.download_only:
            x = [[name, 'Downloaded'] for name in self.pkgs_done]
        else:
            x = [[name, utils.human_time(time)] for name, time in zip(self.pkgs_done, ACBSVariables.get('timings'))]
        print(utils.format_column(x))
        return

    def build_main(self, target, tmp_dir_loc=[], skipbuild=False):
        skipbuild = skipbuild or self.download_only
        if not self.isgroup:
            tmp_dir_loc.clear()
        try:
            pkg_slug = os.path.basename(target)
        except Exception:
            pkg_slug = target
            self.pkg_data.name = pkg_slug
        ACBSVariables.get('pending').append(pkg_slug)
        parser = Parser(
            pkg_name=pkg_slug, spec_file_loc=os.path.abspath(target))
        if not tmp_dir_loc:
            self.pkg_data.update(parser.parse_abbs_spec())
        repo_dir = os.path.abspath(os.path.join(self.tree_loc, target))
        if not skipbuild:
            defines_loc = 'defines' if self.isgroup else 'autobuild/defines'
            self.pkg_data.update(parser.parse_ab3_defines(
                os.path.join(self.tree_loc, target, defines_loc)))
            try_build = Dependencies().process_deps(
                self.pkg_data.build_deps, self.pkg_data.run_deps, pkg_slug)
            if try_build:
                if set(try_build).intersection(ACBSVariables.get('pending')):
                    # Suspect this is dependency loop
                    err_msg = 'Dependency loop: %s' % '<->'.join(
                        ACBSVariables.get('pending'))
                    utils.err_msg(err_msg)
                    raise ACBSGeneralError(err_msg)
                self.new_build_thread(try_build)
        if not tmp_dir_loc:
            src_fetcher = SourceFetcher(
                self.pkg_data.buffer['abbs_data'], self.dump_loc)
            self.pkg_data.src_name = src_fetcher.fetch_src()
            self.pkg_data.src_path = self.dump_loc
            tmp_dir_loc.append(SourceProcessor(self.pkg_data, self).process())
        if self.isgroup:
            repo_ab_dir = repo_dir
        else:
            repo_ab_dir = os.path.join(repo_dir, 'autobuild/')
        if not skipbuild:
            ab3 = Autobuild(tmp_dir_loc[0], repo_ab_dir, self.pkg_data)
            ab3.copy_abd()
            ab3.timed_start_ab3(rm_abdir=self.isgroup)
        if ACBSVariables.get('pending'):
            ACBSVariables.get('pending').pop()
        self.pkgs_done.append(
            target if not self.isgroup else '%s::%s' % (self.isgroup, target))

    def build_single_pkg(self, single_pkg):
        logging.info('Start building \033[36m{}\033[0m'.format(single_pkg))
        os.chdir(self.tree_loc)
        self.pkg_data.slug = single_pkg
        pkg_type_res = Finder.determine_pkg_type(single_pkg)
        if isinstance(pkg_type_res, dict):
            return self.build_pkg_group(pkg_type_res, single_pkg)  # FIXME
        self.isgroup = False
        self.build_main(single_pkg)
        return 0

    def new_build_thread(self, try_build):
        def slave_thread_build(pkg, shared_error):
            logging.debug(
                'New build thread started for \033[36m{}\033[0m'.format(pkg))
            try:
                new_build_instance = BuildCore(
                    **self.acbs_settings, pkgs_name=[pkg], init=False)
                new_build_instance.tree_loc = self.tree_loc
                new_build_instance.shared_error = shared_error
                new_build_instance.build()
            except Exception as ex:
                shared_error.set()
                raise
            return
        from multiprocessing import Process, Event, Lock
        self.shared_error = Event()
        for sub_pkg in list(try_build):
            dumb_mutex = Lock()
            dumb_mutex.acquire()
            sub_thread = Process(target=slave_thread_build, args=(sub_pkg, self.shared_error))
            sub_thread.start()
            sub_thread.join()
            dumb_mutex.release()
            if self.shared_error.is_set():
                raise ACBSGeneralError(
                    'Sub-build process building \033[36m{}\033[0m \033[93mfailed!\033[0m'.format(sub_pkg))

    def acbs_except_hdr(self, type, value, tb):
        logging.debug('Traceback:\n' + ''.join(traceback.format_tb(tb)))
        if self.isdebug:
            sys.__excepthook__(type, value, tb)
        else:
            print()
            logging.fatal('Oops! \033[93m%s\033[0m: \033[93m%s\033[0m' % (
                str(type.__name__), str(value)))
<|endoftext|>


<commit_before>#!/usr/bin/env python3
import os, argparse, subprocess, json, sys, base64, binascii, time, hashlib, re, copy, textwrap, logging
import dns.resolver, dns.tsigkeyring, dns.update
from configparser import ConfigParser
import urllib.request
from urllib.error import HTTPError

LOGGER = logging.getLogger('acme_dns_tiny')
LOGGER.addHandler(logging.StreamHandler())
LOGGER.setLevel(logging.INFO)

def get_crt(config, log=LOGGER):
    # helper function base64 encode as defined in acme spec
    def _b64(b):
        return base64.urlsafe_b64encode(b).decode("utf8").rstrip("=")

    # helper function to run openssl command
    def _openssl(command, options, communicate=None):
        openssl = subprocess.Popen(["openssl", command] + options,
                                   stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = openssl.communicate(communicate)
        if openssl.returncode != 0:
            raise IOError("OpenSSL Error: {0}".format(err))
        return out

    # helper function to send DNS dynamic update messages
    def _update_dns(rrset, action):
        algorithm = dns.name.from_text("{0}".format(config["TSIGKeyring"]["Algorithm"].lower()))
        dns_update = dns.update.Update(config["DNS"]["zone"], keyring=keyring, keyalgorithm=algorithm)
        if action == "add":
            dns_update.add(rrset.name, rrset)
        elif action == "delete":
            dns_update.delete(rrset.name, rrset)
        resp = dns.query.tcp(dns_update, config["DNS"]["Host"], port=config.getint("DNS", "Port"))
        dns_update = None
        return resp

    # helper function to send signed requests
    def _send_signed_request(url, payload):
        nonlocal jws_nonce
        payload64 = _b64(json.dumps(payload).encode("utf8"))
        protected = copy.deepcopy(jws_header)
        protected["nonce"] = jws_nonce or webclient.open(acme_config["newNonce"]).getheader("Replay-Nonce", None)
        protected["url"] = url
        if url == acme_config["newAccount"]:
            del protected["kid"]
        else:
            del protected["jwk"]
        protected64 = _b64(json.dumps(protected).encode("utf8"))
        signature = _openssl("dgst", ["-sha256", "-sign", config["acmednstiny"]["AccountKeyFile"]],
                             "{0}.{1}".format(protected64, payload64).encode("utf8"))
        data = json.dumps({
            "protected": protected64, "payload": payload64,"signature": _b64(signature)
        })
        try:
            resp = webclient.open(url, data.encode("utf8"))
        except HTTPError as httperror:
            resp = httperror
        finally:
            jws_nonce = resp.getheader("Replay-Nonce", None)
            return resp.getcode(), resp.read(), resp.getheaders()

    # main code
    webclient = urllib.request.build_opener()
    webclient.addheaders = [('User-Agent', 'acme-dns-tiny/2.0'), ('Accept-Language', config["acmednstiny"].get("Language", "en"))]

    log.info("Read ACME directory.")
    directory = webclient.open(config["acmednstiny"]["ACMEDirectory"])
    acme_config = json.loads(directory.read().decode("utf8"))
    terms_service = acme_config.get("meta", {}).get("termsOfService")

    log.info("Prepare DNS keyring and resolver.")
    keyring = dns.tsigkeyring.from_text({config["TSIGKeyring"]["KeyName"]: config["TSIGKeyring"]["KeyValue"]})
    resolver = dns.resolver.Resolver(configure=False)
    resolver.retry_servfail = True
    nameserver = []
    try:
        nameserver = [ipv4_rrset.to_text() for ipv4_rrset in dns.resolver.query(config["DNS"]["Host"], rdtype="A")]
        nameserver = nameserver + [ipv6_rrset.to_text() for ipv6_rrset in dns.resolver.query(config["DNS"]["Host"], rdtype="AAAA")]
    except dns.exception.DNSException as e:
        log.info("A and/or AAAA DNS resources not found for configured dns host: we will use either resource found if exists or directly the DNS Host configuration.")
    if not nameserver:
        nameserver = [config["DNS"]["Host"]]
    resolver.nameservers = nameserver

    log.info("Parsing account key looking for public key.")
    accountkey = _openssl("rsa", ["-in", config["acmednstiny"]["AccountKeyFile"], "-noout", "-text"])
    pub_hex, pub_exp = re.search(
        r"modulus:\r?\n\s+00:([a-f0-9\:\s]+?)\r?\npublicExponent: ([0-9]+)",
        accountkey.decode("utf8"), re.MULTILINE | re.DOTALL).groups()
    pub_exp = "{0:x}".format(int(pub_exp))
    pub_exp = "0{0}".format(pub_exp) if len(pub_exp) % 2 else pub_exp
    jws_header = {
        "alg": "RS256",
        "jwk": {
            "e": _b64(binascii.unhexlify(pub_exp.encode("utf-8"))),
            "kty": "RSA",
            "n": _b64(binascii.unhexlify(re.sub(r"(\s|:)", "", pub_hex).encode("utf-8"))),
        },
        "kid": None,
    }
    accountkey_json = json.dumps(jws_header["jwk"], sort_keys=True, separators=(",", ":"))
    thumbprint = _b64(hashlib.sha256(accountkey_json.encode("utf8")).digest())
    jws_nonce = None

    log.info("Parsing CSR looking for domains.")
    csr = _openssl("req", ["-in", config["acmednstiny"]["CSRFile"], "-noout", "-text"]).decode("utf8")
    domains = set([])
    common_name = re.search(r"Subject:\s*?CN\s*?=\s*?([^\s,;/]+)", csr)
    if common_name is not None:
        domains.add(common_name.group(1))
    subject_alt_names = re.search(r"X509v3 Subject Alternative Name: \r?\n +([^\r\n]+)\r?\n", csr, re.MULTILINE | re.DOTALL)
    if subject_alt_names is not None:
        for san in subject_alt_names.group(1).split(", "):
            if san.startswith("DNS:"):
                domains.add(san[4:])

    log.info("Registering ACME Account.")
    account_request = {}
    account_request["termsOfServiceAgreed"] = True
    account_request["contact"] = config["acmednstiny"].get("Contacts", "").split(';')
    if account_request["contact"] == "":
        del account_request["contact"]

    code, result, headers = _send_signed_request(acme_config["newAccount"], account_request)
    account_info = {}
    if code == 201:
        jws_header["kid"] = dict(headers).get("Location")
        log.info("Registered! (account: '{0}')".format(jws_header["kid"]))
        account_info["termsOfServiceAgreed"] = True
        account_info["contact"] = account_request["contact"]
    elif code == 200:
        jws_header["kid"] = dict(headers).get("Location")
        log.info("Already registered! (account: '{0}')".format(jws_header["kid"]))

        code, result, headers = _send_signed_request(jws_header["kid"], {})
        account_info = json.loads(result.decode("utf8"))
    else:
        raise ValueError("Error registering: {0} {1}".format(code, result))

    log.info("Update contact information if needed.")
    if (set(account_request["contact"]) != set(account_info["contact"])):
        code, result, headers = _send_signed_request(jws_header["kid"], account_request)
        if code == 200:
            log.info("Account updated with latest contact informations.")
        else:
            raise ValueError("Error register update: {0} {1}".format(code, result))

    # new order
    log.info("Certification issuance: ask for a new Order")
    new_order = { "identifiers": [{"type": "dns", "value": domain} for domain in domains]}
    code, result, headers = _send_signed_request(acme_config["newOrder"], new_order)
    order = json.loads(result.decode("utf8"))
    if code == 201:
        order_location = dict(headers).get("Location")
        log.info("Order received: {0}".format(order_location))
        if order["status"] != "pending":
            raise ValueError("Order status is not pending, we can't use it: {0}".format(order))
    elif (code == 403
        and order["type"] == "urn:ietf:params:acme:error:userActionRequired"):
        raise ValueError("Order creation failed ({0}). Read Terms of Service ({1}), then follow your CA instructions: {2}".format(order["detail"], dict(headers)["Link"], order["instance"]))
    else:
        raise ValueError("Error getting new Order: {0} {1}".format(code, result))

    # complete each authorization challenge
    for authz in order["authorizations"]:
        log.info("Completing authz: {0}".format(authz))

        # get new challenge
        resp = webclient.open(authz)
        authorization = json.loads(resp.read().decode("utf8"))
        if resp.getcode() != 200:
            raise ValueError("Error requesting challenges: {0} {1}".format(resp.getcode(), authorization))
        domain = authorization["identifier"]["value"]

        log.info("Create and install DNS TXT challenge resource for: {0}".format(domain))
        challenge = [c for c in authorization["challenges"] if c["type"] == "dns-01"][0]
        token = re.sub(r"[^A-Za-z0-9_\-]", "_", challenge["token"])
        keyauthorization = "{0}.{1}".format(token, thumbprint)
        keydigest64 = _b64(hashlib.sha256(keyauthorization.encode("utf8")).digest())
        dnsrr_domain = "_acme-challenge.{0}.".format(domain)
        dnsrr_set = dns.rrset.from_text(dnsrr_domain, 300, "IN", "TXT",  '"{0}"'.format(keydigest64))
        try:
            _update_dns(dnsrr_set, "add")
        except dns.exception.DNSException as dnsexception:
            raise ValueError("Error updating DNS records: {0} : {1}".format(type(dnsexception).__name__, str(dnsexception)))

        log.info("Waiting for {0} seconds before starting self challenge check.".format(config["acmednstiny"].getint("CheckChallengeDelay")))
        time.sleep(config["acmednstiny"].getint("CheckChallengeDelay"))
        challenge_verified = False
        number_check_fail = 1
        while challenge_verified is False:
            try:
                log.info('Try {0}: Check resource with value "{1}" exits on nameservers: {2}'.format(number_check_fail, keydigest64, resolver.nameservers))
                challenges = resolver.query(dnsrr_domain, rdtype="TXT")
                for response in challenges.rrset:
                    log.info(".. Found value {0}".format(response.to_text()))
                    challenge_verified = challenge_verified or response.to_text() == '"{0}"'.format(keydigest64)
            except dns.exception.DNSException as dnsexception:
                log.info("Info: retry, because a DNS error occurred while checking challenge: {0} : {1}".format(type(dnsexception).__name__, dnsexception))
            finally:
                if challenge_verified is False:
                    if number_check_fail >= 10:
                        raise ValueError("Error checking challenge, value not found: {0}".format(keydigest64))
                    number_check_fail = number_check_fail + 1
                    time.sleep(2)

        log.info("Ask ACME server to perform checks.")
        code, result, headers = _send_signed_request(challenge["url"], {"keyAuthorization": keyauthorization})
        if code != 200:
            raise ValueError("Error triggering challenge: {0} {1}".format(code, result))

        log.info("Waiting challenge to be verified.")
        try:
            while True:
                try:
                    resp = webclient.open(challenge["url"])
                    challenge_status = json.loads(resp.read().decode("utf8"))
                except IOError as e:
                    raise ValueError("Error checking challenge: {0} {1}".format(
                        e.code, json.loads(e.read().decode("utf8"))))
                if challenge_status["status"] == "pending":
                    time.sleep(2)
                elif challenge_status["status"] == "valid":
                    log.info("Domain {0} verified!".format(domain))
                    break
                else:
                    raise ValueError("{0} challenge did not pass: {1}".format(
                        domain, challenge_status))
        finally:
            _update_dns(dnsrr_set, "delete")

    log.info("Finalizing the order...")
    resp = webclient.open(order_location)
    finalize = json.loads(resp.read().decode("utf8"))
    csr_der = _b64(_openssl("req", ["-in", config["acmednstiny"]["CSRFile"], "-outform", "DER"]))
    code, result, headers = _send_signed_request(order["finalize"], {"csr": csr_der})
    if code != 200:
        raise ValueError("Error while sending the CSR: {0} {1}".format(code, result))

    while True:
        try:
            resp = webclient.open(order_location)
            finalize = json.loads(resp.read().decode("utf8"))
        except IOError as e:
            raise ValueError("Error finalizing order: {0} {1}".format(
                e.code, json.loads(e.read().decode("utf8"))))

        if finalize["status"] == "processing":
            time.sleep(resp.getheader("Retry-After", 2))
        elif finalize["status"] == "valid":
            log.info("Order finalized!")
            break
        else:
            raise ValueError("Finalizing order {0} got errors: {1}".format(
                domain, finalize))

    resp = webclient.open(finalize["certificate"])
    if resp.getcode() != 200:
        raise ValueError("Finalizing order {0} got errors: {1}".format(
            resp.getcode(), resp.read.decode("utf8")))
    certchain = resp.read().decode("utf8")

    log.info("Certificate signed and chain received: {0}".format(finalize["certificate"]))
    return certchain

def main(argv):
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""
This script automates the process of getting a signed TLS certificate
chain from Let's Encrypt using the ACME protocol and its DNS verification.
It will need to have access to your private account key and dns server
so PLEASE READ THROUGH IT!
It's around 300 lines, so it won't take long.

===Example Usage===
python3 acme_dns_tiny.py ./example.ini > chain.crt
See example.ini file to configure correctly this script.
===================
"""
    )
    parser.add_argument("--quiet", action="store_const", const=logging.ERROR, help="suppress output except for errors")
    parser.add_argument("configfile", help="path to your configuration file")
    args = parser.parse_args(argv)

    config = ConfigParser()
    config.read_dict({"acmednstiny": {"ACMEDirectory": "https://acme-staging-v02.api.letsencrypt.org/directory",
                                      "CheckChallengeDelay": 3},
                      "DNS": {"Port": "53"}})
    config.read(args.configfile)

    if (set(["accountkeyfile", "csrfile", "acmedirectory", "checkchallengedelay"]) - set(config.options("acmednstiny"))
        or set(["keyname", "keyvalue", "algorithm"]) - set(config.options("TSIGKeyring"))
        or set(["zone", "host", "port"]) - set(config.options("DNS"))):
        raise ValueError("Some required settings are missing.")

    LOGGER.setLevel(args.quiet or LOGGER.level)
    signed_crt = get_crt(config, log=LOGGER)
    sys.stdout.write(signed_crt)

if __name__ == "__main__":  # pragma: no cover
    main(sys.argv[1:])
<commit_msg>Log informations updated<commit_after>#!/usr/bin/env python3
import os, argparse, subprocess, json, sys, base64, binascii, time, hashlib, re, copy, textwrap, logging
import dns.resolver, dns.tsigkeyring, dns.update
from configparser import ConfigParser
import urllib.request
from urllib.error import HTTPError

LOGGER = logging.getLogger('acme_dns_tiny')
LOGGER.addHandler(logging.StreamHandler())
LOGGER.setLevel(logging.INFO)

def get_crt(config, log=LOGGER):
    # helper function base64 encode as defined in acme spec
    def _b64(b):
        return base64.urlsafe_b64encode(b).decode("utf8").rstrip("=")

    # helper function to run openssl command
    def _openssl(command, options, communicate=None):
        openssl = subprocess.Popen(["openssl", command] + options,
                                   stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        out, err = openssl.communicate(communicate)
        if openssl.returncode != 0:
            raise IOError("OpenSSL Error: {0}".format(err))
        return out

    # helper function to send DNS dynamic update messages
    def _update_dns(rrset, action):
        algorithm = dns.name.from_text("{0}".format(config["TSIGKeyring"]["Algorithm"].lower()))
        dns_update = dns.update.Update(config["DNS"]["zone"], keyring=keyring, keyalgorithm=algorithm)
        if action == "add":
            dns_update.add(rrset.name, rrset)
        elif action == "delete":
            dns_update.delete(rrset.name, rrset)
        resp = dns.query.tcp(dns_update, config["DNS"]["Host"], port=config.getint("DNS", "Port"))
        dns_update = None
        return resp

    # helper function to send signed requests
    def _send_signed_request(url, payload):
        nonlocal jws_nonce
        payload64 = _b64(json.dumps(payload).encode("utf8"))
        protected = copy.deepcopy(jws_header)
        protected["nonce"] = jws_nonce or webclient.open(acme_config["newNonce"]).getheader("Replay-Nonce", None)
        protected["url"] = url
        if url == acme_config["newAccount"]:
            del protected["kid"]
        else:
            del protected["jwk"]
        protected64 = _b64(json.dumps(protected).encode("utf8"))
        signature = _openssl("dgst", ["-sha256", "-sign", config["acmednstiny"]["AccountKeyFile"]],
                             "{0}.{1}".format(protected64, payload64).encode("utf8"))
        data = json.dumps({
            "protected": protected64, "payload": payload64,"signature": _b64(signature)
        })
        try:
            resp = webclient.open(url, data.encode("utf8"))
        except HTTPError as httperror:
            resp = httperror
        finally:
            jws_nonce = resp.getheader("Replay-Nonce", None)
            return resp.getcode(), resp.read(), resp.getheaders()

    # main code
    webclient = urllib.request.build_opener()
    webclient.addheaders = [('User-Agent', 'acme-dns-tiny/2.0'), ('Accept-Language', config["acmednstiny"].get("Language", "en"))]

    log.info("Fetch informations from the ACME directory.")
    directory = webclient.open(config["acmednstiny"]["ACMEDirectory"])
    acme_config = json.loads(directory.read().decode("utf8"))
    terms_service = acme_config.get("meta", {}).get("termsOfService")

    log.info("Prepare DNS keyring and resolver.")
    keyring = dns.tsigkeyring.from_text({config["TSIGKeyring"]["KeyName"]: config["TSIGKeyring"]["KeyValue"]})
    resolver = dns.resolver.Resolver(configure=False)
    resolver.retry_servfail = True
    nameserver = []
    try:
        nameserver = [ipv4_rrset.to_text() for ipv4_rrset in dns.resolver.query(config["DNS"]["Host"], rdtype="A")]
        nameserver = nameserver + [ipv6_rrset.to_text() for ipv6_rrset in dns.resolver.query(config["DNS"]["Host"], rdtype="AAAA")]
    except dns.exception.DNSException as e:
        log.info("A and/or AAAA DNS resources not found for configured dns host: we will use either resource found if one exists or directly the DNS Host configuration.")
    if not nameserver:
        nameserver = [config["DNS"]["Host"]]
    resolver.nameservers = nameserver

    log.info("Read account key.")
    accountkey = _openssl("rsa", ["-in", config["acmednstiny"]["AccountKeyFile"], "-noout", "-text"])
    pub_hex, pub_exp = re.search(
        r"modulus:\r?\n\s+00:([a-f0-9\:\s]+?)\r?\npublicExponent: ([0-9]+)",
        accountkey.decode("utf8"), re.MULTILINE | re.DOTALL).groups()
    pub_exp = "{0:x}".format(int(pub_exp))
    pub_exp = "0{0}".format(pub_exp) if len(pub_exp) % 2 else pub_exp
    jws_header = {
        "alg": "RS256",
        "jwk": {
            "e": _b64(binascii.unhexlify(pub_exp.encode("utf-8"))),
            "kty": "RSA",
            "n": _b64(binascii.unhexlify(re.sub(r"(\s|:)", "", pub_hex).encode("utf-8"))),
        },
        "kid": None,
    }
    accountkey_json = json.dumps(jws_header["jwk"], sort_keys=True, separators=(",", ":"))
    thumbprint = _b64(hashlib.sha256(accountkey_json.encode("utf8")).digest())
    jws_nonce = None

    log.info("Read CSR to find domains to validate.")
    csr = _openssl("req", ["-in", config["acmednstiny"]["CSRFile"], "-noout", "-text"]).decode("utf8")
    domains = set([])
    common_name = re.search(r"Subject:\s*?CN\s*?=\s*?([^\s,;/]+)", csr)
    if common_name is not None:
        domains.add(common_name.group(1))
    subject_alt_names = re.search(r"X509v3 Subject Alternative Name: \r?\n +([^\r\n]+)\r?\n", csr, re.MULTILINE | re.DOTALL)
    if subject_alt_names is not None:
        for san in subject_alt_names.group(1).split(", "):
            if san.startswith("DNS:"):
                domains.add(san[4:])

    log.info("Register ACME Account.")
    account_request = {}
    account_request["termsOfServiceAgreed"] = True
    account_request["contact"] = config["acmednstiny"].get("Contacts", "").split(';')
    if account_request["contact"] == "":
        del account_request["contact"]

    code, result, headers = _send_signed_request(acme_config["newAccount"], account_request)
    account_info = {}
    if code == 201:
        jws_header["kid"] = dict(headers).get("Location")
        log.info("Registered! (account: '{0}')".format(jws_header["kid"]))
        account_info["termsOfServiceAgreed"] = True
        account_info["contact"] = account_request["contact"]
    elif code == 200:
        jws_header["kid"] = dict(headers).get("Location")
        log.debug("  - Account is already registered: '{0}'".format(jws_header["kid"]))

        code, result, headers = _send_signed_request(jws_header["kid"], {})
        account_info = json.loads(result.decode("utf8"))
    else:
        raise ValueError("Error registering account: {0} {1}".format(code, result))

    log.info("Update contact information if needed.")
    if (set(account_request["contact"]) != set(account_info["contact"])):
        code, result, headers = _send_signed_request(jws_header["kid"], account_request)
        if code == 200:
            log.debug("  - Account updated with latest contact informations.")
        else:
            raise ValueError("Error registering updates for the account: {0} {1}".format(code, result))

    # new order
    log.info("Request to the ACME server an order to validate domains.")
    new_order = { "identifiers": [{"type": "dns", "value": domain} for domain in domains]}
    code, result, headers = _send_signed_request(acme_config["newOrder"], new_order)
    order = json.loads(result.decode("utf8"))
    if code == 201:
        order_location = dict(headers).get("Location")
        log.debug("  - Order received: {0}".format(order_location))
        if order["status"] != "pending":
            raise ValueError("Order status is not pending, we can't use it: {0}".format(order))
    elif (code == 403
        and order["type"] == "urn:ietf:params:acme:error:userActionRequired"):
        raise ValueError("Order creation failed ({0}). Read Terms of Service ({1}), then follow your CA instructions: {2}".format(order["detail"], dict(headers)["Link"], order["instance"]))
    else:
        raise ValueError("Error getting new Order: {0} {1}".format(code, result))

    # complete each authorization challenge
    for authz in order["authorizations"]:
        log.info("Process challenge for authorization: {0}".format(authz))

        # get new challenge
        resp = webclient.open(authz)
        authorization = json.loads(resp.read().decode("utf8"))
        if resp.getcode() != 200:
            raise ValueError("Error fetching challenges: {0} {1}".format(resp.getcode(), authorization))
        domain = authorization["identifier"]["value"]

        log.info("Install DNS TXT resource for domain: {0}".format(domain))
        challenge = [c for c in authorization["challenges"] if c["type"] == "dns-01"][0]
        token = re.sub(r"[^A-Za-z0-9_\-]", "_", challenge["token"])
        keyauthorization = "{0}.{1}".format(token, thumbprint)
        keydigest64 = _b64(hashlib.sha256(keyauthorization.encode("utf8")).digest())
        dnsrr_domain = "_acme-challenge.{0}.".format(domain)
        dnsrr_set = dns.rrset.from_text(dnsrr_domain, 300, "IN", "TXT",  '"{0}"'.format(keydigest64))
        try:
            _update_dns(dnsrr_set, "add")
        except dns.exception.DNSException as dnsexception:
            raise ValueError("Error updating DNS records: {0} : {1}".format(type(dnsexception).__name__, str(dnsexception)))

        log.info("Waiting for {0} seconds before starting self challenge check.".format(config["acmednstiny"].getint("CheckChallengeDelay")))
        time.sleep(config["acmednstiny"].getint("CheckChallengeDelay"))
        challenge_verified = False
        number_check_fail = 1
        while challenge_verified is False:
            try:
                log.debug('Self test (try: {0}): Check resource with value "{1}" exits on nameservers: {2}'.format(number_check_fail, keydigest64, resolver.nameservers))
                challenges = resolver.query(dnsrr_domain, rdtype="TXT")
                for response in challenges.rrset:
                    log.debug("  - Found value {0}".format(response.to_text()))
                    challenge_verified = challenge_verified or response.to_text() == '"{0}"'.format(keydigest64)
            except dns.exception.DNSException as dnsexception:
                log.debug("  - Will retry as a DNS error occurred while checking challenge: {0} : {1}".format(type(dnsexception).__name__, dnsexception))
            finally:
                if challenge_verified is False:
                    if number_check_fail >= 10:
                        raise ValueError("Error checking challenge, value not found: {0}".format(keydigest64))
                    number_check_fail = number_check_fail + 1
                    time.sleep(2)

        log.info("Ask ACME server to validate thise challenge.")
        code, result, headers = _send_signed_request(challenge["url"], {"keyAuthorization": keyauthorization})
        if code != 200:
            raise ValueError("Error triggering challenge: {0} {1}".format(code, result))
        try:
            while True:
                try:
                    resp = webclient.open(challenge["url"])
                    challenge_status = json.loads(resp.read().decode("utf8"))
                except IOError as e:
                    raise ValueError("Error during challenge validation: {0} {1}".format(
                        e.code, json.loads(e.read().decode("utf8"))))
                if challenge_status["status"] == "pending":
                    time.sleep(2)
                elif challenge_status["status"] == "valid":
                    log.info("ACME has verified challenge for domain: {0}".format(domain))
                    break
                else:
                    raise ValueError("Challenge for domain {0} did not pass: {1}".format(
                        domain, challenge_status))
        finally:
            _update_dns(dnsrr_set, "delete")

    log.info("Request to finalize the order (all chalenge have been completed)")
    resp = webclient.open(order_location)
    finalize = json.loads(resp.read().decode("utf8"))
    csr_der = _b64(_openssl("req", ["-in", config["acmednstiny"]["CSRFile"], "-outform", "DER"]))
    code, result, headers = _send_signed_request(order["finalize"], {"csr": csr_der})
    if code != 200:
        raise ValueError("Error while sending the CSR: {0} {1}".format(code, result))

    while True:
        try:
            resp = webclient.open(order_location)
            finalize = json.loads(resp.read().decode("utf8"))
        except IOError as e:
            raise ValueError("Error finalizing order: {0} {1}".format(
                e.code, json.loads(e.read().decode("utf8"))))

        if finalize["status"] == "processing":
            time.sleep(resp.getheader("Retry-After", 2))
        elif finalize["status"] == "valid":
            log.info("Order finalized!")
            break
        else:
            raise ValueError("Finalizing order {0} got errors: {1}".format(
                domain, finalize))

    resp = webclient.open(finalize["certificate"])
    if resp.getcode() != 200:
        raise ValueError("Finalizing order {0} got errors: {1}".format(
            resp.getcode(), resp.read.decode("utf8")))
    certchain = resp.read().decode("utf8")

    log.info("Certificate signed and chain received: {0}".format(finalize["certificate"]))
    return certchain

def main(argv):
    parser = argparse.ArgumentParser(
        formatter_class=argparse.RawDescriptionHelpFormatter,
        description="""
This script automates the process of getting a signed TLS certificate
chain from any CA using the ACME protocol and its DNS verification.
It will need to have access to your private ACME account key and dns server
so PLEASE READ THROUGH IT!
It's around 300 lines, so it won't take long.

===Example Usage===
python3 acme_dns_tiny.py ./example.ini > chain.crt
See example.ini file to configure correctly this script.
===================
"""
    )
    parser.add_argument("--quiet", action="store_const", const=logging.ERROR, help="suppress output except for errors")
    parser.add_argument("configfile", help="path to your configuration file")
    args = parser.parse_args(argv)

    config = ConfigParser()
    config.read_dict({"acmednstiny": {"ACMEDirectory": "https://acme-staging-v02.api.letsencrypt.org/directory",
                                      "CheckChallengeDelay": 3},
                      "DNS": {"Port": "53"}})
    config.read(args.configfile)

    if (set(["accountkeyfile", "csrfile", "acmedirectory", "checkchallengedelay"]) - set(config.options("acmednstiny"))
        or set(["keyname", "keyvalue", "algorithm"]) - set(config.options("TSIGKeyring"))
        or set(["zone", "host", "port"]) - set(config.options("DNS"))):
        raise ValueError("Some required settings are missing.")

    LOGGER.setLevel(args.quiet or LOGGER.level)
    signed_crt = get_crt(config, log=LOGGER)
    sys.stdout.write(signed_crt)

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


<commit_before>#! /usr/bin/env python2
# vim:fileencoding=UTF-8

import sys
import os
import re
import tweepy
import readline
import pprint
import traceback
import tempfile
import subprocess
import xml.sax.saxutils

DEFAULT_AUTHFILE = os.path.expanduser("~") + "/.twitter_auth"

class TweetShell:
	def __init__(self, authfile):
		self.__authfile = authfile
		a = map(lambda l: re.compile("^(.*?)(//.*)?$").search(l.strip()).group(0).rstrip().split(),
		        open(self.__authfile))
		a = filter(lambda l: len(l) == 5, a)
		self.__auth = {}
		for t in a:
			if t[0] in self.__auth:
				sys.exit("Duplicated screen_name in the authfile.")
			self.__auth[t[0]] = tweepy.OAuthHandler(t[1], t[2])
			self.__auth[t[0]].set_access_token(t[3], t[4])
		self.__commands = {
				"new_auth"    : self.__new_auth,
				"login"       : self.__login,
				"whoami"      : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).me().__dict__),
				"whois"       : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).get_user(screen_name = argv[0]).__dict__),
				"update_stdin": lambda *argv: self.__update(sys.stdin.read(), *argv),
				"update"      : lambda *argv: self.__update_with_editor(*argv),
				"update_photo": lambda *argv: self.__update_with_photo(*argv),
				"favor"       : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).create_favorite(*map(lambda x: int(x), argv)).__dict__),
				"unfavor"     : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_favorite(*map(lambda x: int(x), argv)).__dict__),
				"retweet"     : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).retweet(*map(lambda x: int(x), argv)).__dict__),
				"follow_id"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).create_friendship(*map(lambda x: int(x), argv)).__dict__),
				"follow_sn"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).create_friendship(*argv).__dict__),
				"unfollow_id"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_friendship(*map(lambda x: int(x), argv)).__dict__),
				"unfollow_sn"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_friendship(*argv).__dict__),
				"profile_id"  : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).get_user(*map(lambda x: int(x), argv)).__dict__),
				"profile_sn"  : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).get_user(*argv).__dict__),
				"user_sn"     : lambda sn:    sys.stdout.write("\n".join(map(self.__tl_stringify, reversed(tweepy.API(self.__current_user).user_timeline(count=200, screen_name=sn))))),
				"user_id"     : lambda uid:    sys.stdout.write("\n".join(map(self.__tl_stringify, reversed(tweepy.API(self.__current_user).user_timeline(count=200, id=uid))))),
				"home"        : lambda *argv: self.__timeline(u"home_timeline", *argv),
				"mentions"    : lambda *argv: self.__timeline(u"mentions_timeline", *argv),
				"favorites"   : lambda *argv: self.__timeline(u"favorites", *argv),
				"destroy"     : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_status(*map(lambda x: int(x), argv)).__dict__),
				"shell"       : lambda *argv: subprocess.call(argv),
				"help"        : lambda *args: sys.stdout.writelines(sorted(map(lambda e: e + "\n", self.__commands.keys()))),
			}
		#Synonims
		self.__commands["?"]  = self.__commands["help"]
		# States & Settings
		self.__current_user   = None
		self.__prompt         = "> "
		self.__sn             = ""
		self.__tl_format      = u"""\
@{screen_name} {id} :
{status}
{date} via {source}
"""
		self.__tl_rted_format = u"""\
@{screen_name} (RTed by @{rter_screen_name}) {id} :
{status}
{date} via {source}
"""
	def shell_loop(self):
		try:
			while True:
				try:
					line = raw_input(self.__prompt if sys.stdin.isatty() else "")
					line = line.strip()
					self.eval(line)
				except EOFError:
					raise
				except KeyboardInterrupt:
					sys.stdout.write("KeyboardInterrupt\n")
				except:
					sys.stdout.write(traceback.format_exc(sys.exc_info()[2]))
		except EOFError:
			sys.stdout.write("\n")
	def __update(self, text, *argv):
		_argv = []
		l     = len(argv)
		if l >= 1:
			_argv.append(int(argv[0]))
		sys.stdout.write(self.__tl_stringify(tweepy.API(self.__current_user).update_status(text, *_argv)))
	def __update_with_photo(self, *argv):
		_argv = []
		tmp = tempfile.mktemp(prefix = 'twshell-%s-' % self.__sn)
		edt = os.getenv("EDITOR", "vi")
		subprocess.call([edt, tmp])
		if not os.path.exists(tmp):
			sys.stdout.write("Interrupt\n")
			return
		fp = open(tmp)
		sys.stdout.write(self.__tl_stringify(
			tweepy.API(self.__current_user).update_with_media(argv[0], status = fp.read().rstrip(), in_reply_to_status_id = int(argv[1]))\
				if len(argv) >= 2 else\
				tweepy.API(self.__current_user).update_with_media(argv[0], status = fp.read().rstrip())))
		fp.close()
		os.remove(tmp)
	def __update_with_editor(self, *argv):
		tmp = tempfile.mktemp(prefix = 'twshell-%s-' % self.__sn)
		edt = os.getenv("EDITOR", "vi")
		while True:
			subprocess.call([edt, tmp])
			if not os.path.exists(tmp):
				sys.stdout.write("Interrupt\n")
				return
			fp = open(tmp)
			txt = fp.read()
			print(txt)
			sys.stdout.write("OK? ([y]es/[c]ancel/or else to re-edit)> ")
			answer = sys.stdin.readline().strip()
			if answer == "y":
				self.__update(txt, *argv)
				break
			if answer == "c":
				break
		fp.close()
		os.remove(tmp)
	def __new_auth(self, *arv):
		pair = []
		if self.__auth:
			sys.stdout.write("Please select screen_name which you want to use same consumer as (or EMPTY, new one) > ")
			line = sys.stdin.readline().strip()
			if line:
				if line not in self.__auth:
					sys.stdout.write("You do NOT have @%s in auth dict.\n" % line)
					return
				pair.append(self.__auth[line]._consumer.key)
				pair.append(self.__auth[line]._consumer.secret)
		if not pair:
			sys.stdout.write("Consumer Key   : ")
			pair.append(sys.stdin.readline().rstrip())
			sys.stdout.write("Consumer Secret: ")
			pair.append(sys.stdin.readline().rstrip())
		auth = tweepy.OAuthHandler(*pair)
		sys.stdout.write("Please access: %s\n" % auth.get_authorization_url())
		sys.stdout.write("Input PIN: ")
		auth.get_access_token(sys.stdin.readline().strip())
		self.__auth[auth.get_username()] = auth
		fp = open(self.__authfile, "r")
		org = filter(lambda l: len(re.compile("^(.*?)(//.*)?$").search(l.strip()).group(0).rstrip().split()) != 5, fp)
		fp.close()
		fp = open(self.__authfile, "w")
		fp.writelines(org + ["%s %s %s %s %s\n" % (k, v._consumer.key, v._consumer.secret, v.access_token.key, v.access_token.secret) for k, v in self.__auth.items()])
	def __login(self, *argv):
		term = open(os.ctermid(), "r+w")
		if argv:
			sn = argv[0]
		else:
			for sn in self.__auth.keys():
				term.write("* @%s\n" % sn)
			term.write("WHICH?> ")
			sn = term.readline().strip()
		if sn not in self.__auth.keys():
			term.write("%s was NOT found.\n" % sn)
			return
		self.__current_user = self.__auth[sn]
		self.__prompt = "%s> " % sn
		self.__sn     = sn
	def __tl_stringify(self, status):
		if u"retweeted_status" in status.__dict__:
			return self.__tl_rted_format.format(
					screen_name      = status.retweeted_status.author.screen_name,
					rter_screen_name = status.author.screen_name,
					id               = status.retweeted_status.id,
					status           = xml.sax.saxutils.unescape(status.retweeted_status.text),
					date             = u"%s" % status.retweeted_status.created_at,
					source           = status.retweeted_status.source,
				)
		else:
			return self.__tl_format.format(
					screen_name = status.author.screen_name,
					id          = status.id,
					status      = xml.sax.saxutils.unescape(status.text),
					date        = u"%s" % status.created_at,
					source      = status.source,
				)
	def __timeline(self, tl_type, *argv):
		sys.stdout.write("\n".join(map(self.__tl_stringify, reversed(tweepy.API.__dict__[tl_type](tweepy.API(self.__current_user), count=200, *argv)))))
	def eval(self, *commands):
		for command in map(lambda l: l.strip(), commands):
			m = map(lambda e: e.group("quoted").replace("''", "'") if e.group("quoted") else e.group("literal"),
				re.compile("(?:'(?P<quoted>(?:''|[^'])*)'|(?P<literal>\S+))").finditer(command))
			if not m:
				return
			if (m[0] not in self.__commands):
				sys.stdout.write("Command not found.\n")
				return
			self.__commands[m[0]](*m[1:len(m)])
	@classmethod
	def ask_yn(c, messace = ""):
		fp = open(os.ctermid(), "r+w")
		fp.write("y/n> ")
		return True if fp.readline().rstrip().lower() == "y" else False
<commit_msg>RetweetedのidをRetweeter側に変更<commit_after>#! /usr/bin/env python2
# vim:fileencoding=UTF-8

import sys
import os
import re
import tweepy
import readline
import pprint
import traceback
import tempfile
import subprocess
import xml.sax.saxutils

DEFAULT_AUTHFILE = os.path.expanduser("~") + "/.twitter_auth"

class TweetShell:
	def __init__(self, authfile):
		self.__authfile = authfile
		a = map(lambda l: re.compile("^(.*?)(//.*)?$").search(l.strip()).group(0).rstrip().split(),
		        open(self.__authfile))
		a = filter(lambda l: len(l) == 5, a)
		self.__auth = {}
		for t in a:
			if t[0] in self.__auth:
				sys.exit("Duplicated screen_name in the authfile.")
			self.__auth[t[0]] = tweepy.OAuthHandler(t[1], t[2])
			self.__auth[t[0]].set_access_token(t[3], t[4])
		self.__commands = {
				"new_auth"    : self.__new_auth,
				"login"       : self.__login,
				"whoami"      : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).me().__dict__),
				"whois"       : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).get_user(screen_name = argv[0]).__dict__),
				"update_stdin": lambda *argv: self.__update(sys.stdin.read(), *argv),
				"update"      : lambda *argv: self.__update_with_editor(*argv),
				"update_photo": lambda *argv: self.__update_with_photo(*argv),
				"favor"       : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).create_favorite(*map(lambda x: int(x), argv)).__dict__),
				"unfavor"     : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_favorite(*map(lambda x: int(x), argv)).__dict__),
				"retweet"     : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).retweet(*map(lambda x: int(x), argv)).__dict__),
				"follow_id"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).create_friendship(*map(lambda x: int(x), argv)).__dict__),
				"follow_sn"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).create_friendship(*argv).__dict__),
				"unfollow_id"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_friendship(*map(lambda x: int(x), argv)).__dict__),
				"unfollow_sn"   : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_friendship(*argv).__dict__),
				"profile_id"  : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).get_user(*map(lambda x: int(x), argv)).__dict__),
				"profile_sn"  : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).get_user(*argv).__dict__),
				"user_sn"     : lambda sn:    sys.stdout.write("\n".join(map(self.__tl_stringify, reversed(tweepy.API(self.__current_user).user_timeline(count=200, screen_name=sn))))),
				"user_id"     : lambda uid:    sys.stdout.write("\n".join(map(self.__tl_stringify, reversed(tweepy.API(self.__current_user).user_timeline(count=200, id=uid))))),
				"home"        : lambda *argv: self.__timeline(u"home_timeline", *argv),
				"mentions"    : lambda *argv: self.__timeline(u"mentions_timeline", *argv),
				"favorites"   : lambda *argv: self.__timeline(u"favorites", *argv),
				"destroy"     : lambda *argv: pprint.pprint(tweepy.API(self.__current_user).destroy_status(*map(lambda x: int(x), argv)).__dict__),
				"shell"       : lambda *argv: subprocess.call(argv),
				"help"        : lambda *args: sys.stdout.writelines(sorted(map(lambda e: e + "\n", self.__commands.keys()))),
			}
		#Synonims
		self.__commands["?"]  = self.__commands["help"]
		# States & Settings
		self.__current_user   = None
		self.__prompt         = "> "
		self.__sn             = ""
		self.__tl_format      = u"""\
@{screen_name} {id} :
{status}
{date} via {source}
"""
		self.__tl_rted_format = u"""\
@{screen_name} (RTed by @{rter_screen_name}) {id} :
{status}
{date} via {source}
"""
	def shell_loop(self):
		try:
			while True:
				try:
					line = raw_input(self.__prompt if sys.stdin.isatty() else "")
					line = line.strip()
					self.eval(line)
				except EOFError:
					raise
				except KeyboardInterrupt:
					sys.stdout.write("KeyboardInterrupt\n")
				except:
					sys.stdout.write(traceback.format_exc(sys.exc_info()[2]))
		except EOFError:
			sys.stdout.write("\n")
	def __update(self, text, *argv):
		_argv = []
		l     = len(argv)
		if l >= 1:
			_argv.append(int(argv[0]))
		sys.stdout.write(self.__tl_stringify(tweepy.API(self.__current_user).update_status(text, *_argv)))
	def __update_with_photo(self, *argv):
		_argv = []
		tmp = tempfile.mktemp(prefix = 'twshell-%s-' % self.__sn)
		edt = os.getenv("EDITOR", "vi")
		subprocess.call([edt, tmp])
		if not os.path.exists(tmp):
			sys.stdout.write("Interrupt\n")
			return
		fp = open(tmp)
		sys.stdout.write(self.__tl_stringify(
			tweepy.API(self.__current_user).update_with_media(argv[0], status = fp.read().rstrip(), in_reply_to_status_id = int(argv[1]))\
				if len(argv) >= 2 else\
				tweepy.API(self.__current_user).update_with_media(argv[0], status = fp.read().rstrip())))
		fp.close()
		os.remove(tmp)
	def __update_with_editor(self, *argv):
		tmp = tempfile.mktemp(prefix = 'twshell-%s-' % self.__sn)
		edt = os.getenv("EDITOR", "vi")
		while True:
			subprocess.call([edt, tmp])
			if not os.path.exists(tmp):
				sys.stdout.write("Interrupt\n")
				return
			fp = open(tmp)
			txt = fp.read()
			print(txt)
			sys.stdout.write("OK? ([y]es/[c]ancel/or else to re-edit)> ")
			answer = sys.stdin.readline().strip()
			if answer == "y":
				self.__update(txt, *argv)
				break
			if answer == "c":
				break
		fp.close()
		os.remove(tmp)
	def __new_auth(self, *arv):
		pair = []
		if self.__auth:
			sys.stdout.write("Please select screen_name which you want to use same consumer as (or EMPTY, new one) > ")
			line = sys.stdin.readline().strip()
			if line:
				if line not in self.__auth:
					sys.stdout.write("You do NOT have @%s in auth dict.\n" % line)
					return
				pair.append(self.__auth[line]._consumer.key)
				pair.append(self.__auth[line]._consumer.secret)
		if not pair:
			sys.stdout.write("Consumer Key   : ")
			pair.append(sys.stdin.readline().rstrip())
			sys.stdout.write("Consumer Secret: ")
			pair.append(sys.stdin.readline().rstrip())
		auth = tweepy.OAuthHandler(*pair)
		sys.stdout.write("Please access: %s\n" % auth.get_authorization_url())
		sys.stdout.write("Input PIN: ")
		auth.get_access_token(sys.stdin.readline().strip())
		self.__auth[auth.get_username()] = auth
		fp = open(self.__authfile, "r")
		org = filter(lambda l: len(re.compile("^(.*?)(//.*)?$").search(l.strip()).group(0).rstrip().split()) != 5, fp)
		fp.close()
		fp = open(self.__authfile, "w")
		fp.writelines(org + ["%s %s %s %s %s\n" % (k, v._consumer.key, v._consumer.secret, v.access_token.key, v.access_token.secret) for k, v in self.__auth.items()])
	def __login(self, *argv):
		term = open(os.ctermid(), "r+w")
		if argv:
			sn = argv[0]
		else:
			for sn in self.__auth.keys():
				term.write("* @%s\n" % sn)
			term.write("WHICH?> ")
			sn = term.readline().strip()
		if sn not in self.__auth.keys():
			term.write("%s was NOT found.\n" % sn)
			return
		self.__current_user = self.__auth[sn]
		self.__prompt = "%s> " % sn
		self.__sn     = sn
	def __tl_stringify(self, status):
		if u"retweeted_status" in status.__dict__:
			return self.__tl_rted_format.format(
					screen_name      = status.retweeted_status.author.screen_name,
					rter_screen_name = status.author.screen_name,
					id               = status.id,
					status           = xml.sax.saxutils.unescape(status.retweeted_status.text),
					date             = u"%s" % status.retweeted_status.created_at,
					source           = status.retweeted_status.source,
				)
		else:
			return self.__tl_format.format(
					screen_name = status.author.screen_name,
					id          = status.id,
					status      = xml.sax.saxutils.unescape(status.text),
					date        = u"%s" % status.created_at,
					source      = status.source,
				)
	def __timeline(self, tl_type, *argv):
		sys.stdout.write("\n".join(map(self.__tl_stringify, reversed(tweepy.API.__dict__[tl_type](tweepy.API(self.__current_user), count=200, *argv)))))
	def eval(self, *commands):
		for command in map(lambda l: l.strip(), commands):
			m = map(lambda e: e.group("quoted").replace("''", "'") if e.group("quoted") else e.group("literal"),
				re.compile("(?:'(?P<quoted>(?:''|[^'])*)'|(?P<literal>\S+))").finditer(command))
			if not m:
				return
			if (m[0] not in self.__commands):
				sys.stdout.write("Command not found.\n")
				return
			self.__commands[m[0]](*m[1:len(m)])
	@classmethod
	def ask_yn(c, messace = ""):
		fp = open(os.ctermid(), "r+w")
		fp.write("y/n> ")
		return True if fp.readline().rstrip().lower() == "y" else False
<|endoftext|>


<commit_before>import math

class CodeRing:
    def __init__(self):
        self.encoded = []
        self.password = []
        pass

    def encodeChr(self, character):
        return ord(character)

    def decodeChr(self, integer):
        return chr(integer)

    def encodeLine(self, line):
        encMSG = []
        for index, char in enumerate(line):
            intEncoded = self.encodeChr(char)
            mutator = self.password[index % len(self.password)]
            intEncoded = ((intEncoded * intEncoded) + (mutator * mutator))/89 #89 is prime, makes the number smaller, can mutate some chars due to rounding though. Not an issue in most cases though...
            encMSG.append(intEncoded-97) #removing prime number to lower the integer values below 256
        return encMSG

    def decodeLine(self, tupleEnc):
        buff = []
        if self.password == []:
            return "could not decode wihout password"
        for index, char in enumerate(tupleEnc.split(",")):
            intEncoded = int(char)+97 # adding back
            intEncoded = intEncoded * 89
            mutator = self.password[index % len(self.password)]
            mutator = mutator * mutator
            intDecoded = intEncoded - mutator
            decoded = math.sqrt(intDecoded)
            buff.append(chr(int(decoded)+1)) #undoing the mutations seems to round everything down. Adding one takes care of that.
        return ''.join(buff)

    def setPass(self, password):
        self.password = [self.encodeChr(x) for x in password]
        pass
<commit_msg>added example uses, made easier to use<commit_after>import math

class CodeRing:
    def __init__(self):
        self.encoded = []
        self.password = []
        pass

    def encodeChr(self, character):
        return ord(character)

    def decodeChr(self, integer):
        return chr(integer)

    def encodeLine(self, line):
        encMSG = []
        for index, char in enumerate(line):
            intEncoded = self.encodeChr(char)
            mutator = self.password[index % len(self.password)]
            intEncoded = ((intEncoded * intEncoded) + (mutator * mutator))/89 #89 is prime, makes the number smaller, can mutate some chars due to rounding though. Not an issue in most cases though...
            encMSG.append(intEncoded-97) #removing prime number to lower the integer values below 256
        return encMSG

    def decodeLine(self, tupleEnc):
        if type(tupleEnc) == list:
            print 'found list type'
            tupleEnc = ','.join(str(x) for x in tupleEnc)
        buff = []
        if self.password == []:
            return "could not decode wihout password"
        for index, char in enumerate(tupleEnc.split(",")):
            intEncoded = int(char)+97 # adding back
            intEncoded = intEncoded * 89
            mutator = self.password[index % len(self.password)]
            mutator = mutator * mutator
            intDecoded = intEncoded - mutator
            decoded = math.sqrt(intDecoded)
            buff.append(chr(int(decoded)+1)) #undoing the mutations seems to round everything down. Adding one takes care of that.
        return ''.join(buff)

    def setPass(self, password):
        self.password = [self.encodeChr(x) for x in password]
        pass

if __name__=="__main__":
    c = CodeRing()
    c.setPass("HWre0IXm0Juln6vial4ik4xpqKdV7tOOdymXw7BcohwxWFXABogTc8loMbPfWAs")
    enctuple = c.encodeLine("grapefruits hunted the crazy people furiously quick now")
    decstring = c.decodeLine(enctuple)
    print enctuple
    print decstring

<|endoftext|>


<commit_before><commit_msg>Trying to remove user_modules.<commit_after><|endoftext|>


<commit_before>"""distutils

The main package for the Python Module Distribution Utilities.  Normally
used from a setup script as

   from distutils.core import setup

   setup (...)
"""

# This module should be kept compatible with Python 2.1.

__revision__ = "$Id$"

# Distutils version
#
# Please coordinate with Marc-Andre Lemburg <mal@egenix.com> when adding
# new features to distutils that would warrant bumping the version number.
#
# In general, major and minor version should loosely follow the Python
# version number the distutils code was shipped with.
#

#--start constants--
__version__ = "2.6.4"
#--end constants--
<commit_msg>Bump to 2.6.5 rc 1.<commit_after>"""distutils

The main package for the Python Module Distribution Utilities.  Normally
used from a setup script as

   from distutils.core import setup

   setup (...)
"""

# This module should be kept compatible with Python 2.1.

__revision__ = "$Id$"

# Distutils version
#
# Please coordinate with Marc-Andre Lemburg <mal@egenix.com> when adding
# new features to distutils that would warrant bumping the version number.
#
# In general, major and minor version should loosely follow the Python
# version number the distutils code was shipped with.
#

#--start constants--
__version__ = "2.6.5rc1"
#--end constants--
<|endoftext|>


<commit_before>"""Position reconstruction algorithm using chi square gamma distribution minimization"""

import numpy as np
from scipy.optimize import fmin_powell

from pax import plugin
from pax.datastructure import ReconstructedPosition
from pax.utils import InterpolatingMap, data_file_name

class PosRecChiSquareGamma(plugin.TransformPlugin):

    """Position reconstruction by minimization of chi square gamma function.

    Class to reconstruct S2's x, y and chi square gamma
    (as 'goodness-of-fit' parameter) using the S2 signal in top PMTs, simulated
    LCE map, gains and QE's.

    Additionally calculate chi square gamma for reconstructed positions by other
    algorithms.
    """

    def startup(self):
        """Initialize reconstruction algorithm

        Determine which PMTs to use for reconstruction.
        Load LCE maps, gains and QE's
        """

        # Set mode for algorithm
        self.mode = self.config['mode']
        self.log.debug("Startup, mode: %s" % self.mode)
        if not self.mode == 'no_reconstruct' and not self.mode == 'only_reconstruct' and not self.mode == 'full':
            raise RuntimeError("Bad choice 'mode'")

        # List of integers of which PMTs to use, this algorithm uses the top pmt array to reconstruct
        self.pmts = self.config['channels_top']

        # Set area threshold, minimum area for a peak to be reconstructed
        self.area_threshold = self.config['area_threshold']

        # Set the TPC radius squared
        self.tpc_radius_squared = self.config['tpc_radius']**2

        # (x,y) Locations of these PMTs.  This is stored as a dictionary such
        # that self.pmt_locations[int] = {'x' : int, 'y' : int, 'z' : None}
        self.pmt_locations = self.config['pmt_locations']

        # Load LCE maps
        self.lce_map_file_name = self.config['lce_map_file_name']
        self.maps = InterpolatingMap(data_file_name(self.lce_map_file_name))

        # Load gains (gains from config file, gain error 0.5 pe for all pmts for now)
        self.gains = self.config['gains']
        self.gain_errors = [0.5 for i in range(99)]  # test for now, make config variable

        # Load QE (for now use arbitrary values)
        self.qes = [0.3 for i in range(99)]  # test for now, make config variable
        self.qe_errors = [0.009 for i in range(99)]  # test for now, make config variable

        # Number of pmts (minus dead pmts)
        self.n_channels = 0
        for pmt in self.pmts:
            if not self.gains[pmt] == 0:
                self.n_channels += 1

        # Number of degrees of freedom, n_channels - model degrees of freedom (x,y) - 1
        self.ndf = self.n_channels - 2 - 1

        # Log the total number of calls and success rate to debug, see shutdown
        self.total_rec_calls = 0
        self.total_rec_success = 0

    def function_chi_square_gamma(self, position):
        """Return Chi Square Gamma value for position x,y

        This function is passed to the minimizer, the lowest
        function value will yield the reconstructed position"""

        x = position[0]
        y = position[1]

        # Cutoff value at TPC radius, chi_square_gamma is +infinity here
        if x ** 2 + y ** 2 > self.tpc_radius_squared:
            return float('inf')

        function_value = 0

        # Iterate over all pmts, adding each pmts contribution to function_value
        for pmt in self.pmts:
            # Exclude dead pmts
            if self.gains[pmt] == 0:
                continue

            photons_in_pmt = self.hits[pmt] / self.qes[pmt]
            pmt_error = (self.qe_errors[pmt] / self.qes[pmt]) ** 4 + (self.gain_errors[pmt] / self.gains[pmt]) ** 2

            # Lookup value from LCE map for pmt at position x,y
            map_value = self.maps.get_value(x, y, map_name=str(pmt))

            term_numerator = (photons_in_pmt + min(photons_in_pmt, 1) - self.area_photons * map_value) ** 2
            term_denominator = photons_in_pmt ** 2 * pmt_error + self.area_photons * map_value + 1.0

            function_value += term_numerator / term_denominator

        return function_value

    def transform_event(self, event):
        """Reconstruct the position of S2s in an event.

        For each S2 peak, append a chi_square_gamma value and ndf to existing
        ReconstructedPosition objects.
        Secondly minimize and find the minimum chi_square_gamma and reconstructed
        positions and append a ReconstructedPosition object for those.
        """

        # For every S2 peak found in the event
        for peak in event.S2s():
            # This is an array where every i-th element is how many pe
            # were seen by the i-th PMT
            self.hits = peak.area_per_channel

            # Total number of detected photons in the top array (pe/qe)
            self.area_photons = 0.0

            # Highest number of pe seen in one PMT and the id of this PMT
            max_pmt = 0
            max_pmt_id = 0

            # Calculate which pmt has maximum signal
            for pmt in self.pmts:
                # Exclude dead pmts
                if self.gains[pmt] == 0 or str(self.hits[pmt]) == 'nan' or self.qes[pmt] == 0:
                    continue

                if self.hits[pmt] > max_pmt:
                    max_pmt = self.hits[pmt]
                    max_pmt_id = pmt

                self.area_photons += self.hits[pmt] / self.qes[pmt]

            # Start position for minimizer, if no weighted sum position is present
            # use max pmt location as minimizer start position
            start_x = self.pmt_locations[max_pmt_id]['x']
            start_y = self.pmt_locations[max_pmt_id]['y']

            if not self.mode == 'only_reconstruct':
                # Check for reconstructed positions by other algorithms, if found, calculate chi_square_gamma for those
                for position in peak.reconstructed_positions:
                    if position.algorithm == self.name:
                        continue

                    position.goodness_of_fit = self.function_chi_square_gamma([position.x, position.y])
                    position.ndf = self.ndf
                    self.log.debug("Found reconstructed position by other algorithm"
                                   " %s x: %f y: %f, appending chi_square_gamma: %f ndf: %d"
                                   % (position.algorithm, position.x, position.y, position.goodness_of_fit, self.ndf))

                    # If a neural net position is already calculated for this peak, use it as start position
                    if position.algorithm == 'NeuralNet' and not self.mode == 'no_reconstruct':
                        self.log.debug('Using NeuralNet position as minimizer start position')

                        start_x = position.x
                        start_y = position.y

            if self.mode == 'no_reconstruct':
                return event

            # Only reconstruct peak if it has an area of more then area_threshold pe
            # Setting a very low threshold is bad for speed
            if peak.area < self.area_threshold:
                self.log.debug("Peak area below threshold, skipping to the next peak")
                continue

            # Start minimization
            self.log.debug("Starting minimizer for position reconstruction")

            # Minimize chi_square_gamma function, fmin_powell is the call to the SciPy minimizer
            # It takes the function to minimize, starting position and several options
            # It returns the optimal values for the position (xopt) and function value (fopt)
            # A warnflag tells if the maximum number of iterations was exceeded
            #    warnflag 0, OK
            #    warnflag 1, maximum functions evaluations exceeded
            #    warnflag 2, maximum iterations exceeded
            xopt, fopt, direc, iter, funcalls, warnflag = fmin_powell(self.function_chi_square_gamma,
                                                                      [start_x, start_y],
                                                                      args=(),
                                                                      xtol=0.0001,
                                                                      ftol=0.0001,
                                                                      maxiter=5,
                                                                      maxfun=None,
                                                                      full_output=1,
                                                                      disp=0,
                                                                      direc=None,
                                                                      retall=0)

            self.log.debug("Minimizer warnflag: %d" % warnflag)

            self.total_rec_calls += 1
            if not warnflag:
                self.total_rec_success += 1

            x = float(xopt[0])
            y = float(xopt[1])
            chi_square_gamma = float(fopt)

            # Correct strange error where numpy arrays are passed as fopt
            if isinstance(chi_square_gamma, np.ndarray):
                self.log.warning("Help! Optimizer gave me an ndarray instead of float... "
                                 "To be precise, I got this: %s" % str(chi_square_gamma))
                chi_square_gamma = float('nan')

            # If the minimizer failed do not report a position
            if chi_square_gamma == float('nan'):
                x = y = float('nan')

            self.log.debug("Reconstructed event at x: %f y: %f chi_square_gamma:"
                           " %f ndf: %d" % (x, y, chi_square_gamma, self.ndf))

            # Create a reconstructed position object
            rp = ReconstructedPosition({'x': x,
                                        'y': y,
                                        'goodness_of_fit': chi_square_gamma,
                                        'ndf': self.ndf,
                                        'algorithm': self.name})

            # Append our reconstructed position object to the peak
            peak.reconstructed_positions.append(rp)

        # Return the event such that the next plugin can work on it
        return event

    def shutdown(self):
        if not self.total_rec_calls:
            return

        self.log.debug("Total number of reconstruct calls: %d" % self.total_rec_calls)
        self.log.debug("Success rate: %f" % (self.total_rec_success / self.total_rec_calls))
<commit_msg>Add renormalization of the LCE map values This enables chi_square_gamma to also properly reconstruct below dead PMTs Initial search direction for the minimizer is set to center of the TPC, this avoids edge effects<commit_after>"""Position reconstruction algorithm using chi square gamma distribution minimization"""

import numpy as np
from scipy.optimize import fmin_powell

from pax import plugin
from pax.datastructure import ReconstructedPosition
from pax.utils import InterpolatingMap, data_file_name

class PosRecChiSquareGamma(plugin.TransformPlugin):

    """Position reconstruction by minimization of chi square gamma function.

    Class to reconstruct S2's x, y and chi square gamma
    (as 'goodness-of-fit' parameter) using the S2 signal in top PMTs, simulated
    LCE map, gains and QE's.

    Additionally calculate chi square gamma for reconstructed positions by other
    algorithms.
    """

    def startup(self):
        """Initialize reconstruction algorithm

        Determine which PMTs to use for reconstruction.
        Load LCE maps, gains and QE's
        """

        # Set mode for algorithm
        self.mode = self.config['mode']
        self.log.debug("Startup, mode: %s" % self.mode)
        if not self.mode == 'no_reconstruct' and not self.mode == 'only_reconstruct' and not self.mode == 'full':
            raise RuntimeError("Bad choice 'mode'")

        # List of integers of which PMTs to use, this algorithm uses the top pmt array to reconstruct
        self.pmts = self.config['channels_top']

        # Set area threshold, minimum area for a peak to be reconstructed
        self.area_threshold = self.config['area_threshold']

        # Set the TPC radius squared
        self.tpc_radius_squared = self.config['tpc_radius']**2

        # (x,y) Locations of these PMTs.  This is stored as a dictionary such
        # that self.pmt_locations[int] = {'x' : int, 'y' : int, 'z' : None}
        self.pmt_locations = self.config['pmt_locations']

        # Load LCE maps
        self.lce_map_file_name = self.config['lce_map_file_name']
        self.maps = InterpolatingMap(data_file_name(self.lce_map_file_name))

        # Load gains (gains from config file, gain error 0.5 pe for all pmts for now)
        self.gains = self.config['gains']
        self.gain_errors = [0.5 for i in range(99)]  # test for now, make config variable

        # Load QE (for now use arbitrary values)
        self.qes = [0.3 for i in range(99)]  # test for now, make config variable
        self.qe_errors = [0.009 for i in range(99)]  # test for now, make config variable

        # Number of pmts (minus dead pmts)
        self.n_channels = 0
        for pmt in self.pmts:
            if not self.gains[pmt] == 0:
                self.n_channels += 1

        # Number of degrees of freedom, n_channels - model degrees of freedom (x,y) - 1
        self.ndf = self.n_channels - 2 - 1

        # Log the total number of calls and success rate to debug, see shutdown
        self.total_rec_calls = 0
        self.total_rec_success = 0

    def function_chi_square_gamma(self, position):
        """Return Chi Square Gamma value for position x,y

        This function is passed to the minimizer, the lowest
        function value will yield the reconstructed position"""

        x = position[0]
        y = position[1]

        # Cutoff value at TPC radius, chi_square_gamma is +infinity here
        if x ** 2 + y ** 2 > self.tpc_radius_squared:
            return float('inf')

        function_value = 0

        # Get all LCE map values for the live PMTs at position x,y
        map_values = []
        for pmt in self.pmts:
            if self.gains[pmt] == 0:
                continue

            map_values.append(max(0, self.maps.get_value(x,
                                                         y,
                                                         map_name=str(pmt))))

        # Renormalize the values
        map_values = np.array(map_values) / np.array(map_values).sum()
        index = 0

        # Iterate over all pmts, adding each pmts contribution to function_value
        for pmt in self.pmts:
            # Exclude dead pmts
            if self.gains[pmt] == 0:
                continue

            photons_in_pmt = self.hits[pmt] / self.qes[pmt]
            pmt_error = (self.qe_errors[pmt] / self.qes[pmt]) ** 4 + (self.gain_errors[pmt] / self.gains[pmt]) ** 2

            # Lookup value from LCE map for pmt at position x,y
            map_value = map_values[index]
            index += 1

            term_numerator = (photons_in_pmt + min(photons_in_pmt, 1) - self.area_photons * map_value) ** 2
            term_denominator = photons_in_pmt ** 2 * pmt_error + self.area_photons * map_value + 1.0

            function_value += term_numerator / term_denominator

        return function_value

    def transform_event(self, event):
        """Reconstruct the position of S2s in an event.

        For each S2 peak, append a chi_square_gamma value and ndf to existing
        ReconstructedPosition objects.
        Secondly minimize and find the minimum chi_square_gamma and reconstructed
        positions and append a ReconstructedPosition object for those.
        """

        # For every S2 peak found in the event
        for peak in event.S2s():
            # This is an array where every i-th element is how many pe
            # were seen by the i-th PMT
            self.hits = peak.area_per_channel

            # Total number of detected photons in the top array (pe/qe)
            self.area_photons = 0.0

            # Highest number of pe seen in one PMT and the id of this PMT
            max_pmt = 0
            max_pmt_id = 0

            # Calculate which pmt has maximum signal
            for pmt in self.pmts:
                # Exclude dead pmts
                if self.gains[pmt] == 0 or str(self.hits[pmt]) == 'nan' or self.qes[pmt] == 0:
                    continue

                if self.hits[pmt] > max_pmt:
                    max_pmt = self.hits[pmt]
                    max_pmt_id = pmt

                self.area_photons += self.hits[pmt] / self.qes[pmt]

            # Start position for minimizer, if no weighted sum position is present
            # use max pmt location as minimizer start position
            start_x = self.pmt_locations[max_pmt_id]['x']
            start_y = self.pmt_locations[max_pmt_id]['y']

            if not self.mode == 'only_reconstruct':
                # Check for reconstructed positions by other algorithms, if found, calculate chi_square_gamma for those
                for position in peak.reconstructed_positions:
                    if position.algorithm == self.name:
                        continue

                    position.goodness_of_fit = self.function_chi_square_gamma([position.x, position.y])
                    position.ndf = self.ndf
                    self.log.debug("Found reconstructed position by other algorithm"
                                   " %s x: %f y: %f, appending chi_square_gamma: %f ndf: %d"
                                   % (position.algorithm, position.x, position.y, position.goodness_of_fit, self.ndf))

                    # If a neural net position is already calculated for this peak, use it as start position
                    if position.algorithm == 'NeuralNet' and not self.mode == 'no_reconstruct':
                        self.log.debug('Using NeuralNet position as minimizer start position')

                        start_x = position.x
                        start_y = position.y

            if self.mode == 'no_reconstruct':
                return event

            # Only reconstruct peak if it has an area of more then area_threshold pe
            # Setting a very low threshold is bad for speed
            if peak.area < self.area_threshold:
                self.log.debug("Peak area below threshold, skipping to the next peak")
                continue

            # Set initial search direction of the minimizer to center to avoid edge effects
            direc = np.array([[-np.sign(start_x), 0],
                              [0, -np.sign(start_y)]])

            # Start minimization
            self.log.debug("Starting minimizer for position reconstruction")

            # Minimize chi_square_gamma function, fmin_powell is the call to the SciPy minimizer
            # It takes the function to minimize, starting position and several options
            # It returns the optimal values for the position (xopt) and function value (fopt)
            # A warnflag tells if the maximum number of iterations was exceeded
            #    warnflag 0, OK
            #    warnflag 1, maximum functions evaluations exceeded
            #    warnflag 2, maximum iterations exceeded
            xopt, fopt, direc, iter, funcalls, warnflag = fmin_powell(self.function_chi_square_gamma,
                                                                      [start_x, start_y],
                                                                      args=(),
                                                                      xtol=0.0001,
                                                                      ftol=0.0001,
                                                                      maxiter=5,
                                                                      maxfun=None,
                                                                      full_output=1,
                                                                      disp=0,
                                                                      direc=direc,
                                                                      retall=0)

            self.log.debug("Minimizer warnflag: %d" % warnflag)

            self.total_rec_calls += 1
            if not warnflag:
                self.total_rec_success += 1

            x = float(xopt[0])
            y = float(xopt[1])
            chi_square_gamma = float(fopt)

            # Correct strange error where numpy arrays are passed as fopt
            if isinstance(chi_square_gamma, np.ndarray):
                self.log.warning("Help! Optimizer gave me an ndarray instead of float... "
                                 "To be precise, I got this: %s" % str(chi_square_gamma))
                chi_square_gamma = float('nan')

            # If the minimizer failed do not report a position
            if chi_square_gamma == float('nan'):
                x = y = float('nan')

            self.log.debug("Reconstructed event at x: %f y: %f chi_square_gamma:"
                           " %f ndf: %d" % (x, y, chi_square_gamma, self.ndf))

            # Create a reconstructed position object
            rp = ReconstructedPosition({'x': x,
                                        'y': y,
                                        'goodness_of_fit': chi_square_gamma,
                                        'ndf': self.ndf,
                                        'algorithm': self.name})

            # Append our reconstructed position object to the peak
            peak.reconstructed_positions.append(rp)

        # Return the event such that the next plugin can work on it
        return event

    def shutdown(self):
        if not self.total_rec_calls:
            return

        self.log.debug("Total number of reconstruct calls: %d" % self.total_rec_calls)
        self.log.debug("Success rate: %f" % (self.total_rec_success / self.total_rec_calls))
<|endoftext|>