<commit_before># -*- coding: utf8 -*-
"""
.. module:: burpui.server
    :platform: Unix
    :synopsis: Burp-UI server module.

.. moduleauthor:: Ziirish <hi+burpui@ziirish.me>

"""
import os
import sys
import logging
import traceback

from .misc.auth.handler import UserAuthHandler
from .config import config

from datetime import timedelta
from flask import Flask
from six import iteritems

G_PORT = 5000
G_BIND = u'::'
G_REFRESH = 180
G_LIVEREFRESH = 5
G_SSL = False
G_SINGLE = True
G_SSLCERT = u''
G_SSLKEY = u''
G_VERSION = 2
G_AUTH = [u'basic']
G_ACL = u'none'
G_STORAGE = u''
G_CACHE = u''
G_SESSION = u''
G_REDIS = u''
G_LIMITER = False
G_RATIO = u'20/minute'
G_CELERY = False
G_SCOOKIE = True
G_DEMO = False
G_APPSECRET = u'random'
G_COOKIETIME = 14
G_SESSIONTIME = 5
G_DATABASE = u''
G_PREFIX = u''
G_NO_SERVER_RESTORE = False

class BUIServer(Flask):
    """
    The :class:`burpui.server.BUIServer` class provides the ``Burp-UI`` server.
    """
    gunicorn = False

    defaults = {
        'Global': {
            'port': G_PORT,
            'bind': G_BIND,
            'ssl': G_SSL,
            'standalone': G_SINGLE,
            'single': G_SINGLE,
            'sslcert': G_SSLCERT,
            'sslkey': G_SSLKEY,
            'version': G_VERSION,
            'auth': G_AUTH,
            'acl': G_ACL,
            'prefix': G_PREFIX,
            'demo': G_DEMO,
        },
        'UI': {
            'refresh': G_REFRESH,
            'liverefresh': G_LIVEREFRESH,
        },
        'Security': {
            'scookie': G_SCOOKIE,
            'appsecret': G_APPSECRET,
            'cookietime': G_COOKIETIME,
            'sessiontime': G_SESSIONTIME,
        },
        'Production': {
            'storage': G_STORAGE,
            'session': G_SESSION,
            'cache': G_CACHE,
            'redis': G_REDIS,
            'celery': G_CELERY,
            'database': G_DATABASE,
            'limiter': G_LIMITER,
            'ratio': G_RATIO,
        },
        'Experimental': {
            'noserverrestore': G_NO_SERVER_RESTORE,
        }
    }

    def __init__(self):
        """The :class:`burpui.server.BUIServer` class provides the ``Burp-UI``
        server.

        :param app: The Flask application to launch
        """
        super(BUIServer, self).__init__('burpui')
        self.init = False
        # We cannot override the Flask's logger so we use our own
        self._logger = logging.getLogger('burp-ui')
        self._logger.disabled = True
        self.conf = config
        # switch the flask config with our magic config object
        self.conf.update(self.config)
        self.config = self.conf

    def enable_logger(self, enable=True):
        """Enable or disable the logger"""
        self._logger.disabled = not enable

    @property
    def logger(self):
        return self._logger

    def setup(self, conf=None, unittest=False, cli=False):
        """The :func:`burpui.server.BUIServer.setup` functions is used to setup
        the whole server by parsing the configuration file and loading the
        different backends.

        :param conf: Path to a configuration file
        :type conf: str

        :param unittest: Whether we are unittesting or not (used to avoid burp2
                         strict requirements checks)
        :type unittest: bool

        :param cli: Whether we are in cli mode or not
        :type cli: bool
        """
        self.sslcontext = None
        if not conf:
            conf = self.config['CFG']

        if not conf:
            raise IOError('No configuration file found')

        # Raise exception if errors are encountered during parsing
        self.conf.parse(conf, True, self.defaults)
        self.conf.default_section('Global')

        self.port = self.config['BUI_PORT'] = self.conf.safe_get(
            'port',
            'integer'
        )
        self.demo = self.config['BUI_DEMO'] = self.conf.safe_get(
            'demo',
            'boolean'
        )
        self.bind = self.config['BUI_BIND'] = self.conf.safe_get('bind')
        version = self.conf.safe_get('version', 'integer')
        if unittest and version != 1:
            version = 1
        self.vers = self.config['BUI_VERS'] = version
        self.ssl = self.config['BUI_SSL'] = self.conf.safe_get(
            'ssl',
            'boolean'
        )
        # option standalone has been renamed for less confusion
        key = 'standalone' if 'standalone' in \
            self.conf.conf.get(self.conf.section, {}) else 'single'
        self.standalone = self.config['STANDALONE'] = self.conf.safe_get(
            key,
            'boolean'
        )
        self.sslcert = self.config['BUI_SSLCERT'] = self.conf.safe_get(
            'sslcert'
        )
        self.sslkey = self.config['BUI_SSLKEY'] = self.conf.safe_get(
            'sslkey'
        )
        self.prefix = self.config['BUI_PREFIX'] = self.conf.safe_get(
            'prefix'
        )
        if self.prefix and not self.prefix.startswith('/'):
            if self.prefix.lower() != 'none':
                self.logger.warning("'prefix' must start with a '/'!")
            self.prefix = self.config['BUI_PREFIX'] = ''

        self.auth = self.config['BUI_AUTH'] = self.conf.safe_get(
            'auth',
            'string_lower_list'
        )

        # UI options
        self.config['REFRESH'] = self.conf.safe_get(
            'refresh',
            'integer',
            'UI'
        )
        self.config['LIVEREFRESH'] = self.conf.safe_get(
            'liverefresh',
            'integer',
            'UI'
        )

        # Production options
        self.storage = self.config['BUI_STORAGE'] = self.conf.safe_get(
            'storage',
            section='Production'
        )
        self.cache_db = self.config['BUI_CACHE_DB'] = self.conf.safe_get(
            'cache',
            section='Production'
        )
        self.session_db = self.config['BUI_SESSION_DB'] = self.conf.safe_get(
            'session',
            section='Production'
        )
        self.redis = self.config['BUI_REDIS'] = self.conf.safe_get(
            'redis',
            section='Production'
        )
        self.limiter = self.config['BUI_LIMITER'] = self.conf.safe_get(
            'limiter',
            'boolean_or_string',
            section='Production'
        )
        if isinstance(self.limiter, bool) and not self.limiter:
            self.limiter = self.config['BUI_LIMITER'] = 'none'
        self.ratio = self.config['BUI_RATIO'] = self.conf.safe_get(
            'ratio',
            section='Production'
        )
        self.use_celery = self.config['BUI_CELERY'] = self.conf.safe_get(
            'celery',
            'boolean_or_string',
            section='Production'
        )
        self.database = self.config['SQLALCHEMY_DATABASE_URI'] = \
            self.conf.safe_get(
                'database',
                'boolean_or_string',
                section='Production'
        )
        self.config['WITH_LIMIT'] = False
        if isinstance(self.database, bool):
            self.config['WITH_SQL'] = self.database
        else:
            self.config['WITH_SQL'] = self.database and \
                self.database.lower() != 'none'
        if isinstance(self.use_celery, bool):
            self.config['WITH_CELERY'] = self.use_celery
        else:
            self.config['WITH_CELERY'] = self.use_celery and \
                self.use_celery.lower() != 'none'

        # Experimental options
        self.noserverrestore = self.conf.safe_get(
            'noserverrestore',
            'boolean',
            section='Experimental'
        )

        # Security options
        self.scookie = self.config['BUI_SCOOKIE'] = self.conf.safe_get(
            'scookie',
            'boolean',
            section='Security'
        )
        self.config['SECRET_KEY'] = self.conf.safe_get(
            'appsecret',
            section='Security'
        )
        days = self.conf.safe_get('cookietime', 'integer', section='Security') \
            or G_COOKIETIME
        self.config['REMEMBER_COOKIE_DURATION'] = \
            self.config['PERMANENT_SESSION_LIFETIME'] = timedelta(
                days=days
        )
        self.config['REMEMBER_COOKIE_NAME'] = 'remember_token'
        days = self.conf.safe_get(
            'sessiontime',
            'integer',
            section='Security'
        ) or G_SESSIONTIME
        self.config['SESSION_INACTIVE'] = timedelta(days=days)

        self.logger.info('burp version: {}'.format(self.vers))
        self.logger.info('listen port: {}'.format(self.port))
        self.logger.info('bind addr: {}'.format(self.bind))
        self.logger.info('use ssl: {}'.format(self.ssl))
        self.logger.info('standalone: {}'.format(self.standalone))
        self.logger.info('sslcert: {}'.format(self.sslcert))
        self.logger.info('sslkey: {}'.format(self.sslkey))
        self.logger.info('prefix: {}'.format(self.prefix))
        self.logger.info('secure cookie: {}'.format(self.scookie))
        self.logger.info(
            'cookietime: {}'.format(self.config['REMEMBER_COOKIE_DURATION'])
        )
        self.logger.info(
            'session inactive: {}'.format(self.config['SESSION_INACTIVE'])
        )
        self.logger.info('refresh: {}'.format(self.config['REFRESH']))
        self.logger.info('liverefresh: {}'.format(self.config['LIVEREFRESH']))
        self.logger.info('auth: {}'.format(self.auth))
        self.logger.info('celery: {}'.format(self.use_celery))
        self.logger.info('redis: {}'.format(self.redis))
        self.logger.info('database: {}'.format(self.database))
        self.logger.info('with SQL: {}'.format(self.config['WITH_SQL']))
        self.logger.info('with Celery: {}'.format(self.config['WITH_CELERY']))
        self.logger.info('demo: {}'.format(self.config['BUI_DEMO']))

        self.init = True
        if not cli:
            self.load_modules()

    def load_modules(self, strict=True):
        """Load the extensions"""
        if self.auth and 'none' not in self.auth:
            try:
                self.uhandler = UserAuthHandler(self)
                for back, err in iteritems(self.uhandler.errors):
                    self.logger.critical(
                        'Unable to load \'{}\' authentication backend:\n{}'
                        .format(back, err)
                    )
            except ImportError as e:
                self.logger.critical(
                    'Import Exception, module \'{0}\': {1}'.format(
                        self.auth,
                        str(e)
                    )
                )
                raise e
            self.acl_engine = self.config['BUI_ACL'] = self.conf.safe_get(
                'acl'
            )
            self.config['LOGIN_DISABLED'] = False
        else:
            self.config['LOGIN_DISABLED'] = True
            # No login => no ACL
            self.acl_engine = self.config['BUI_ACL'] = 'none'
            self.auth = self.config['BUI_AUTH'] = 'none'

        if self.acl_engine and self.acl_engine.lower() != 'none':
            try:
                # Try to load submodules from our current environment
                # first
                sys.path.insert(
                    0,
                    os.path.dirname(os.path.abspath(__file__))
                )
                mod = __import__(
                    'burpui.misc.acl.{0}'.format(
                        self.acl_engine.lower()
                    ),
                    fromlist=['ACLloader']
                )
                ACLloader = mod.ACLloader
                self.acl_handler = ACLloader(self)
                # for development purpose only
                # from .misc.acl.interface import BUIacl
                # self.acl = BUIacl
                # self.acl = self.acl_handler.acl
            except Exception as e:
                self.logger.critical(
                    'Import Exception, module \'{0}\': {1}'.format(
                        self.acl_engine,
                        str(e)
                    )
                )
                raise e
        else:
            self.acl_handler = False
            self.acl = False

        self.logger.info('acl: {}'.format(self.acl_engine))

        if self.standalone:
            module = 'burpui.misc.backend.burp{0}'.format(self.vers)
        else:
            module = 'burpui.misc.backend.multi'

        # This is used for development purpose only
        from .misc.backend.burp1 import Burp as BurpGeneric
        self.client = BurpGeneric(dummy=True)
        self.strict = strict
        try:
            # Try to load submodules from our current environment
            # first
            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
            mod = __import__(module, fromlist=['Burp'])
            Client = mod.Burp
            self.client = Client(self, conf=self.conf)
        except Exception as e:
            self.logger.critical(traceback.format_exc())
            self.logger.critical(
                'Failed loading backend for Burp version {0}: {1}'.format(
                    self.vers,
                    str(e)
                )
            )
            sys.exit(2)

    @property
    def acl(self):
        """ACL module

        :returns: :class:`burpui.misc.acl.interface.BUIacl`
        """
        if self.acl_engine and self.acl_engine.lower() != 'none':
            # refresh acl to detect config changes
            from .misc.acl.interface import BUIacl  # noqa
            acl = self.acl_handler.acl  # type: BUIacl
            return acl
        return None

    def get_send_file_max_age(self, name):
        """Provides default cache_timeout for the send_file() functions."""
        if name:
            lname = name.lower()
            extensions = ['js', 'css', 'woff']
            for ext in extensions:
                if lname.endswith('.{}'.format(ext)):
                    return 3600 * 24 * 30  # 30 days
        return Flask.get_send_file_max_age(self, name)

    def manual_run(self):
        """The :func:`burpui.server.BUIServer.manual_run` functions is used to
        actually launch the ``Burp-UI`` server.

        .. deprecated:: 0.4.0
           You should now run the Flask's run command.
        """
        if not self.init:
            self.setup()

        if self.ssl:
            self.sslcontext = (self.sslcert, self.sslkey)

        if self.sslcontext:
            self.config['SSL'] = True
            self.run(
                host=self.bind,
                port=self.port,
                debug=self.config['DEBUG'],
                ssl_context=self.sslcontext
            )
        else:
            self.run(host=self.bind, port=self.port, debug=self.config['DEBUG'])
<commit_msg>fix: attribute changed to property without a setter<commit_after># -*- coding: utf8 -*-
"""
.. module:: burpui.server
    :platform: Unix
    :synopsis: Burp-UI server module.

.. moduleauthor:: Ziirish <hi+burpui@ziirish.me>

"""
import os
import sys
import logging
import traceback

from .misc.auth.handler import UserAuthHandler
from .config import config

from datetime import timedelta
from flask import Flask
from six import iteritems

G_PORT = 5000
G_BIND = u'::'
G_REFRESH = 180
G_LIVEREFRESH = 5
G_SSL = False
G_SINGLE = True
G_SSLCERT = u''
G_SSLKEY = u''
G_VERSION = 2
G_AUTH = [u'basic']
G_ACL = u'none'
G_STORAGE = u''
G_CACHE = u''
G_SESSION = u''
G_REDIS = u''
G_LIMITER = False
G_RATIO = u'20/minute'
G_CELERY = False
G_SCOOKIE = True
G_DEMO = False
G_APPSECRET = u'random'
G_COOKIETIME = 14
G_SESSIONTIME = 5
G_DATABASE = u''
G_PREFIX = u''
G_NO_SERVER_RESTORE = False

class BUIServer(Flask):
    """
    The :class:`burpui.server.BUIServer` class provides the ``Burp-UI`` server.
    """
    gunicorn = False

    defaults = {
        'Global': {
            'port': G_PORT,
            'bind': G_BIND,
            'ssl': G_SSL,
            'standalone': G_SINGLE,
            'single': G_SINGLE,
            'sslcert': G_SSLCERT,
            'sslkey': G_SSLKEY,
            'version': G_VERSION,
            'auth': G_AUTH,
            'acl': G_ACL,
            'prefix': G_PREFIX,
            'demo': G_DEMO,
        },
        'UI': {
            'refresh': G_REFRESH,
            'liverefresh': G_LIVEREFRESH,
        },
        'Security': {
            'scookie': G_SCOOKIE,
            'appsecret': G_APPSECRET,
            'cookietime': G_COOKIETIME,
            'sessiontime': G_SESSIONTIME,
        },
        'Production': {
            'storage': G_STORAGE,
            'session': G_SESSION,
            'cache': G_CACHE,
            'redis': G_REDIS,
            'celery': G_CELERY,
            'database': G_DATABASE,
            'limiter': G_LIMITER,
            'ratio': G_RATIO,
        },
        'Experimental': {
            'noserverrestore': G_NO_SERVER_RESTORE,
        }
    }

    def __init__(self):
        """The :class:`burpui.server.BUIServer` class provides the ``Burp-UI``
        server.

        :param app: The Flask application to launch
        """
        super(BUIServer, self).__init__('burpui')
        self.init = False
        # We cannot override the Flask's logger so we use our own
        self._logger = logging.getLogger('burp-ui')
        self._logger.disabled = True
        self.conf = config
        # switch the flask config with our magic config object
        self.conf.update(self.config)
        self.config = self.conf

    def enable_logger(self, enable=True):
        """Enable or disable the logger"""
        self._logger.disabled = not enable

    @property
    def logger(self):
        return self._logger

    def setup(self, conf=None, unittest=False, cli=False):
        """The :func:`burpui.server.BUIServer.setup` functions is used to setup
        the whole server by parsing the configuration file and loading the
        different backends.

        :param conf: Path to a configuration file
        :type conf: str

        :param unittest: Whether we are unittesting or not (used to avoid burp2
                         strict requirements checks)
        :type unittest: bool

        :param cli: Whether we are in cli mode or not
        :type cli: bool
        """
        self.sslcontext = None
        if not conf:
            conf = self.config['CFG']

        if not conf:
            raise IOError('No configuration file found')

        # Raise exception if errors are encountered during parsing
        self.conf.parse(conf, True, self.defaults)
        self.conf.default_section('Global')

        self.port = self.config['BUI_PORT'] = self.conf.safe_get(
            'port',
            'integer'
        )
        self.demo = self.config['BUI_DEMO'] = self.conf.safe_get(
            'demo',
            'boolean'
        )
        self.bind = self.config['BUI_BIND'] = self.conf.safe_get('bind')
        version = self.conf.safe_get('version', 'integer')
        if unittest and version != 1:
            version = 1
        self.vers = self.config['BUI_VERS'] = version
        self.ssl = self.config['BUI_SSL'] = self.conf.safe_get(
            'ssl',
            'boolean'
        )
        # option standalone has been renamed for less confusion
        key = 'standalone' if 'standalone' in \
            self.conf.conf.get(self.conf.section, {}) else 'single'
        self.standalone = self.config['STANDALONE'] = self.conf.safe_get(
            key,
            'boolean'
        )
        self.sslcert = self.config['BUI_SSLCERT'] = self.conf.safe_get(
            'sslcert'
        )
        self.sslkey = self.config['BUI_SSLKEY'] = self.conf.safe_get(
            'sslkey'
        )
        self.prefix = self.config['BUI_PREFIX'] = self.conf.safe_get(
            'prefix'
        )
        if self.prefix and not self.prefix.startswith('/'):
            if self.prefix.lower() != 'none':
                self.logger.warning("'prefix' must start with a '/'!")
            self.prefix = self.config['BUI_PREFIX'] = ''

        self.auth = self.config['BUI_AUTH'] = self.conf.safe_get(
            'auth',
            'string_lower_list'
        )

        # UI options
        self.config['REFRESH'] = self.conf.safe_get(
            'refresh',
            'integer',
            'UI'
        )
        self.config['LIVEREFRESH'] = self.conf.safe_get(
            'liverefresh',
            'integer',
            'UI'
        )

        # Production options
        self.storage = self.config['BUI_STORAGE'] = self.conf.safe_get(
            'storage',
            section='Production'
        )
        self.cache_db = self.config['BUI_CACHE_DB'] = self.conf.safe_get(
            'cache',
            section='Production'
        )
        self.session_db = self.config['BUI_SESSION_DB'] = self.conf.safe_get(
            'session',
            section='Production'
        )
        self.redis = self.config['BUI_REDIS'] = self.conf.safe_get(
            'redis',
            section='Production'
        )
        self.limiter = self.config['BUI_LIMITER'] = self.conf.safe_get(
            'limiter',
            'boolean_or_string',
            section='Production'
        )
        if isinstance(self.limiter, bool) and not self.limiter:
            self.limiter = self.config['BUI_LIMITER'] = 'none'
        self.ratio = self.config['BUI_RATIO'] = self.conf.safe_get(
            'ratio',
            section='Production'
        )
        self.use_celery = self.config['BUI_CELERY'] = self.conf.safe_get(
            'celery',
            'boolean_or_string',
            section='Production'
        )
        self.database = self.config['SQLALCHEMY_DATABASE_URI'] = \
            self.conf.safe_get(
                'database',
                'boolean_or_string',
                section='Production'
        )
        self.config['WITH_LIMIT'] = False
        if isinstance(self.database, bool):
            self.config['WITH_SQL'] = self.database
        else:
            self.config['WITH_SQL'] = self.database and \
                self.database.lower() != 'none'
        if isinstance(self.use_celery, bool):
            self.config['WITH_CELERY'] = self.use_celery
        else:
            self.config['WITH_CELERY'] = self.use_celery and \
                self.use_celery.lower() != 'none'

        # Experimental options
        self.noserverrestore = self.conf.safe_get(
            'noserverrestore',
            'boolean',
            section='Experimental'
        )

        # Security options
        self.scookie = self.config['BUI_SCOOKIE'] = self.conf.safe_get(
            'scookie',
            'boolean',
            section='Security'
        )
        self.config['SECRET_KEY'] = self.conf.safe_get(
            'appsecret',
            section='Security'
        )
        days = self.conf.safe_get('cookietime', 'integer', section='Security') \
            or G_COOKIETIME
        self.config['REMEMBER_COOKIE_DURATION'] = \
            self.config['PERMANENT_SESSION_LIFETIME'] = timedelta(
                days=days
        )
        self.config['REMEMBER_COOKIE_NAME'] = 'remember_token'
        days = self.conf.safe_get(
            'sessiontime',
            'integer',
            section='Security'
        ) or G_SESSIONTIME
        self.config['SESSION_INACTIVE'] = timedelta(days=days)

        self.logger.info('burp version: {}'.format(self.vers))
        self.logger.info('listen port: {}'.format(self.port))
        self.logger.info('bind addr: {}'.format(self.bind))
        self.logger.info('use ssl: {}'.format(self.ssl))
        self.logger.info('standalone: {}'.format(self.standalone))
        self.logger.info('sslcert: {}'.format(self.sslcert))
        self.logger.info('sslkey: {}'.format(self.sslkey))
        self.logger.info('prefix: {}'.format(self.prefix))
        self.logger.info('secure cookie: {}'.format(self.scookie))
        self.logger.info(
            'cookietime: {}'.format(self.config['REMEMBER_COOKIE_DURATION'])
        )
        self.logger.info(
            'session inactive: {}'.format(self.config['SESSION_INACTIVE'])
        )
        self.logger.info('refresh: {}'.format(self.config['REFRESH']))
        self.logger.info('liverefresh: {}'.format(self.config['LIVEREFRESH']))
        self.logger.info('auth: {}'.format(self.auth))
        self.logger.info('celery: {}'.format(self.use_celery))
        self.logger.info('redis: {}'.format(self.redis))
        self.logger.info('database: {}'.format(self.database))
        self.logger.info('with SQL: {}'.format(self.config['WITH_SQL']))
        self.logger.info('with Celery: {}'.format(self.config['WITH_CELERY']))
        self.logger.info('demo: {}'.format(self.config['BUI_DEMO']))

        self.init = True
        if not cli:
            self.load_modules()

    def load_modules(self, strict=True):
        """Load the extensions"""
        if self.auth and 'none' not in self.auth:
            try:
                self.uhandler = UserAuthHandler(self)
                for back, err in iteritems(self.uhandler.errors):
                    self.logger.critical(
                        'Unable to load \'{}\' authentication backend:\n{}'
                        .format(back, err)
                    )
            except ImportError as e:
                self.logger.critical(
                    'Import Exception, module \'{0}\': {1}'.format(
                        self.auth,
                        str(e)
                    )
                )
                raise e
            self.acl_engine = self.config['BUI_ACL'] = self.conf.safe_get(
                'acl'
            )
            self.config['LOGIN_DISABLED'] = False
        else:
            self.config['LOGIN_DISABLED'] = True
            # No login => no ACL
            self.acl_engine = self.config['BUI_ACL'] = 'none'
            self.auth = self.config['BUI_AUTH'] = 'none'

        if self.acl_engine and self.acl_engine.lower() != 'none':
            try:
                # Try to load submodules from our current environment
                # first
                sys.path.insert(
                    0,
                    os.path.dirname(os.path.abspath(__file__))
                )
                mod = __import__(
                    'burpui.misc.acl.{0}'.format(
                        self.acl_engine.lower()
                    ),
                    fromlist=['ACLloader']
                )
                ACLloader = mod.ACLloader
                self.acl_handler = ACLloader(self)
            except Exception as e:
                self.logger.critical(
                    'Import Exception, module \'{0}\': {1}'.format(
                        self.acl_engine,
                        str(e)
                    )
                )
                raise e
        else:
            self.acl_handler = False

        self.logger.info('acl: {}'.format(self.acl_engine))

        if self.standalone:
            module = 'burpui.misc.backend.burp{0}'.format(self.vers)
        else:
            module = 'burpui.misc.backend.multi'

        # This is used for development purpose only
        from .misc.backend.burp1 import Burp as BurpGeneric
        self.client = BurpGeneric(dummy=True)
        self.strict = strict
        try:
            # Try to load submodules from our current environment
            # first
            sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
            mod = __import__(module, fromlist=['Burp'])
            Client = mod.Burp
            self.client = Client(self, conf=self.conf)
        except Exception as e:
            self.logger.critical(traceback.format_exc())
            self.logger.critical(
                'Failed loading backend for Burp version {0}: {1}'.format(
                    self.vers,
                    str(e)
                )
            )
            sys.exit(2)

    @property
    def acl(self):
        """ACL module

        :returns: :class:`burpui.misc.acl.interface.BUIacl`
        """
        if self.acl_engine and self.acl_engine.lower() != 'none':
            # refresh acl to detect config changes
            from .misc.acl.interface import BUIacl  # noqa
            acl = self.acl_handler.acl  # type: BUIacl
            return acl
        return None

    def get_send_file_max_age(self, name):
        """Provides default cache_timeout for the send_file() functions."""
        if name:
            lname = name.lower()
            extensions = ['js', 'css', 'woff']
            for ext in extensions:
                if lname.endswith('.{}'.format(ext)):
                    return 3600 * 24 * 30  # 30 days
        return Flask.get_send_file_max_age(self, name)

    def manual_run(self):
        """The :func:`burpui.server.BUIServer.manual_run` functions is used to
        actually launch the ``Burp-UI`` server.

        .. deprecated:: 0.4.0
           You should now run the Flask's run command.
        """
        if not self.init:
            self.setup()

        if self.ssl:
            self.sslcontext = (self.sslcert, self.sslkey)

        if self.sslcontext:
            self.config['SSL'] = True
            self.run(
                host=self.bind,
                port=self.port,
                debug=self.config['DEBUG'],
                ssl_context=self.sslcontext
            )
        else:
            self.run(host=self.bind, port=self.port, debug=self.config['DEBUG'])
<|endoftext|>


<commit_before>import requests, iso8601, os, json, logging

from flask import Flask, jsonify, json, request, redirect, abort, make_response
from flask import render_template
import urllib2
import base64
import time
import urlparse

from totalimpactwebapp import app, util
from totalimpactwebapp.models import Github
from totalimpactwebapp import pretty_date

logger = logging.getLogger("tiwebapp.views")
mymixpanel = mixpanel.Mixpanel(token=os.getenv("MIXPANEL_TOKEN"))
mixpanel_token = os.getenv("MIXPANEL_TOKEN")

@app.before_request
def log_ip_address():
    if request.endpoint != "static":
        ip_address = request.remote_addr
        logger.info("%30s IP address calling %s %s" % (ip_address, request.method, request.url))

@app.after_request
def add_crossdomain_header(resp):
    resp.headers['Access-Control-Allow-Origin'] = "*"
    resp.headers['Access-Control-Allow-Methods'] = "POST, GET, OPTIONS, PUT, DELETE"
    resp.headers['Access-Control-Allow-Headers'] = "Content-Type"

    return resp

# static pages
@app.route('/')
def home():
    return render_template(
    'index.html',
    page_title="tell the full story of your research impact",
    body_class="homepage",
    mixpanel_token=os.environ["MIXPANEL_TOKEN"],
    api_root=os.environ["API_ROOT"]
    )

@app.route("/embed/impactstory.js")
@app.route("/embed/v1/impactstory.js")
def impactstory_dot_js():

    badges_template = render_template("js-template-badges.html").replace("\n", "")
    ich_def = open(os.path.dirname(__file__) + "/static/js/icanhaz.min.js", "r").read()
    ti_item_def = open(os.path.dirname(__file__) + "/static/js/ti-item.js", "r").read()

    rendered = render_template(
        "impactstory.js",
        ich=ich_def,
        ti_item=ti_item_def,
        badges_template=badges_template,
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        webapp_root=os.environ["WEBAPP_ROOT"]
    )
    resp = make_response(rendered)
    """
    There is no standard way to indicate you're sending back javascript;
    This seems the most recommended one, though. See
    http://stackoverflow.com/questions/2706290/why-write-script-type-text-javascript-when-the-mime-type-is-set-by-the-serve
    and http://www.ietf.org/rfc/rfc4329.txt
     """
    resp.headers["Content-Type"] = "application/javascript; charset=utf-8"
    return resp

@app.route("/embed/test")
def embed_test():
    return render_template(
        "sample-embed-internal-test.html",
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        webapp_root = os.environ["WEBAPP_ROOT"]
    )

@app.route('/about')
def about():
    return render_template(
        'about.html',
        page_title="about",
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"]
        )

@app.route('/faq')
def faq():
    # get the table of items and identifiers
    which_items_loc = os.path.join(
        os.path.dirname(__file__),
        "static",
        "whichartifacts.html"
        )
    which_item_types = open(which_items_loc).read()

    # get the static_meta info for each metric
    try:
        r = requests.get("http://" + os.environ["API_ROOT"] +'/provider')
        metadata = json.loads(r.text)
    except requests.ConnectionError:
        metadata = {}

    return render_template(
        'faq.html',
        page_title="faq",
        which_artifacts=which_item_types,
        provider_metadata=metadata,
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"]
        )

@app.route('/api-docs')
def apidocs():
    return render_template(
        'api-docs.html',
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        webapp_root = os.environ["WEBAPP_ROOT"],
        page_title="api & widget"
        )

@app.route('/pricing')
def pricing():
    return render_template(
        'pricing.html',
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        webapp_root = os.environ["WEBAPP_ROOT"],
        page_title="pricing"
        )

@app.route('/create')
def collection_create():
    return render_template(
        'create-collection.html',
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        page_title="create collection",
        body_class="create-collection"
        )

@app.route('/collection/<collection_id>')
def collection_report(collection_id):
    url = "http://{api_root}/collection/{collection_id}?include_items=0".format(
        api_root=os.getenv("API_ROOT"),
        collection_id=collection_id
    )

    r = requests.get(url)
    if r.status_code == 200:
        collection = json.loads(r.text)
        return render_template(
            'report.html',
            mixpanel_token=os.environ["MIXPANEL_TOKEN"],
            api_root=os.environ["API_ROOT"],
            api_key=os.environ["API_KEY"],
            request_url=request.url,
            page_title=collection["title"],
            body_class="report",
            report_id=collection["_id"],
            report_id_namespace="impactstory_collection_id",
            api_query="collection/"+collection["_id"]
        )
    else:
        abort(404, "This collection doesn't seem to exist yet. "+url)

@app.route('/item/<ns>/<path:id>')
def item_report(ns, id):
    url = "http://{api_root}/v1/item/{ns}/{id}?key={api_key}".format(
        api_root=os.getenv("API_ROOT"),
        ns=ns,
        id=id,
        api_key=os.environ["API_KEY"]
    )

    r = requests.get(url)
    if r.status_code <= 210: # allow unfinished items
        item = json.loads(r.text)
        return render_template(
            'report.html',
            mixpanel_token=os.environ["MIXPANEL_TOKEN"],
            api_root=os.environ["API_ROOT"],
            api_key=os.environ["API_KEY"],
            request_url=request.url,
            page_title="",
            body_class="report",
            report_id=id,
            report_id_namespace=ns,
            api_query="item/{ns}/{id}".format(ns=ns, id=id)
        )
    else:
        abort(404, "This item doesn't seem to exist yet. "+url)

@app.route('/vitals', methods=["POST"])
def vitals():
    """
    Logs reporting stats from the embed code to mixpanel

    Gets a "vitals" object that has a url and a list of "params" objects.
    Each widget on a page sends its own params object; in practice, these will
    almost certainly be identical since users aren't likely to set different
    options for widgets on the same page. It's a good way to count widgets on
    that page, though.

    For documentation on the keys in the the params object, see the default
    params listed at the head of impactstory.js's main() function, and also the
    api-docs page.
    """

    vitals = request.json
    properties = {
                    'token': mixpanel_token,
                    'time': int(time.time()),
                    'ip': request.remote_addr,
                    "$referring_domain": urlparse.urlsplit(request.referrer).netloc,
                    "$referrer" : request.referrer,
                    "$os": request.user_agent.platform,
                    "$browser": request.user_agent.browser,
                    "Embeds per page": len(vitals["allParams"])
                    "Host page": vitals["url"],
                    "API Key": vitals["allParams"][0]["api-key"],
                }

    for embed_param in vitals["allParams"][0]:
        properties["Embed:"+embed_param] = vitals["allParams"][0][embed_param]

    #logger.debug("properties are {properties}".format(
    #    properties=properties))

    mixpanel_params = {"event": "Impression:embedpage", "properties": properties}
    mixpanel_data = base64.b64encode(json.dumps(mixpanel_params))
    mixpanel_resp = urllib2.urlopen("http://api.mixpanel.com/track/?data=%s" % mixpanel_data)
    mixpanel_response_read = mixpanel_resp.read()
    logger.info("mixpanel response:{response}".format(
        response = mixpanel_response_read))

    resp = make_response("duly noted. carry on.", 200)
    return resp

@app.route('/wospicker', methods=["GET"])
def wospicker():
    num_total = int(request.args.get("total"))
    num_subset = int(request.args.get("subset"))

    pages_and_ids = util.pickWosItems(num_total, num_subset)

    resp = make_response(json.dumps(pages_and_ids, indent=4), 200)
    resp.mimetype = "application/json"
    return resp

<commit_msg>change embedpage event back to embed<commit_after>import requests, iso8601, os, json, logging

from flask import Flask, jsonify, json, request, redirect, abort, make_response
from flask import render_template
import urllib2
import base64
import time
import urlparse

from totalimpactwebapp import app, util
from totalimpactwebapp.models import Github
from totalimpactwebapp import pretty_date

logger = logging.getLogger("tiwebapp.views")
mymixpanel = mixpanel.Mixpanel(token=os.getenv("MIXPANEL_TOKEN"))
mixpanel_token = os.getenv("MIXPANEL_TOKEN")

@app.before_request
def log_ip_address():
    if request.endpoint != "static":
        ip_address = request.remote_addr
        logger.info("%30s IP address calling %s %s" % (ip_address, request.method, request.url))

@app.after_request
def add_crossdomain_header(resp):
    resp.headers['Access-Control-Allow-Origin'] = "*"
    resp.headers['Access-Control-Allow-Methods'] = "POST, GET, OPTIONS, PUT, DELETE"
    resp.headers['Access-Control-Allow-Headers'] = "Content-Type"

    return resp

# static pages
@app.route('/')
def home():
    return render_template(
    'index.html',
    page_title="tell the full story of your research impact",
    body_class="homepage",
    mixpanel_token=os.environ["MIXPANEL_TOKEN"],
    api_root=os.environ["API_ROOT"]
    )

@app.route("/embed/impactstory.js")
@app.route("/embed/v1/impactstory.js")
def impactstory_dot_js():

    badges_template = render_template("js-template-badges.html").replace("\n", "")
    ich_def = open(os.path.dirname(__file__) + "/static/js/icanhaz.min.js", "r").read()
    ti_item_def = open(os.path.dirname(__file__) + "/static/js/ti-item.js", "r").read()

    rendered = render_template(
        "impactstory.js",
        ich=ich_def,
        ti_item=ti_item_def,
        badges_template=badges_template,
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        webapp_root=os.environ["WEBAPP_ROOT"]
    )
    resp = make_response(rendered)
    """
    There is no standard way to indicate you're sending back javascript;
    This seems the most recommended one, though. See
    http://stackoverflow.com/questions/2706290/why-write-script-type-text-javascript-when-the-mime-type-is-set-by-the-serve
    and http://www.ietf.org/rfc/rfc4329.txt
     """
    resp.headers["Content-Type"] = "application/javascript; charset=utf-8"
    return resp

@app.route("/embed/test")
def embed_test():
    return render_template(
        "sample-embed-internal-test.html",
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        webapp_root = os.environ["WEBAPP_ROOT"]
    )

@app.route('/about')
def about():
    return render_template(
        'about.html',
        page_title="about",
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"]
        )

@app.route('/faq')
def faq():
    # get the table of items and identifiers
    which_items_loc = os.path.join(
        os.path.dirname(__file__),
        "static",
        "whichartifacts.html"
        )
    which_item_types = open(which_items_loc).read()

    # get the static_meta info for each metric
    try:
        r = requests.get("http://" + os.environ["API_ROOT"] +'/provider')
        metadata = json.loads(r.text)
    except requests.ConnectionError:
        metadata = {}

    return render_template(
        'faq.html',
        page_title="faq",
        which_artifacts=which_item_types,
        provider_metadata=metadata,
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"]
        )

@app.route('/api-docs')
def apidocs():
    return render_template(
        'api-docs.html',
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        webapp_root = os.environ["WEBAPP_ROOT"],
        page_title="api & widget"
        )

@app.route('/pricing')
def pricing():
    return render_template(
        'pricing.html',
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        webapp_root = os.environ["WEBAPP_ROOT"],
        page_title="pricing"
        )

@app.route('/create')
def collection_create():
    return render_template(
        'create-collection.html',
        mixpanel_token=os.environ["MIXPANEL_TOKEN"],
        api_root=os.environ["API_ROOT"],
        page_title="create collection",
        body_class="create-collection"
        )

@app.route('/collection/<collection_id>')
def collection_report(collection_id):
    url = "http://{api_root}/collection/{collection_id}?include_items=0".format(
        api_root=os.getenv("API_ROOT"),
        collection_id=collection_id
    )

    r = requests.get(url)
    if r.status_code == 200:
        collection = json.loads(r.text)
        return render_template(
            'report.html',
            mixpanel_token=os.environ["MIXPANEL_TOKEN"],
            api_root=os.environ["API_ROOT"],
            api_key=os.environ["API_KEY"],
            request_url=request.url,
            page_title=collection["title"],
            body_class="report",
            report_id=collection["_id"],
            report_id_namespace="impactstory_collection_id",
            api_query="collection/"+collection["_id"]
        )
    else:
        abort(404, "This collection doesn't seem to exist yet. "+url)

@app.route('/item/<ns>/<path:id>')
def item_report(ns, id):
    url = "http://{api_root}/v1/item/{ns}/{id}?key={api_key}".format(
        api_root=os.getenv("API_ROOT"),
        ns=ns,
        id=id,
        api_key=os.environ["API_KEY"]
    )

    r = requests.get(url)
    if r.status_code <= 210: # allow unfinished items
        item = json.loads(r.text)
        return render_template(
            'report.html',
            mixpanel_token=os.environ["MIXPANEL_TOKEN"],
            api_root=os.environ["API_ROOT"],
            api_key=os.environ["API_KEY"],
            request_url=request.url,
            page_title="",
            body_class="report",
            report_id=id,
            report_id_namespace=ns,
            api_query="item/{ns}/{id}".format(ns=ns, id=id)
        )
    else:
        abort(404, "This item doesn't seem to exist yet. "+url)

@app.route('/vitals', methods=["POST"])
def vitals():
    """
    Logs reporting stats from the embed code to mixpanel

    Gets a "vitals" object that has a url and a list of "params" objects.
    Each widget on a page sends its own params object; in practice, these will
    almost certainly be identical since users aren't likely to set different
    options for widgets on the same page. It's a good way to count widgets on
    that page, though.

    For documentation on the keys in the the params object, see the default
    params listed at the head of impactstory.js's main() function, and also the
    api-docs page.
    """

    vitals = request.json
    properties = {
                    'token': mixpanel_token,
                    'time': int(time.time()),
                    'ip': request.remote_addr,
                    "$referring_domain": urlparse.urlsplit(request.referrer).netloc,
                    "$referrer" : request.referrer,
                    "$os": request.user_agent.platform,
                    "$browser": request.user_agent.browser,
                    "Embeds per page": len(vitals["allParams"])
                    "Host page": vitals["url"],
                    "API Key": vitals["allParams"][0]["api-key"],
                }

    for embed_param in vitals["allParams"][0]:
        properties["Embed:"+embed_param] = vitals["allParams"][0][embed_param]

    #logger.debug("properties are {properties}".format(
    #    properties=properties))

    mixpanel_params = {"event": "Impression:embed", "properties": properties}
    mixpanel_data = base64.b64encode(json.dumps(mixpanel_params))
    mixpanel_resp = urllib2.urlopen("http://api.mixpanel.com/track/?data=%s" % mixpanel_data)
    mixpanel_response_read = mixpanel_resp.read()
    logger.info("mixpanel response:{response}".format(
        response = mixpanel_response_read))

    resp = make_response("duly noted. carry on.", 200)
    return resp

@app.route('/wospicker', methods=["GET"])
def wospicker():
    num_total = int(request.args.get("total"))
    num_subset = int(request.args.get("subset"))

    pages_and_ids = util.pickWosItems(num_total, num_subset)

    resp = make_response(json.dumps(pages_and_ids, indent=4), 200)
    resp.mimetype = "application/json"
    return resp

<|endoftext|>


<commit_before>#!/usr/bin/env python3
#
# Tests if the Beeler-Reuter AP (toy) model works.
#
# This file is part of PINTS.
#  Copyright (c) 2017-2018, University of Oxford.
#  For licensing information, see the LICENSE file distributed with the PINTS
#  software package.
#
import unittest
import numpy as np
import pints
import pints.toy

# Consistent unit testing in Python 2 and 3
try:
    unittest.TestCase.assertRaisesRegex
except AttributeError:
    unittest.TestCase.assertRaisesRegex = unittest.TestCase.assertRaisesRegexp

class TestActionPotentialModel(unittest.TestCase):
    """
    Tests if the Beeler-Reuter AP (toy) model works.
    """

    def test_ap(self):
        model = pints.toy.ActionPotentialModel()
        p0 = model.suggested_parameters()
        self.assertEqual(len(p0), model.n_parameters())
        times = model.suggested_times()
        self.assertTrue(np.all(times[1:] > times[:-1]))
        values = model.simulate(p0, times)
        self.assertEqual(len(times), len(values))
        self.assertEqual(values.shape[1], model.n_outputs())

        # Try simulating all states
        model.simulate_all_states(p0, times)

        # Test setting and getting init cond.
        self.assertFalse(np.all(model.initial_conditions() == [-80, 1e-5]))
        model.set_initial_conditions([-80, 1e-5])
        self.assertTrue(np.all(model.initial_conditions() == [-80, 1e-5]))

        # Initial conditions cannot be negative
        self.assertRaisesRegex(
            ValueError, 'cannot be negative',
            pints.toy.ActionPotentialModel, [-80, -1])

if __name__ == '__main__':
    unittest.main()
<commit_msg>Added value-based tests for BR model<commit_after>#!/usr/bin/env python3
#
# Tests if the Beeler-Reuter AP (toy) model works.
#
# This file is part of PINTS.
#  Copyright (c) 2017-2018, University of Oxford.
#  For licensing information, see the LICENSE file distributed with the PINTS
#  software package.
#
import unittest
import numpy as np
import pints
import pints.toy

# Consistent unit testing in Python 2 and 3
try:
    unittest.TestCase.assertRaisesRegex
except AttributeError:
    unittest.TestCase.assertRaisesRegex = unittest.TestCase.assertRaisesRegexp

class TestActionPotentialModel(unittest.TestCase):
    """
    Tests if the Beeler-Reuter AP (toy) model works.
    """

    def test_creation(self):

        # Test creation
        model = pints.toy.ActionPotentialModel()

        # Test setting and getting init cond.
        self.assertFalse(np.all(model.initial_conditions() == [-80, 1e-5]))
        model.set_initial_conditions([-80, 1e-5])
        self.assertTrue(np.all(model.initial_conditions() == [-80, 1e-5]))

        # Initial conditions cannot be negative
        self.assertRaisesRegex(
            ValueError, 'cannot be negative',
            pints.toy.ActionPotentialModel, [-80, -1])

    def test_suggestions(self):

        model = pints.toy.ActionPotentialModel()
        p0 = model.suggested_parameters()
        self.assertEqual(len(p0), model.n_parameters())

        times = model.suggested_times()
        self.assertTrue(np.all(times[1:] > times[:-1]))

    def test_simulation(self):

        model = pints.toy.ActionPotentialModel()
        times = model.suggested_times()
        p0 = model.suggested_parameters()

        # Test simulating all eight states
        states = model.simulate_all_states(p0, times)
        self.assertEqual(len(states.shape), 2)
        self.assertEqual(states.shape[0], len(times))
        self.assertEqual(states.shape[1], 8)

        # Test initial state
        x0 = np.array([-84.622, 2e-7, 0.01, 0.99, 0.98, 0.003, 0.99, 0.0004])
        self.assertTrue(np.all(states[0] == x0))

        # Test state during AP (at 100ms, with stimulus applied at t=0ms)
        # Reference values taken from a Myokit simulation with the same model
        i100 = 200
        self.assertEqual(times[i100], 100)
        x100 = [
            1.09411249975881226e+01,
            6.14592181872724475e-06,
            9.93428572073381311e-01,
            2.17659582339377205e-11,
            -3.78573517405393772e-11,
            9.72515482289432853e-01,
            7.67322448537725688e-01,
            2.34776564989968184e-01,
        ]
        self.assertAlmostEqual(states[i100][0], x100[0], places=2)
        self.assertAlmostEqual(states[i100][1], x100[1], places=2)
        self.assertAlmostEqual(states[i100][2], x100[2], places=2)
        self.assertAlmostEqual(states[i100][3], x100[3], places=2)
        self.assertAlmostEqual(states[i100][4], x100[4], places=2)
        self.assertAlmostEqual(states[i100][5], x100[5], places=2)
        self.assertAlmostEqual(states[i100][6], x100[6], places=2)
        self.assertAlmostEqual(states[i100][7], x100[7], places=2)

        # Test simulation outputting only observables
        partial = model.simulate(p0, times)
        self.assertEqual(len(times), len(partial))
        self.assertEqual(partial.shape[1], model.n_outputs())

        partial = np.array(partial)
        states = np.array(states)
        self.assertTrue(np.all(partial[:, 0] == states[:, 0]))
        self.assertTrue(np.all(partial[:, 1] == states[:, 1]))

if __name__ == '__main__':
    unittest.main()
<|endoftext|>


<commit_before>"""Converts tracks or albums to external directory
"""
import logging
import os
import threading
import shutil
from subprocess import Popen, PIPE

import imghdr

from beets.plugins import BeetsPlugin
from beets import ui, library, util, mediafile
from beets.util.functemplate import Template

log = logging.getLogger('beets')
DEVNULL = open(os.devnull, 'wb')
conf = {}

def _embed(path, items):
    """Embed an image file, located at `path`, into each item.
    """
    data = open(util.syspath(path), 'rb').read()
    kindstr = imghdr.what(None, data)
    if kindstr not in ('jpeg', 'png'):
        log.error('A file of type %s is not allowed as coverart.' % kindstr)
    return
    log.debug('Embedding album art.')
    for item in items:
        try:
            f = mediafile.MediaFile(syspath(item.path))
        except mediafile.UnreadableFileError as exc:
            log.warn('Could not embed art in {0}: {1}'.format(
                repr(item.path), exc
            ))
            continue
        f.art = data
        f.save()

class encodeThread(threading.Thread):
    def __init__(self, source, dest, artpath):
        threading.Thread.__init__(self)
        self.source = source
        self.dest = dest
        self.artpath = artpath

    def run(self):
        sema.acquire()
        self.encode()
        sema.release()

        dest_item = library.Item.from_path(self.source)
        dest_item.path = self.dest
        dest_item.write()
        if self.artpath:
            _embed(self.artpath,[dest_item])

    def encode(self):
        log.info('Started encoding '+ self.source)
        temp_dest = self.dest + '~'

        source_ext = os.path.splitext(self.source)[1].lower()
        if source_ext == '.flac':
            decode = Popen([conf['flac'], '-c', '-d', '-s', self.source],
                           stdout=PIPE)
            encode = Popen([conf['lame']] + conf['opts'] + ['-', temp_dest],
                       stdin=decode.stdout, stderr=DEVNULL)
            decode.stdout.close()
            encode.communicate()
        elif source_ext == '.mp3':
            encode = Popen([conf['lame']] + conf['opts'] + ['--mp3input'] +
                 [self.source, temp_dest], close_fds=True, stderr=DEVNULL)
            encode.communicate()
        else:
            log.error('Only converting from FLAC or MP3 implemented')
            return
        shutil.move(temp_dest, self.dest)
        log.info('Finished encoding '+ self.source)

def convert_item(lib, item, dest_dir, artpath):
    if item.format != 'FLAC' and item.format != 'MP3':
        log.info('Skipping {0} : not supported format'.format(item.path))
        return

    dest = os.path.join(dest_dir,lib.destination(item, fragment = True))
    dest = os.path.splitext(dest)[0] + '.mp3'

    if not os.path.exists(dest):
        util.mkdirall(dest)
        if item.format == 'MP3' and item.bitrate < 1000*conf['max_bitrate']:
            log.info('Copying {0}'.format(item.path))
            shutil.copy(item.path, dest)
            if artpath:
                _embed(artpath,[library.Item.from_path(dest)])
        else:
            thread = encodeThread(item.path, dest, artpath)
            thread.start()
    else:
        log.info('Skipping {0} : target file exists'.format(item.path))

def convert_func(lib, config, opts, args):
    global sema

    dest = opts.dest if opts.dest is not None else conf['dest']
    if not dest:
        log.error('No destination set')
        return
    threads = opts.threads if opts.threads is not None else conf['threads']
    sema = threading.BoundedSemaphore(threads)

    fmt = '$albumartist - $album' if opts.album \
                                  else '$artist - $album - $title'
    ui.commands.list_items(lib, ui.decargs(args), opts.album, False, fmt)

    if not ui.input_yn("Convert? (Y/n)"):
        return

    if opts.album:
        for album in lib.albums(ui.decargs(args)):
            for item in album.items():
                convert_item(lib, item, dest, album.artpath)
    else:
        for item in lib.items(ui.decargs(args)):
            convert_item(lib, item, dest, lib.get_album(item).artpath)

class ConvertPlugin(BeetsPlugin):
    def configure(self, config):
        conf['dest'] = ui.config_val(config, 'convert', 'dest', None)
        conf['threads'] = ui.config_val(config, 'convert', 'threads', 2)
        conf['flac'] = ui.config_val(config, 'convert', 'flac', 'flac')
        conf['lame'] = ui.config_val(config, 'convert', 'lame', 'lame')
        conf['opts'] = ui.config_val(config, 'convert',
                                     'opts', '-V2').split(' ')
        conf['max_bitrate'] = int(ui.config_val(config, 'convert',
                                                'max_bitrate','500'))

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-a', '--album', action='store_true',
                        help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                        help='change the number of threads (default 2)')
        cmd.parser.add_option('-d', '--dest', action='store',
                        help='set the destination directory')
        cmd.func = convert_func
        return [cmd]
<commit_msg>Add option to disable embedding<commit_after>"""Converts tracks or albums to external directory
"""
import logging
import os
import threading
import shutil
from subprocess import Popen, PIPE

import imghdr

from beets.plugins import BeetsPlugin
from beets import ui, library, util, mediafile
from beets.util.functemplate import Template

log = logging.getLogger('beets')
DEVNULL = open(os.devnull, 'wb')
conf = {}

def _embed(path, items):
    """Embed an image file, located at `path`, into each item.
    """
    data = open(util.syspath(path), 'rb').read()
    kindstr = imghdr.what(None, data)
    if kindstr not in ('jpeg', 'png'):
        log.error('A file of type %s is not allowed as coverart.' % kindstr)
    return
    log.debug('Embedding album art.')
    for item in items:
        try:
            f = mediafile.MediaFile(syspath(item.path))
        except mediafile.UnreadableFileError as exc:
            log.warn('Could not embed art in {0}: {1}'.format(
                repr(item.path), exc
            ))
            continue
        f.art = data
        f.save()

class encodeThread(threading.Thread):
    def __init__(self, source, dest, artpath):
        threading.Thread.__init__(self)
        self.source = source
        self.dest = dest
        self.artpath = artpath

    def run(self):
        sema.acquire()
        self.encode()
        sema.release()

        dest_item = library.Item.from_path(self.source)
        dest_item.path = self.dest
        dest_item.write()
        if self.artpath and conf['embed']:
            _embed(self.artpath,[dest_item])

    def encode(self):
        log.info('Started encoding '+ self.source)
        temp_dest = self.dest + '~'

        source_ext = os.path.splitext(self.source)[1].lower()
        if source_ext == '.flac':
            decode = Popen([conf['flac'], '-c', '-d', '-s', self.source],
                           stdout=PIPE)
            encode = Popen([conf['lame']] + conf['opts'] + ['-', temp_dest],
                       stdin=decode.stdout, stderr=DEVNULL)
            decode.stdout.close()
            encode.communicate()
        elif source_ext == '.mp3':
            encode = Popen([conf['lame']] + conf['opts'] + ['--mp3input'] +
                 [self.source, temp_dest], close_fds=True, stderr=DEVNULL)
            encode.communicate()
        else:
            log.error('Only converting from FLAC or MP3 implemented')
            return
        shutil.move(temp_dest, self.dest)
        log.info('Finished encoding '+ self.source)

def convert_item(lib, item, dest_dir, artpath):
    if item.format != 'FLAC' and item.format != 'MP3':
        log.info('Skipping {0} : not supported format'.format(item.path))
        return

    dest = os.path.join(dest_dir,lib.destination(item, fragment = True))
    dest = os.path.splitext(dest)[0] + '.mp3'

    if not os.path.exists(dest):
        util.mkdirall(dest)
        if item.format == 'MP3' and item.bitrate < 1000*conf['max_bitrate']:
            log.info('Copying {0}'.format(item.path))
            shutil.copy(item.path, dest)
            if artpath and conf['embed']:
                _embed(artpath,[library.Item.from_path(dest)])
        else:
            thread = encodeThread(item.path, dest, artpath)
            thread.start()
    else:
        log.info('Skipping {0} : target file exists'.format(item.path))

def convert_func(lib, config, opts, args):
    global sema

    dest = opts.dest if opts.dest is not None else conf['dest']
    if not dest:
        log.error('No destination set')
        return
    threads = opts.threads if opts.threads is not None else conf['threads']
    sema = threading.BoundedSemaphore(threads)

    fmt = '$albumartist - $album' if opts.album \
                                  else '$artist - $album - $title'
    ui.commands.list_items(lib, ui.decargs(args), opts.album, False, fmt)

    if not ui.input_yn("Convert? (Y/n)"):
        return

    if opts.album:
        for album in lib.albums(ui.decargs(args)):
            for item in album.items():
                convert_item(lib, item, dest, album.artpath)
    else:
        for item in lib.items(ui.decargs(args)):
            convert_item(lib, item, dest, lib.get_album(item).artpath)

class ConvertPlugin(BeetsPlugin):
    def configure(self, config):
        conf['dest'] = ui.config_val(config, 'convert', 'dest', None)
        conf['threads'] = ui.config_val(config, 'convert', 'threads', 2)
        conf['flac'] = ui.config_val(config, 'convert', 'flac', 'flac')
        conf['lame'] = ui.config_val(config, 'convert', 'lame', 'lame')
        conf['opts'] = ui.config_val(config, 'convert',
                                     'opts', '-V2').split(' ')
        conf['max_bitrate'] = int(ui.config_val(config, 'convert',
                                                'max_bitrate','500'))
        conf['embed'] = ui.config_val(config, 'convert', 'embed', True,
                                      vtype = bool)

    def commands(self):
        cmd = ui.Subcommand('convert', help='convert to external location')
        cmd.parser.add_option('-a', '--album', action='store_true',
                        help='choose albums instead of tracks')
        cmd.parser.add_option('-t', '--threads', action='store', type='int',
                        help='change the number of threads (default 2)')
        cmd.parser.add_option('-d', '--dest', action='store',
                        help='set the destination directory')
        cmd.func = convert_func
        return [cmd]
<|endoftext|>


<commit_before># coding=utf8
"""
admin.py - Willie Admin Module
Copyright 2010-2011, Michael Yanovich, Alek Rollyson, and Edward Powell
Copyright © 2012, Elad Alfassa <elad@fedoraproject.org>
Licensed under the Eiffel Forum License 2.

http://willie.dftba.net/

"""
from __future__ import unicode_literals

import re
from willie.module import commands, priority, OP
from willie.tools import Nick

def setup(bot):
    #Having a db means pref's exists. Later, we can just use `if bot.db`.
    if bot.db and not bot.db.preferences.has_columns('topic_mask'):
        bot.db.preferences.add_columns(['topic_mask'])

@commands('op')
def op(bot, trigger):
    """
    Command to op users in a room. If no nick is given,
    willie will op the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] >= OP:
        nick = trigger.group(2)
        channel = trigger.sender
        if not nick:
            nick = trigger.nick
        bot.write(['MODE', channel, "+o", nick])

@commands('deop')
def deop(bot, trigger):
    """
    Command to deop users in a room. If no nick is given,
    willie will deop the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] >= OP:
        nick = trigger.group(2)
        channel = trigger.sender
        if not nick:
            nick = trigger.nick
        bot.write(['MODE', channel, "-o", nick])

@commands('voice')
def voice(bot, trigger):
    """
    Command to voice users in a room. If no nick is given,
    willie will voice the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] >= OP:
        nick = trigger.group(2)
        channel = trigger.sender
        if not nick:
            nick = trigger.nick
        bot.write(['MODE', channel, "+v", nick])

@commands('devoice')
def devoice(bot, trigger):
    """
    Command to devoice users in a room. If no nick is given,
    willie will devoice the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] >= OP:
        nick = trigger.group(2)
        channel = trigger.sender
        if not nick:
            nick = trigger.nick
        bot.write(['MODE', channel, "-v", nick])

@commands('kick')
@priority('high')
def kick(bot, trigger):
    """
    Kick a user from the channel.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    nick = opt
    channel = trigger.sender
    reasonidx = 2
    if not opt.is_nick():
        if argc < 3:
            return
        nick = text[2]
        channel = opt
        reasonidx = 3
    reason = ' '.join(text[reasonidx:])
    if nick != bot.config.nick:
        bot.write(['KICK', channel, nick, reason])

def configureHostMask(mask):
    if mask == '*!*@*':
        return mask
    if re.match('^[^.@!/]+$', mask) is not None:
        return '%s!*@*' % mask
    if re.match('^[^@!]+$', mask) is not None:
        return '*!*@%s' % mask

    m = re.match('^([^!@]+)@$', mask)
    if m is not None:
        return '*!%s@*' % m.group(1)

    m = re.match('^([^!@]+)@([^@!]+)$', mask)
    if m is not None:
        return '*!%s@%s' % (m.group(1), m.group(2))

    m = re.match('^([^!@]+)!(^[!@]+)@?$', mask)
    if m is not None:
        return '%s!%s@*' % (m.group(1), m.group(2))
    return ''

@commands('ban')
@priority('high')
def ban(bot, trigger):
    """
    This give admins the ability to ban a user.
    The bot must be a Channel Operator for this command to work.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    banmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        channel = opt
        banmask = text[2]
    banmask = configureHostMask(banmask)
    if banmask == '':
        return
    bot.write(['MODE', channel, '+b', banmask])

@commands('unban')
def unban(bot, trigger):
    """
    This give admins the ability to unban a user.
    The bot must be a Channel Operator for this command to work.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    banmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        channel = opt
        banmask = text[2]
    banmask = configureHostMask(banmask)
    if banmask == '':
        return
    bot.write(['MODE', channel, '-b', banmask])

@commands('quiet')
def quiet(bot, trigger):
    """
    This gives admins the ability to quiet a user.
    The bot must be a Channel Operator for this command to work
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    quietmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        quietmask = text[2]
        channel = opt
    quietmask = configureHostMask(quietmask)
    if quietmask == '':
        return
    bot.write(['MODE', channel, '+q', quietmask])

@commands('unquiet')
def unquiet(bot, trigger):
    """
   This gives admins the ability to unquiet a user.
   The bot must be a Channel Operator for this command to work
   """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    quietmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        quietmask = text[2]
        channel = opt
    quietmask = configureHostMask(quietmask)
    if quietmask == '':
        return
    bot.write(['MODE', opt, '-q', quietmask])

@commands('kickban', 'kb')
@priority('high')
def kickban(bot, trigger):
    """
   This gives admins the ability to kickban a user.
   The bot must be a Channel Operator for this command to work
   .kickban [#chan] user1 user!*@* get out of here
   """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    text = trigger.group().split()
    argc = len(text)
    if argc < 4:
        return
    opt = Nick(text[1])
    nick = opt
    mask = text[2]
    reasonidx = 3
    if not opt.is_nick():
        if argc < 5:
            return
        channel = opt
        nick = text[2]
        mask = text[3]
        reasonidx = 4
    reason = ' '.join(text[reasonidx:])
    mask = configureHostMask(mask)
    if mask == '':
        return
    bot.write(['MODE', channel, '+b', mask])
    bot.write(['KICK', channel, nick, ' :', reason])

@commands('topic')
def topic(bot, trigger):
    """
    This gives ops the ability to change the topic.
    """
    purple, green, bold = '\x0306', '\x0310', '\x02'
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    text = trigger.group(2)
    if text == '':
        return
    channel = trigger.sender.lower()

    narg = 1
    mask = None
    if bot.db and channel in bot.db.preferences:
        mask = bot.db.preferences.get(channel, 'topic_mask')
        narg = len(re.findall('%s', mask))
    if not mask or mask == '':
        mask = purple + 'Welcome to: ' + green + channel + purple \
            + ' | ' + bold + 'Topic: ' + bold + green + '%s'

    top = trigger.group(2)
    text = tuple()
    if top:
        text = tuple(unicode.split(top, '~', narg))

    if len(text) != narg:
        message = "Not enough arguments. You gave " + str(len(text)) + ', it requires ' + str(narg) + '.'
        return bot.say(message)
    topic = mask % text

    bot.write(('TOPIC', channel + ' :' + topic))

@commands('tmask')
def set_mask(bot, trigger):
    """
    Set the mask to use for .topic in the current channel. %s is used to allow
    substituting in chunks of text.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if not bot.db:
        bot.say("I'm afraid I can't do that.")
    else:
        bot.db.preferences.update(trigger.sender.lower(), {'topic_mask': trigger.group(2)})
        bot.say("Gotcha, " + trigger.nick)

@commands('showmask')
def show_mask(bot, trigger):
    """Show the topic mask for the current channel."""
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if not bot.db:
        bot.say("I'm afraid I can't do that.")
    elif trigger.sender.lower() in bot.db.preferences:
        bot.say(bot.db.preferences.get(trigger.sender.lower(), 'topic_mask'))
    else:
        bot.say("%s")
<commit_msg>Add code for #564 as a new branch<commit_after># coding=utf8
"""
admin.py - Willie Admin Module
Copyright 2010-2011, Michael Yanovich, Alek Rollyson, and Edward Powell
Copyright © 2012, Elad Alfassa <elad@fedoraproject.org>
Licensed under the Eiffel Forum License 2.

http://willie.dftba.net/

"""
from __future__ import unicode_literals

import re
from willie.module import commands, priority, OP, HALFOP
from willie.tools import Nick

def setup(bot):
    #Having a db means pref's exists. Later, we can just use `if bot.db`.
    if bot.db and not bot.db.preferences.has_columns('topic_mask'):
        bot.db.preferences.add_columns(['topic_mask'])

@commands('op')
def op(bot, trigger):
    """
    Command to op users in a room. If no nick is given,
    willie will op the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < OP:
        return bot.reply("I'm not a channel operator!")
    nick = trigger.group(2)
    channel = trigger.sender
    if not nick:
        nick = trigger.nick
    bot.write(['MODE', channel, "+o", nick])

@commands('deop')
def deop(bot, trigger):
    """
    Command to deop users in a room. If no nick is given,
    willie will deop the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < OP:
        return bot.reply("I'm not a channel operator!")
    nick = trigger.group(2)
    channel = trigger.sender
    if not nick:
        nick = trigger.nick
    bot.write(['MODE', channel, "-o", nick])

@commands('voice')
def voice(bot, trigger):
    """
    Command to voice users in a room. If no nick is given,
    willie will voice the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < HALFOP:
        return bot.reply("I'm not a channel operator!")
    nick = trigger.group(2)
    channel = trigger.sender
    if not nick:
        nick = trigger.nick
    bot.write(['MODE', channel, "+v", nick])

@commands('devoice')
def devoice(bot, trigger):
    """
    Command to devoice users in a room. If no nick is given,
    willie will devoice the nick who sent the command
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < HALFOP:
        return bot.reply("I'm not a channel operator!")
    nick = trigger.group(2)
    channel = trigger.sender
    if not nick:
        nick = trigger.nick
    bot.write(['MODE', channel, "-v", nick])

@commands('kick')
@priority('high')
def kick(bot, trigger):
    """
    Kick a user from the channel.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < HALFOP:
        return bot.reply("I'm not a channel operator!")
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    nick = opt
    channel = trigger.sender
    reasonidx = 2
    if not opt.is_nick():
        if argc < 3:
            return
        nick = text[2]
        channel = opt
        reasonidx = 3
    reason = ' '.join(text[reasonidx:])
    if nick != bot.config.nick:
        bot.write(['KICK', channel, nick, reason])

def configureHostMask(mask):
    if mask == '*!*@*':
        return mask
    if re.match('^[^.@!/]+$', mask) is not None:
        return '%s!*@*' % mask
    if re.match('^[^@!]+$', mask) is not None:
        return '*!*@%s' % mask

    m = re.match('^([^!@]+)@$', mask)
    if m is not None:
        return '*!%s@*' % m.group(1)

    m = re.match('^([^!@]+)@([^@!]+)$', mask)
    if m is not None:
        return '*!%s@%s' % (m.group(1), m.group(2))

    m = re.match('^([^!@]+)!(^[!@]+)@?$', mask)
    if m is not None:
        return '%s!%s@*' % (m.group(1), m.group(2))
    return ''

@commands('ban')
@priority('high')
def ban(bot, trigger):
    """
    This give admins the ability to ban a user.
    The bot must be a Channel Operator for this command to work.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < HALFOP:
        return bot.reply("I'm not a channel operator!")
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    banmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        channel = opt
        banmask = text[2]
    banmask = configureHostMask(banmask)
    if banmask == '':
        return
    bot.write(['MODE', channel, '+b', banmask])

@commands('unban')
def unban(bot, trigger):
    """
    This give admins the ability to unban a user.
    The bot must be a Channel Operator for this command to work.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < HALFOP:
        return bot.reply("I'm not a channel operator!")
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    banmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        channel = opt
        banmask = text[2]
    banmask = configureHostMask(banmask)
    if banmask == '':
        return
    bot.write(['MODE', channel, '-b', banmask])

@commands('quiet')
def quiet(bot, trigger):
    """
    This gives admins the ability to quiet a user.
    The bot must be a Channel Operator for this command to work.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < OP:
        return bot.reply("I'm not a channel operator!")
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    quietmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        quietmask = text[2]
        channel = opt
    quietmask = configureHostMask(quietmask)
    if quietmask == '':
        return
    bot.write(['MODE', channel, '+q', quietmask])

@commands('unquiet')
def unquiet(bot, trigger):
    """
   This gives admins the ability to unquiet a user.
   The bot must be a Channel Operator for this command to work.
   """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < OP:
        return bot.reply("I'm not a channel operator!")
    text = trigger.group().split()
    argc = len(text)
    if argc < 2:
        return
    opt = Nick(text[1])
    quietmask = opt
    channel = trigger.sender
    if not opt.is_nick():
        if argc < 3:
            return
        quietmask = text[2]
        channel = opt
    quietmask = configureHostMask(quietmask)
    if quietmask == '':
        return
    bot.write(['MODE', opt, '-q', quietmask])

@commands('kickban', 'kb')
@priority('high')
def kickban(bot, trigger):
    """
   This gives admins the ability to kickban a user.
   The bot must be a Channel Operator for this command to work.
   .kickban [#chan] user1 user!*@* get out of here
   """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < HALFOP:
        return bot.reply("I'm not a channel operator!")
    text = trigger.group().split()
    argc = len(text)
    if argc < 4:
        return
    opt = Nick(text[1])
    nick = opt
    mask = text[2]
    reasonidx = 3
    if not opt.is_nick():
        if argc < 5:
            return
        channel = opt
        nick = text[2]
        mask = text[3]
        reasonidx = 4
    reason = ' '.join(text[reasonidx:])
    mask = configureHostMask(mask)
    if mask == '':
        return
    bot.write(['MODE', channel, '+b', mask])
    bot.write(['KICK', channel, nick, ' :', reason])

@commands('topic')
def topic(bot, trigger):
    """
    This gives ops the ability to change the topic.
    The bot must be a Channel Operator for this command to work.
    """
    purple, green, bold = '\x0306', '\x0310', '\x02'
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if bot.privileges[trigger.sender][bot.nick] < HALFOP:
        return bot.reply("I'm not a channel operator!")
    text = trigger.group(2)
    if text == '':
        return
    channel = trigger.sender.lower()

    narg = 1
    mask = None
    if bot.db and channel in bot.db.preferences:
        mask = bot.db.preferences.get(channel, 'topic_mask')
        narg = len(re.findall('%s', mask))
    if not mask or mask == '':
        mask = purple + 'Welcome to: ' + green + channel + purple \
            + ' | ' + bold + 'Topic: ' + bold + green + '%s'

    top = trigger.group(2)
    text = tuple()
    if top:
        text = tuple(unicode.split(top, '~', narg))

    if len(text) != narg:
        message = "Not enough arguments. You gave " + str(len(text)) + ', it requires ' + str(narg) + '.'
        return bot.say(message)
    topic = mask % text

    bot.write(('TOPIC', channel + ' :' + topic))

@commands('tmask')
def set_mask(bot, trigger):
    """
    Set the mask to use for .topic in the current channel. %s is used to allow
    substituting in chunks of text.
    """
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if not bot.db:
        bot.say("I'm afraid I can't do that.")
    else:
        bot.db.preferences.update(trigger.sender.lower(), {'topic_mask': trigger.group(2)})
        bot.say("Gotcha, " + trigger.nick)

@commands('showmask')
def show_mask(bot, trigger):
    """Show the topic mask for the current channel."""
    if bot.privileges[trigger.sender][trigger.nick] < OP:
        return
    if not bot.db:
        bot.say("I'm afraid I can't do that.")
    elif trigger.sender.lower() in bot.db.preferences:
        bot.say(bot.db.preferences.get(trigger.sender.lower(), 'topic_mask'))
    else:
        bot.say("%s")
<|endoftext|>


<commit_before>__author__ = 'pavel'

import argparse
import sys

from itertools import permutations
from multiprocessing import Pool, cpu_count
from rdkit import Chem

from mol_context import get_std_context_core_permutations

def create_db(conn):
    cur = conn.cursor()
    cur.execute("""CREATE TABLE IF NOT EXISTS core_table
                  (
                    std_context_smi TEXT NOT NULL,
                    std_core_smi TEXT NOT NULL,
                    smart_core_smi TEXT NOT NULL
                  )""")
    cur.execute("DELETE FROM core_table")
    cur.execute("DROP INDEX IF EXISTS std_context_smi_idx")
    cur.execute("DROP INDEX IF EXISTS std_core_smi_idx")
    conn.commit()

def insert_db(cur, items):
    cur.executemany("INSERT INTO core_table VALUES(?,?,?)", items)

def compress_db(conn):
    cur = conn.cursor()
    cur.execute("DELETE FROM core_table WHERE rowid NOT IN (SELECT MIN(rowid) "
                "FROM core_table GROUP BY std_context_smi, std_core_smi)")
    conn.commit()

def index_db(conn):
    cur = conn.cursor()
    cur.execute("CREATE INDEX std_context_smi_idx ON core_table (std_context_smi)")
    cur.execute("CREATE INDEX std_core_smi_idx ON core_table (std_core_smi)")
    conn.commit()

def process_line(line):
    # returns env_smi, core_smi, heavy_atoms_num, core_smarts

    output = []
    smi, id, core, context = line.strip().split(',')

    if (not core and not context) or (keep_mols_set and id not in keep_mols_set):
        return output
    else:
        # one split
        if not core:
            residues = context.split('.')
            if len(residues) == 2:
                for context, core in permutations(residues, 2):
                    if context == '[H][*:1]':   # ignore such cases
                        continue
                    mm = Chem.MolFromSmiles(core, sanitize=False)
                    num_heavy_atoms = mm.GetNumHeavyAtoms() if mm else float('inf')
                    if num_heavy_atoms <= max_heavy_atoms:
                        env, cores = get_std_context_core_permutations(context, core, radius, keep_stereo)
                        if env and cores:
                            # for 1 cut cores will always contain 1 item
                            if not store_comp_id:
                                output.append((env, cores[0], num_heavy_atoms))
                            else:
                                output.append((env, cores[0], num_heavy_atoms, id))
            else:
                sys.stderr.write('more than two fragments in context (%s) where core is empty' % context)
                sys.stderr.flush()
        # two or more splits
        else:
            mm = Chem.MolFromSmiles(core, sanitize=False)
            num_heavy_atoms = mm.GetNumHeavyAtoms() if mm else float('inf')
            if num_heavy_atoms <= max_heavy_atoms:
                env, cores = get_std_context_core_permutations(context, core, radius, keep_stereo)
                if env and cores:
                    for c in cores:
                        if not store_comp_id:
                            output.append((env, c, num_heavy_atoms))
                        else:
                            output.append((env, c, num_heavy_atoms, id))
        return output

def init(keep_mols):
    global keep_mols_set
    keep_mols_set = set([line.strip() for line in open(keep_mols).readlines()]) if keep_mols else set()

def main(input_fname, output_fname, keep_mols, radius, keep_stereo, max_heavy_atoms, ncpu, store_comp_id, verbose):

    # radius and remove_stereo are supplied to process_context_core via global environment (ugly but working solution)

    ncpu = min(cpu_count(), max(ncpu, 1))
    p = Pool(ncpu, initializer=init, initargs=(keep_mols,))

    try:
        with open(output_fname, 'wt') as out:

            with open(input_fname) as f:

                for i, res in enumerate(p.imap_unordered(process_line, f, chunksize=1000), 1):

                    for item in res:
                        if item:
                            out.write(','.join(map(str, item)) + '\n')

                    if verbose and i % 1000 == 0:
                        sys.stderr.write('\r%i lines passed' % i)
                        sys.stderr.flush()

    finally:
        p.close()

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Create text file for fragment replacement from fragmented molecules '
                                                 'obtained with fragmentation.py. '
                                                 'The output may contain duplicated lines which should be filtered out '
                                                 'externally.')
    parser.add_argument('-i', '--input', metavar='frags.txt', required=True,
                        help='fragmented molecules.')
    parser.add_argument('-o', '--out', metavar='output.txt', required=True,
                        help='output text file.')
    parser.add_argument('-k', '--keep_mols', metavar='molnames.txt', required=False, default=None,
                        help='file with mol names to keep. Molecules which are not in the list will be ignored.')
    parser.add_argument('-r', '--radius', metavar='NUMBER', required=False, default=1,
                        help='radius of molecular context (in bonds) which will be taken into account. Default: 1.')
    parser.add_argument('-a', '--max_heavy_atoms', metavar='NUMBER', required=False, default=20,
                        help='maximum number of heavy atoms in cores. If the number of atoms exceeds the limit '
                             'fragment will be discarded. Default: 20.')
    parser.add_argument('-s', '--keep_stereo', action='store_true', default=False,
                        help='set this flag if you want to keep stereo in context and core parts.')
    parser.add_argument('-c', '--ncpu', metavar='NUMBER', required=False, default=1,
                        help='number of cpus used for computation. Default: 1.')
    parser.add_argument('--store_comp_id', action='store_true', default=False,
                        help='store compound id in output (only for debug).')
    parser.add_argument('-v', '--verbose', action='store_true', default=False,
                        help='print progress.')

    args = vars(parser.parse_args())
    for o, v in args.items():
        if o == "input": input_fname = v
        if o == "out": output_fname = v
        if o == "verbose": verbose = v
        if o == "radius": radius = int(v)
        if o == "keep_stereo": keep_stereo = v
        if o == "ncpu": ncpu = int(v)
        if o == "max_heavy_atoms": max_heavy_atoms = int(v)
        if o == "keep_mols": keep_mols = v
        if o == "store_comp_id": store_comp_id = v

    main(input_fname=input_fname,
         output_fname=output_fname,
         keep_mols=keep_mols,
         radius=radius,
         keep_stereo=keep_stereo,
         max_heavy_atoms=max_heavy_atoms,
         ncpu=ncpu,
         store_comp_id=store_comp_id,
         verbose=verbose)
<commit_msg>Remove obsolete code related to DB creation from frag_to_env_mp.py<commit_after>__author__ = 'pavel'

import argparse
import sys

from itertools import permutations
from multiprocessing import Pool, cpu_count
from rdkit import Chem

from mol_context import get_std_context_core_permutations

def process_line(line):
    # returns env_smi, core_smi, heavy_atoms_num, core_smarts

    output = []
    smi, id, core, context = line.strip().split(',')

    if (not core and not context) or (keep_mols_set and id not in keep_mols_set):
        return output
    else:
        # one split
        if not core:
            residues = context.split('.')
            if len(residues) == 2:
                for context, core in permutations(residues, 2):
                    if context == '[H][*:1]':   # ignore such cases
                        continue
                    mm = Chem.MolFromSmiles(core, sanitize=False)
                    num_heavy_atoms = mm.GetNumHeavyAtoms() if mm else float('inf')
                    if num_heavy_atoms <= max_heavy_atoms:
                        env, cores = get_std_context_core_permutations(context, core, radius, keep_stereo)
                        if env and cores:
                            # for 1 cut cores will always contain 1 item
                            if not store_comp_id:
                                output.append((env, cores[0], num_heavy_atoms))
                            else:
                                output.append((env, cores[0], num_heavy_atoms, id))
            else:
                sys.stderr.write('more than two fragments in context (%s) where core is empty' % context)
                sys.stderr.flush()
        # two or more splits
        else:
            mm = Chem.MolFromSmiles(core, sanitize=False)
            num_heavy_atoms = mm.GetNumHeavyAtoms() if mm else float('inf')
            if num_heavy_atoms <= max_heavy_atoms:
                env, cores = get_std_context_core_permutations(context, core, radius, keep_stereo)
                if env and cores:
                    for c in cores:
                        if not store_comp_id:
                            output.append((env, c, num_heavy_atoms))
                        else:
                            output.append((env, c, num_heavy_atoms, id))
        return output

def init(keep_mols):
    global keep_mols_set
    keep_mols_set = set([line.strip() for line in open(keep_mols).readlines()]) if keep_mols else set()

def main(input_fname, output_fname, keep_mols, radius, keep_stereo, max_heavy_atoms, ncpu, store_comp_id, verbose):

    # radius and remove_stereo are supplied to process_context_core via global environment (ugly but working solution)

    ncpu = min(cpu_count(), max(ncpu, 1))
    p = Pool(ncpu, initializer=init, initargs=(keep_mols,))

    try:
        with open(output_fname, 'wt') as out:

            with open(input_fname) as f:

                for i, res in enumerate(p.imap_unordered(process_line, f, chunksize=1000), 1):

                    for item in res:
                        if item:
                            out.write(','.join(map(str, item)) + '\n')

                    if verbose and i % 1000 == 0:
                        sys.stderr.write('\r%i lines passed' % i)
                        sys.stderr.flush()

    finally:
        p.close()

if __name__ == '__main__':

    parser = argparse.ArgumentParser(description='Create text file for fragment replacement from fragmented molecules '
                                                 'obtained with fragmentation.py. '
                                                 'The output may contain duplicated lines which should be filtered out '
                                                 'externally.')
    parser.add_argument('-i', '--input', metavar='frags.txt', required=True,
                        help='fragmented molecules.')
    parser.add_argument('-o', '--out', metavar='output.txt', required=True,
                        help='output text file.')
    parser.add_argument('-k', '--keep_mols', metavar='molnames.txt', required=False, default=None,
                        help='file with mol names to keep. Molecules which are not in the list will be ignored.')
    parser.add_argument('-r', '--radius', metavar='NUMBER', required=False, default=1,
                        help='radius of molecular context (in bonds) which will be taken into account. Default: 1.')
    parser.add_argument('-a', '--max_heavy_atoms', metavar='NUMBER', required=False, default=20,
                        help='maximum number of heavy atoms in cores. If the number of atoms exceeds the limit '
                             'fragment will be discarded. Default: 20.')
    parser.add_argument('-s', '--keep_stereo', action='store_true', default=False,
                        help='set this flag if you want to keep stereo in context and core parts.')
    parser.add_argument('-c', '--ncpu', metavar='NUMBER', required=False, default=1,
                        help='number of cpus used for computation. Default: 1.')
    parser.add_argument('--store_comp_id', action='store_true', default=False,
                        help='store compound id in output (only for debug).')
    parser.add_argument('-v', '--verbose', action='store_true', default=False,
                        help='print progress.')

    args = vars(parser.parse_args())
    for o, v in args.items():
        if o == "input": input_fname = v
        if o == "out": output_fname = v
        if o == "verbose": verbose = v
        if o == "radius": radius = int(v)
        if o == "keep_stereo": keep_stereo = v
        if o == "ncpu": ncpu = int(v)
        if o == "max_heavy_atoms": max_heavy_atoms = int(v)
        if o == "keep_mols": keep_mols = v
        if o == "store_comp_id": store_comp_id = v

    main(input_fname=input_fname,
         output_fname=output_fname,
         keep_mols=keep_mols,
         radius=radius,
         keep_stereo=keep_stereo,
         max_heavy_atoms=max_heavy_atoms,
         ncpu=ncpu,
         store_comp_id=store_comp_id,
         verbose=verbose)
<|endoftext|>