<commit_before>import os.path
import webbrowser
import re

try:
    from . import api, msg, utils, reactor, shared as G
    from .handlers import account, credentials
    from .. import editor
    from ..common.exc_fmt import str_e
except (ImportError, ValueError):
    from floo.common.exc_fmt import str_e
    from floo.common.handlers import account, credentials
    from floo.common import api, msg, utils, reactor, shared as G
    from floo import editor

class FlooUI(object):
    agent = None

    def _make_agent(self, owner, workspace, auth, created_workspace, d):
        """@returns new Agent()"""
        raise NotImplemented()

    def user_y_or_n(self, context, prompt, affirmation_txt, cb):
        """@returns True/False"""
        raise NotImplemented()

    def user_select(self, context, prompt, choices_big, choices_small, cb):
        """@returns (choice, index)"""
        raise NotImplemented()

    def user_charfield(self, context, prompt, initial, cb):
        """@returns String"""
        raise NotImplemented()

    def get_a_window(self, abs_path, cb):
        """opens a project in a window or something"""
        raise NotImplemented()

    @utils.inlined_callbacks
    def link_account(self, context, host, cb):
        prompt = 'No credentials found in ~/.floorc.json for %s.  Would you like to download them (opens a browser)?.' % host
        yes = yield self.user_y_or_n, context,  prompt, "Download"
        if not yes:
            return

        agent = credentials.RequestCredentialsHandler()
        if not agent:
            self.error_message('''A configuration error occured earlier. Please go to %s and sign up to use this plugin.\n
    We're really sorry. This should never happen.''' % host)
            return

        agent.once('end', cb)

        try:
            reactor.reactor.connect(agent, host, G.DEFAULT_PORT, True)
        except Exception as e:
            print(str_e(e))

    @utils.inlined_callbacks
    def create_or_link_account(self, context, host, force, cb):
        disable_account_creation = utils.get_persistent_data().get('disable_account_creation')

        if host != "floobits.com":
            self.link_account(context, host, cb)
            return

        if disable_account_creation and not force:
            print('We could not automatically create or link your floobits account. Please go to floobits.com and sign up to use this plugin.')
            return

        if not G.EXPERT_MODE:
            editor.message_dialog('Thank you for installing the Floobits plugin!\n\nLet\'s set up your editor to work with Floobits.')

        choices = [
            'Use an existing Floobits account',
            'Create a new Floobits account',
            'Cancel (see https://floobits.com/help/floorc)'
        ]

        (choice, index) = yield self.user_select, context, 'You need a Floobits account to use Floobits! Do you want to:', choices, None

        if index == -1 or index == 2:
            d = utils.get_persistent_data()
            if not d.get('disable_account_creation'):
                d['disable_account_creation'] = True
                utils.update_persistent_data(d)
                editor.message_dialog('''You can set up a Floobits account at any time under\n\nTools -> Floobits -> Setup''')
            cb(None)
            return

        agent = None
        if index == 0:
            agent = credentials.RequestCredentialsHandler()
        else:
            agent = account.CreateAccountHandler()

        agent.once('end', cb)

        try:
            reactor.reactor.connect(agent, host, G.DEFAULT_PORT, True)
        except Exception as e:
            print(str_e(e))

    @utils.inlined_callbacks
    def remote_connect(self, context, host, owner, workspace, d, get_bufs=False):
        G.PROJECT_PATH = os.path.realpath(d)
        try:
            utils.mkdir(os.path.dirname(G.PROJECT_PATH))
        except Exception as e:
            msg.error("Couldn't create directory", G.PROJECT_PATH, str_e(e))
            return

        auth = G.AUTH.get(host)
        if not auth:
            success = yield self.link_account, context, host
            if not success:
                return
            auth = G.AUTH.get(host)
            if not auth:
                msg.error("Something went really wrong.")
                return
        if self.agent:
            try:
                self.agent.stop()
            except:
                pass

        G.WORKSPACE_WINDOW = yield self.get_a_window, d
        self.agent = self._make_agent(owner, workspace, auth, get_bufs, d)
        reactor.reactor.connect(self.agent, host, G.DEFAULT_PORT, True)
        url = self.agent.workspace_url
        utils.add_workspace_to_persistent_json(owner, workspace, url, d)
        utils.update_recent_workspaces(url)

    @utils.inlined_callbacks
    def create_workspace(self, context, host, owner, name, perms, dir_to_share, cb):
        prompt = 'Workspace name: '

        api_args = {
            'name': name,
            'owner': owner,
        }

        if perms:
            api_args['perms'] = {"AnonymousUser": perms}

        while True:
            new_name = yield self.user_charfield, context, prompt, name
            name = new_name or name
            try:
                api_args['name'] = name
                r = api.create_workspace(host, api_args)
            except Exception as e:
                msg.error('Unable to create workspace ', str_e(e))
                editor.error_message('Unable to create workspace: %s' % str_e(e))
                return

            if r.code < 400:
                workspace_url = 'https://%s/%s/%s' % (host, owner, name)
                msg.log('Created workspace', workspace_url)
                self.remote_connect(context, host, owner, name, dir_to_share, True)
                return

            msg.error('Unable to create workspace', r.body)

            if r.code not in (400, 402, 409):
                try:
                    r.body = r.body['detail']
                except Exception:
                    pass
                editor.error_message('Unable to create workspace: %s' % r.body)
                return

            if r.code == 402:
                try:
                    r.body = r.body['detail']
                except Exception:
                    pass

                yes = yield self.user_y_or_n, context, '%s Open billing settings?' % r.body, "Yes"
                if yes:
                    webbrowser.open('https://%s/%s/settings#billing' % (host, owner))
                return

            if r.code == 400:
                name = re.sub('[^A-Za-z0-9_\-\.]', '-', name)
                prompt = 'Invalid name. Workspace names must match the regex [A-Za-z0-9_\-\.]. Choose another name: '
                continue

            prompt = 'Workspace %s/%s already exists. Choose another name: ' % (owner, name)

    def join_workspace_by_url(self, context, workspace_url, possible_dirs=None):
        try:
            d = utils.parse_url(workspace_url)
        except Exception as e:
            return editor.error_message(str_e(e))

        return self.join_workspace(context, d['host'], d['workspace'], d['owner'], possible_dirs)

    @utils.inlined_callbacks
    def join_workspace(self, context, host, name, owner, possible_dirs=None):
        utils.reload_settings()

        if not utils.can_auth():
            success = yield self.create_or_link_account, context, host, False
            if not success:
                return
            utils.reload_settings()

        possible_dirs = possible_dirs or []
        for d in possible_dirs:
            info = utils.read_floo_file(d)
            if not info:
                continue
            try:
                parsed_url = utils.parse_url(info['url'])
            except Exception:
                parsed_url = None

            if parsed_url and parsed_url['host'] == host and parsed_url['workspace'] == name and parsed_url['owner'] == owner:
                self.remote_connect(context, host, owner, name, d)
                return

        try:
            d = utils.get_persistent_data()['workspaces'][owner][name]['path']
        except Exception:
            d = ''

        if d and os.path.isdir(d):
            self.remote_connect(context, host, owner, name, d)
            return

        d = d or os.path.join(G.SHARE_DIR or G.BASE_DIR, owner, name)
        while True:
            d = yield self.user_charfield, context, 'Save workspace files to: ', d
            if not d:
                return
            d = os.path.realpath(os.path.expanduser(d))
            if not os.path.isdir(d):
                y_or_n = yield self.user_y_or_n, context, '%s is not a directory. Create it? ' % d, "Create Directory"
                if not y_or_n:
                    return
                utils.mkdir(d)
                if not os.path.isdir(d):
                    msg.error("Couldn't create directory", d)
                    continue
            if os.path.isdir(d):
                self.remote_connect(context, host, owner, name, d)
                return

    @utils.inlined_callbacks
    def share_dir(self, context, dir_to_share, ask_about_dir, perms):
        utils.reload_settings()
        if not utils.can_auth():
            success = yield self.create_or_link_account, context, G.DEFAULT_HOST, False
            if not success:
                return
            utils.reload_settings()

        if not dir_to_share:
            if not ask_about_dir:
                raise ValueError("I need a thing to share.")
            dir_to_share = yield self.user_charfield, 'Directory to share:', ask_about_dir

        if not dir_to_share:
            return

        dir_to_share = os.path.expanduser(dir_to_share)
        dir_to_share = os.path.realpath(dir_to_share)
        dir_to_share = utils.unfuck_path(dir_to_share)

        if os.path.isfile(dir_to_share):
            dir_to_share = os.path.dirname(dir_to_share)

        workspace_name = os.path.basename(dir_to_share)
        msg.debug('', workspace_name, dir_to_share)

        if os.path.isfile(dir_to_share):
            # file_to_share = dir_to_share
            dir_to_share = os.path.dirname(dir_to_share)

        try:
            utils.mkdir(dir_to_share)
        except Exception:
            msg.error("The directory", dir_to_share, "doesn't exist and I can't create it.")
            return

        info = utils.read_floo_file(dir_to_share)

        workspace_url = info.get('url')
        if workspace_url:
            try:
                parsed_url = api.prejoin_workspace(workspace_url, dir_to_share, {'perms': perms})
            except ValueError as e:
                pass
            if parsed_url:
                # TODO: make sure we create_flooignore
                self.remote_connect(context, parsed_url['host'], parsed_url['owner'], parsed_url['workspace'], dir_to_share)
                return

        def prejoin(workspace_url):
            try:
                return api.prejoin_workspace(workspace_url, dir_to_share, {'perms': perms})
            except ValueError:
                pass

        parsed_url = utils.get_workspace_by_path(dir_to_share, prejoin)
        if parsed_url:
            self.remote_connect(context, parsed_url['host'], parsed_url['owner'], parsed_url['workspace'], dir_to_share)
            return

        if not G.AUTH:
            msg.warn('no auth')
            return

        hosts = list(G.AUTH.keys())
        if len(hosts) == 1:
            host = hosts[0]
        else:
            (host, index) = yield self.user_select, context, 'Which Floobits account should be used?', hosts, None
            if not host:
                return

        try:
            r = api.get_orgs_can_admin(host)
        except IOError as e:
            editor.error_message('Error getting org list: %s' % str_e(e))
            return

        # orgs = [[org['name'], 'Create workspace owned by %s' % org['name']] for org in r.body]
        # orgs.insert(0, [username, 'Create workspace owned by %s' % username])

        choices = [G.AUTH[host]['username']]
        if r.code >= 400:
            editor.error_message('Error getting org list: %s' % r.body)
        elif r.body:
            choices += [org['name'] for org in r.body]

        if len(choices) == 1:
            owner = choices[0]
        else:
            (owner, index) = yield self.user_select, context, 'Create workspace owned by', choices, None

        yield self.create_workspace, context, host, owner, workspace_name, perms, dir_to_share
<commit_msg>clean up dialogs and fix perms<commit_after>import os.path
import webbrowser
import re

try:
    from . import api, msg, utils, reactor, shared as G
    from .handlers import account, credentials
    from .. import editor
    from ..common.exc_fmt import str_e
except (ImportError, ValueError):
    from floo.common.exc_fmt import str_e
    from floo.common.handlers import account, credentials
    from floo.common import api, msg, utils, reactor, shared as G
    from floo import editor

class FlooUI(object):
    agent = None

    def _make_agent(self, owner, workspace, auth, created_workspace, d):
        """@returns new Agent()"""
        raise NotImplemented()

    def user_y_or_n(self, context, prompt, affirmation_txt, cb):
        """@returns True/False"""
        raise NotImplemented()

    def user_select(self, context, prompt, choices_big, choices_small, cb):
        """@returns (choice, index)"""
        raise NotImplemented()

    def user_charfield(self, context, prompt, initial, cb):
        """@returns String"""
        raise NotImplemented()

    def get_a_window(self, abs_path, cb):
        """opens a project in a window or something"""
        raise NotImplemented()

    @utils.inlined_callbacks
    def link_account(self, context, host, cb):
        prompt = 'No credentials found in ~/.floorc.json for %s.  Would you like to download them (opens a browser)?.' % host
        yes = yield self.user_y_or_n, context,  prompt, "Download"
        if not yes:
            return

        agent = credentials.RequestCredentialsHandler()
        if not agent:
            self.error_message('''A configuration error occured earlier. Please go to %s and sign up to use this plugin.\n
    We're really sorry. This should never happen.''' % host)
            return

        agent.once('end', cb)

        try:
            reactor.reactor.connect(agent, host, G.DEFAULT_PORT, True)
        except Exception as e:
            print(str_e(e))

    @utils.inlined_callbacks
    def create_or_link_account(self, context, host, force, cb):
        disable_account_creation = utils.get_persistent_data().get('disable_account_creation')

        if host != "floobits.com":
            self.link_account(context, host, cb)
            return

        if disable_account_creation and not force:
            print('We could not automatically create or link your floobits account. Please go to floobits.com and sign up to use this plugin.')
            return

        if not G.EXPERT_MODE:
            editor.message_dialog('Thank you for installing the Floobits plugin!\n\nLet\'s set up your editor to work with Floobits.')

        choices = [
            'Use an existing Floobits account',
            'Create a new Floobits account',
            'Cancel (see https://floobits.com/help/floorc)'
        ]

        (choice, index) = yield self.user_select, context, 'You need a Floobits account to use Floobits! Do you want to:', choices, None

        if index == -1 or index == 2:
            d = utils.get_persistent_data()
            if not d.get('disable_account_creation'):
                d['disable_account_creation'] = True
                utils.update_persistent_data(d)
                editor.message_dialog('''You can set up a Floobits account at any time under\n\nTools -> Floobits -> Setup''')
            cb(None)
            return

        agent = None
        if index == 0:
            agent = credentials.RequestCredentialsHandler()
        else:
            agent = account.CreateAccountHandler()

        agent.once('end', cb)

        try:
            reactor.reactor.connect(agent, host, G.DEFAULT_PORT, True)
        except Exception as e:
            print(str_e(e))

    @utils.inlined_callbacks
    def remote_connect(self, context, host, owner, workspace, d, get_bufs=False):
        G.PROJECT_PATH = os.path.realpath(d)
        try:
            utils.mkdir(os.path.dirname(G.PROJECT_PATH))
        except Exception as e:
            msg.error("Couldn't create directory", G.PROJECT_PATH, str_e(e))
            return

        auth = G.AUTH.get(host)
        if not auth:
            success = yield self.link_account, context, host
            if not success:
                return
            auth = G.AUTH.get(host)
            if not auth:
                msg.error("Something went really wrong.")
                return
        if self.agent:
            try:
                self.agent.stop()
            except:
                pass

        G.WORKSPACE_WINDOW = yield self.get_a_window, d
        self.agent = self._make_agent(owner, workspace, auth, get_bufs, d)
        reactor.reactor.connect(self.agent, host, G.DEFAULT_PORT, True)
        url = self.agent.workspace_url
        utils.add_workspace_to_persistent_json(owner, workspace, url, d)
        utils.update_recent_workspaces(url)

    @utils.inlined_callbacks
    def create_workspace(self, context, host, owner, name, api_args, dir_to_share, cb):
        prompt = 'Workspace name: '

        api_args['name'] = name
        api_args['owner'] = owner,

        while True:
            new_name = yield self.user_charfield, context, prompt, name
            name = new_name or name
            try:
                api_args['name'] = name
                r = api.create_workspace(host, api_args)
            except Exception as e:
                msg.error('Unable to create workspace ', str_e(e))
                editor.error_message('Unable to create workspace: %s' % str_e(e))
                return

            if r.code < 400:
                workspace_url = 'https://%s/%s/%s' % (host, owner, name)
                msg.log('Created workspace', workspace_url)
                self.remote_connect(context, host, owner, name, dir_to_share, True)
                return

            msg.error('Unable to create workspace', r.body)

            if r.code not in (400, 402, 409):
                try:
                    r.body = r.body['detail']
                except Exception:
                    pass
                editor.error_message('Unable to create workspace: %s' % r.body)
                return

            if r.code == 402:
                try:
                    r.body = r.body['detail']
                except Exception:
                    pass

                yes = yield self.user_y_or_n, context, '%s Open billing settings?' % r.body, "Yes"
                if yes:
                    webbrowser.open('https://%s/%s/settings#billing' % (host, owner))
                return

            if r.code == 400:
                name = re.sub('[^A-Za-z0-9_\-\.]', '-', name)
                prompt = 'Invalid name. Workspace names must match the regex [A-Za-z0-9_\-\.]. Choose another name: '
                continue

            prompt = 'Workspace %s/%s already exists. Choose another name: ' % (owner, name)

    def join_workspace_by_url(self, context, workspace_url, possible_dirs=None):
        try:
            d = utils.parse_url(workspace_url)
        except Exception as e:
            return editor.error_message(str_e(e))

        return self.join_workspace(context, d['host'], d['workspace'], d['owner'], possible_dirs)

    @utils.inlined_callbacks
    def join_workspace(self, context, host, name, owner, possible_dirs=None):
        utils.reload_settings()

        if not utils.can_auth():
            success = yield self.create_or_link_account, context, host, False
            if not success:
                return
            utils.reload_settings()

        possible_dirs = possible_dirs or []
        for d in possible_dirs:
            info = utils.read_floo_file(d)
            if not info:
                continue
            try:
                parsed_url = utils.parse_url(info['url'])
            except Exception:
                parsed_url = None

            if parsed_url and parsed_url['host'] == host and parsed_url['workspace'] == name and parsed_url['owner'] == owner:
                self.remote_connect(context, host, owner, name, d)
                return

        try:
            d = utils.get_persistent_data()['workspaces'][owner][name]['path']
        except Exception:
            d = ''

        if d and os.path.isdir(d):
            self.remote_connect(context, host, owner, name, d)
            return

        d = d or os.path.join(G.SHARE_DIR or G.BASE_DIR, owner, name)
        while True:
            d = yield self.user_charfield, context, 'Save workspace files to: ', d
            if not d:
                return
            d = os.path.realpath(os.path.expanduser(d))
            if not os.path.isdir(d):
                y_or_n = yield self.user_y_or_n, context, '%s is not a directory. Create it? ' % d, "Create Directory"
                if not y_or_n:
                    return
                utils.mkdir(d)
                if not os.path.isdir(d):
                    msg.error("Couldn't create directory", d)
                    continue
            if os.path.isdir(d):
                self.remote_connect(context, host, owner, name, d)
                return

    @utils.inlined_callbacks
    def share_dir(self, context, dir_to_share, ask_about_dir, api_args):
        utils.reload_settings()
        if not utils.can_auth():
            success = yield self.create_or_link_account, context, G.DEFAULT_HOST, False
            if not success:
                return
            utils.reload_settings()

        if not dir_to_share:
            if not ask_about_dir:
                raise ValueError("I need a thing to share.")
            dir_to_share = yield self.user_charfield, 'Directory to share:', ask_about_dir

        if not dir_to_share:
            return

        dir_to_share = os.path.expanduser(dir_to_share)
        dir_to_share = os.path.realpath(dir_to_share)
        dir_to_share = utils.unfuck_path(dir_to_share)

        if os.path.isfile(dir_to_share):
            dir_to_share = os.path.dirname(dir_to_share)

        workspace_name = os.path.basename(dir_to_share)
        msg.debug('', workspace_name, dir_to_share)

        if os.path.isfile(dir_to_share):
            # file_to_share = dir_to_share
            dir_to_share = os.path.dirname(dir_to_share)

        try:
            utils.mkdir(dir_to_share)
        except Exception:
            msg.error("The directory", dir_to_share, "doesn't exist and I can't create it.")
            return

        info = utils.read_floo_file(dir_to_share)

        def prejoin(workspace_url):
            try:
                return api.prejoin_workspace(workspace_url, dir_to_share, api_args)
            except ValueError:
                pass

        workspace_url = info.get('url')
        if workspace_url:
            parsed_url = prejoin(workspace_url)
            if parsed_url:
                # TODO: make sure we create_flooignorei
                self.remote_connect(context, parsed_url['host'], parsed_url['owner'], parsed_url['workspace'], dir_to_share)
                return

        parsed_url = utils.get_workspace_by_path(dir_to_share, prejoin)
        if parsed_url:
            self.remote_connect(context, parsed_url['host'], parsed_url['owner'], parsed_url['workspace'], dir_to_share)
            return

        if not G.AUTH:
            msg.warn('no auth')
            return

        hosts = list(G.AUTH.keys())
        if len(hosts) == 1:
            host = hosts[0]
        else:
            little = ["%s on %s" % (a['username'], h) for h, a in G.AUTH.items()]
            (host, index) = yield self.user_select, context, 'Which Floobits account should be used?', hosts, little
            if not host:
                return

        try:
            r = api.get_orgs_can_admin(host)
        except IOError as e:
            editor.error_message('Error getting org list: %s' % str_e(e))
            return

        # orgs = [[org['name'], 'Create workspace owned by %s' % org['name']] for org in r.body]
        # orgs.insert(0, [username, 'Create workspace owned by %s' % username])

        choices = [G.AUTH[host]['username']]
        if r.code >= 400:
            editor.error_message('Error getting org list: %s' % r.body)
        elif r.body:
            choices += [org['name'] for org in r.body]

        if len(choices) == 1:
            owner = choices[0]
        else:
            little = ['Create workspace owned by %s' % s for s in choices]
            (owner, index) = yield self.user_select, context, 'Create workspace owned by', choices, little

        yield self.create_workspace, context, host, owner, workspace_name, api_args, dir_to_share
<|endoftext|>


<commit_before>import uuid, json
from esprit import raw, util
from copy import deepcopy

class StoreException(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class DAO(object):

    @classmethod
    def makeid(cls):
        return uuid.uuid4().hex

    def actions(self, conn, action_queue):
        for action in action_queue:
            if action.keys()[0] == "remove":
                self._action_remove(conn, action)
            elif action.keys()[0] == "store":
                self._action_store(conn, action)

    def _action_remove(self, conn, remove_action):
        obj = remove_action.get("remove")
        if "index" not in obj:
            raise StoreException("no index provided for remove action")
        if "id" not in obj and "query" not in obj:
            raise StoreException("no id or query provided for remove action")
        if "id" in obj:
            raw.delete(conn, obj.get("index"), obj.get("id"))
        elif "query" in obj:
            raw.delete_by_query(conn, obj.get("index"), obj.get("query"))

    def _action_store(self, conn, store_action):
        obj = store_action.get("store")
        if "index" not in obj:
            raise StoreException("no index provided for store action")
        if "record" not in obj:
            raise StoreException("no record provided for store action")
        raw.store(conn, obj.get("index"), obj.get("record"), obj.get("id"))

class DomainObject(DAO):
    __type__ = None
    __conn__ = None

    def __init__(self, raw=None):
        self.data = raw if raw is not None else {}

    @property
    def id(self):
        return self.data.get('id', None)

    @id.setter
    def id(self, val):
        self.data["id"] = val

    @property
    def json(self):
        return json.dumps(self.data)

    @property
    def raw(self):
        return self.data

    @classmethod
    def refresh(cls, conn=None):
        if conn is None:
            conn = cls.__conn__
        raw.refresh(conn)

    @classmethod
    def pull(cls, id_, conn=None):
        '''Retrieve object by id.'''
        if conn is None:
            conn = cls.__conn__

        if id_ is None:
            return None
        try:
            resp = raw.get(conn, cls.__type__, id_)
            if resp.status_code == 404:
                return None
            else:
                j = raw.unpack_get(resp)
                return cls(j)
        except:
            return None

    @classmethod
    def query(cls, q='', terms=None, facets=None, conn=None, **kwargs):
        '''Perform a query on backend.

        :param q: maps to query_string parameter if string, or query dict if dict.
        :param terms: dictionary of terms to filter on. values should be lists.
        :param facets: dict of facets to return from the query.
        :param kwargs: any keyword args as per
            http://www.elasticsearch.org/guide/reference/api/search/uri-request.html
        '''
        if conn is None:
            conn = cls.__conn__

        if isinstance(q,dict):
            query = q
            if 'bool' not in query['query']:
                boolean = {'bool':{'must': [] }}
                boolean['bool']['must'].append( query['query'] )
                query['query'] = boolean
            if 'must' not in query['query']['bool']:
                query['query']['bool']['must'] = []
        elif q:
            query = {
                'query': {
                    'bool': {
                        'must': [
                            {'query_string': { 'query': q }}
                        ]
                    }
                }
            }
        else:
            query = {
                'query': {
                    'bool': {
                        'must': [
                            {'match_all': {}}
                        ]
                    }
                }
            }

        if facets:
            if 'facets' not in query:
                query['facets'] = {}
            for k, v in facets.items():
                query['facets'][k] = {"terms":v}

        if terms:
            boolean = {'must': [] }
            for term in terms:
                if not isinstance(terms[term],list): terms[term] = [terms[term]]
                for val in terms[term]:
                    obj = {'term': {}}
                    obj['term'][ term ] = val
                    boolean['must'].append(obj)
            if q and not isinstance(q,dict):
                boolean['must'].append( {'query_string': { 'query': q } } )
            elif q and 'query' in q:
                boolean['must'].append( query['query'] )
            query['query'] = {'bool': boolean}

        for k,v in kwargs.items():
            if k == '_from':
                query['from'] = v
            else:
                query[k] = v

        r = raw.search(conn, cls.__type__, query)
        # r = requests.post(cls.target() + recid + endpoint, data=json.dumps(query))
        return r.json()

    def save(self, conn=None, makeid=True, created=True, updated=True):
        if makeid:
            if "id" not in self.data:
                self.id = self.makeid()
        if created:
            if 'created_date' not in self.data:
                self.data['created_date'] = util.now()
        if updated:
            self.data['last_updated'] = util.now()

        if conn is None:
            conn = self.__conn__
        raw.store(conn, self.__type__, self.data, self.id)

    def delete(self, conn=None):
        if conn is None:
            conn = self.__conn__
        raw.delete(conn, self.__type__, self.id)

    @classmethod
    def iterate(cls, q, page_size=1000, limit=None, wrap=True):
        q["size"] = page_size
        q["from"] = 0
        if "sort" not in q: # to ensure complete coverage on a changing index, sort by id is our best bet
            q["sort"] = [{"id" : {"order" : "asc"}}]
        counter = 0
        while True:
            # apply the limit
            if limit is not None and counter >= limit:
                break

            res = cls.query(q=q)
            rs = [r.get("_source") if "_source" in r else r.get("fields") for r in res.get("hits", {}).get("hits", [])]
            # print counter, len(rs), res.get("hits", {}).get("total"), len(res.get("hits", {}).get("hits", [])), json.dumps(q)
            if len(rs) == 0:
                break
            for r in rs:
                # apply the limit (again)
                if limit is not None and counter >= limit:
                    break
                counter += 1
                if wrap:
                    yield cls(r)
                else:
                    yield r
            q["from"] += page_size

    @classmethod
    def iterall(cls, page_size=1000, limit=None):
        return cls.iterate(deepcopy(all_query), page_size, limit)

########################################################################
## Some useful ES queries
########################################################################

all_query = {
    "query" : {
        "match_all" : { }
    }
}<commit_msg>Add support for should_terms on query on domain object<commit_after>import uuid, json
from esprit import raw, util
from copy import deepcopy

class StoreException(Exception):
    def __init__(self, value):
        self.value = value
    def __str__(self):
        return repr(self.value)

class DAO(object):

    @classmethod
    def makeid(cls):
        return uuid.uuid4().hex

    def actions(self, conn, action_queue):
        for action in action_queue:
            if action.keys()[0] == "remove":
                self._action_remove(conn, action)
            elif action.keys()[0] == "store":
                self._action_store(conn, action)

    def _action_remove(self, conn, remove_action):
        obj = remove_action.get("remove")
        if "index" not in obj:
            raise StoreException("no index provided for remove action")
        if "id" not in obj and "query" not in obj:
            raise StoreException("no id or query provided for remove action")
        if "id" in obj:
            raw.delete(conn, obj.get("index"), obj.get("id"))
        elif "query" in obj:
            raw.delete_by_query(conn, obj.get("index"), obj.get("query"))

    def _action_store(self, conn, store_action):
        obj = store_action.get("store")
        if "index" not in obj:
            raise StoreException("no index provided for store action")
        if "record" not in obj:
            raise StoreException("no record provided for store action")
        raw.store(conn, obj.get("index"), obj.get("record"), obj.get("id"))

class DomainObject(DAO):
    __type__ = None
    __conn__ = None

    def __init__(self, raw=None):
        self.data = raw if raw is not None else {}

    @property
    def id(self):
        return self.data.get('id', None)

    @id.setter
    def id(self, val):
        self.data["id"] = val

    @property
    def json(self):
        return json.dumps(self.data)

    @property
    def raw(self):
        return self.data

    @classmethod
    def refresh(cls, conn=None):
        if conn is None:
            conn = cls.__conn__
        raw.refresh(conn)

    @classmethod
    def pull(cls, id_, conn=None):
        '''Retrieve object by id.'''
        if conn is None:
            conn = cls.__conn__

        if id_ is None:
            return None
        try:
            resp = raw.get(conn, cls.__type__, id_)
            if resp.status_code == 404:
                return None
            else:
                j = raw.unpack_get(resp)
                return cls(j)
        except:
            return None

    @classmethod
    def query(cls, q='', terms=None, should_terms=None, facets=None, conn=None, **kwargs):
        '''Perform a query on backend.

        :param q: maps to query_string parameter if string, or query dict if dict.
        :param terms: dictionary of terms to filter on. values should be lists.
        :param facets: dict of facets to return from the query.
        :param kwargs: any keyword args as per
            http://www.elasticsearch.org/guide/reference/api/search/uri-request.html
        '''
        if conn is None:
            conn = cls.__conn__

        if isinstance(q,dict):
            query = q
            if 'bool' not in query['query']:
                boolean = {'bool':{'must': [] }}
                boolean['bool']['must'].append( query['query'] )
                query['query'] = boolean
            if 'must' not in query['query']['bool']:
                query['query']['bool']['must'] = []
        elif q:
            query = {
                'query': {
                    'bool': {
                        'must': [
                            {'query_string': { 'query': q }}
                        ]
                    }
                }
            }
        else:
            query = {
                'query': {
                    'bool': {
                        'must': [
                            {'match_all': {}}
                        ]
                    }
                }
            }

        if facets:
            if 'facets' not in query:
                query['facets'] = {}
            for k, v in facets.items():
                query['facets'][k] = {"terms":v}

        if terms:
            boolean = {'must': [] }
            for term in terms:
                if not isinstance(terms[term],list): terms[term] = [terms[term]]
                for val in terms[term]:
                    obj = {'term': {}}
                    obj['term'][ term ] = val
                    boolean['must'].append(obj)
            if q and not isinstance(q,dict):
                boolean['must'].append( {'query_string': { 'query': q } } )
            elif q and 'query' in q:
                boolean['must'].append( query['query'] )
            query['query'] = {'bool': boolean}

        for k,v in kwargs.items():
            if k == '_from':
                query['from'] = v
            else:
                query[k] = v

        if should_terms is not None and len(should_terms) > 0:
            for s in should_terms:
                if not isinstance(should_terms[s],list): should_terms[s] = [should_terms[s]]
                query["query"]["bool"]["must"].append({"terms" : {s : should_terms[s]}})

        r = raw.search(conn, cls.__type__, query)
        return r.json()

    def save(self, conn=None, makeid=True, created=True, updated=True):
        if makeid:
            if "id" not in self.data:
                self.id = self.makeid()
        if created:
            if 'created_date' not in self.data:
                self.data['created_date'] = util.now()
        if updated:
            self.data['last_updated'] = util.now()

        if conn is None:
            conn = self.__conn__
        raw.store(conn, self.__type__, self.data, self.id)

    def delete(self, conn=None):
        if conn is None:
            conn = self.__conn__
        raw.delete(conn, self.__type__, self.id)

    @classmethod
    def iterate(cls, q, page_size=1000, limit=None, wrap=True):
        q["size"] = page_size
        q["from"] = 0
        if "sort" not in q: # to ensure complete coverage on a changing index, sort by id is our best bet
            q["sort"] = [{"id" : {"order" : "asc"}}]
        counter = 0
        while True:
            # apply the limit
            if limit is not None and counter >= limit:
                break

            res = cls.query(q=q)
            rs = [r.get("_source") if "_source" in r else r.get("fields") for r in res.get("hits", {}).get("hits", [])]
            # print counter, len(rs), res.get("hits", {}).get("total"), len(res.get("hits", {}).get("hits", [])), json.dumps(q)
            if len(rs) == 0:
                break
            for r in rs:
                # apply the limit (again)
                if limit is not None and counter >= limit:
                    break
                counter += 1
                if wrap:
                    yield cls(r)
                else:
                    yield r
            q["from"] += page_size

    @classmethod
    def iterall(cls, page_size=1000, limit=None):
        return cls.iterate(deepcopy(all_query), page_size, limit)

########################################################################
## Some useful ES queries
########################################################################

all_query = {
    "query" : {
        "match_all" : { }
    }
}<|endoftext|>


<commit_before>#!/usr/bin/env python
"""
Fourmi, a web scraper build to search specific information for a given compound (and it's pseudonyms).

Usage:
    fourmi search <compound>
    fourmi [options] search <compound>
    fourmi list
    fourmi -h | --help
    fourmi --version

Options:
    -h --help                       Show this screen.
    --version                       Show version.
    --verbose                       Verbose logging output.
    --log=<file>                    Save log to an file.
    -o <file> --output=<file>       Output file [default: result.*format*]
    -f <format> --format=<format>   Output formats (supported: csv, json, jsonlines, xml) [default: jsonlines]
    --include=<sourcenames>         Include only sources that match the regular these expressions split by a comma.
    --exclude=<sourcenames>         Exclude the sources that match the regular these expressions split by a comma.
"""

from twisted.internet import reactor
from scrapy.crawler import Crawler
from scrapy import log, signals
from scrapy.utils.project import get_project_settings
import docopt

from FourmiCrawler.spider import FourmiSpider
from sourceloader import SourceLoader

def setup_crawler(searchable, settings, loader):
    spider = FourmiSpider(compound=searchable)
    spider.add_parsers(loader.sources)
    crawler = Crawler(settings)
    crawler.signals.connect(reactor.stop, signal=signals.spider_closed)
    crawler.configure()
    crawler.crawl(spider)
    crawler.start()

def scrapy_settings_manipulation(arguments):
    settings = get_project_settings()

    if arguments["--output"] != 'result.*format*':
        settings.overrides["FEED_URI"] = arguments["--output"]
    elif arguments["--format"] == "jsonlines":
        settings.overrides["FEED_URI"] = "results.json"
    elif arguments["--format"] is not None:
        settings.overrides["FEED_URI"] = "results." + arguments["--format"]

    if arguments["--format"] is not None:
        settings.overrides["FEED_FORMAT"] = arguments["--format"]

    return settings

def start_log(arguments):
    if arguments["--log"] is not None:
        if arguments["--verbose"]:
            log.start(logfile=arguments["--log"], logstdout=False, loglevel=log.DEBUG)
        else:
            log.start(logfile=arguments["--log"], logstdout=True, loglevel=log.WARNING)
    else:
        if arguments["--verbose"]:
            log.start(logstdout=False, loglevel=log.DEBUG)
        else:
            log.start(logstdout=True, loglevel=log.WARNING)

def search(arguments, loader):
    start_log(arguments)
    settings = scrapy_settings_manipulation(arguments)
    setup_crawler(arguments["<compound>"], settings, loader)
    reactor.run()

if __name__ == '__main__':
    arguments = docopt.docopt(__doc__, version='Fourmi - V0.1.0')
    loader = SourceLoader()

    if arguments["search"]:
        search(arguments, loader)
    elif arguments["list"]:
        print "-== Available Sources ==-"
        print str(loader)
<commit_msg>Added the right commands to include and exclude a source<commit_after>#!/usr/bin/env python
"""
Fourmi, a web scraper build to search specific information for a given compound (and it's pseudonyms).

Usage:
    fourmi search <compound>
    fourmi [options] search <compound>
    fourmi list
    fourmi [--include=<sourcename> | --exclude=<sourcename>] list
    fourmi -h | --help
    fourmi --version

Options:
    -h --help                       Show this screen.
    --version                       Show version.
    --verbose                       Verbose logging output.
    --log=<file>                    Save log to an file.
    -o <file> --output=<file>       Output file [default: result.*format*]
    -f <format> --format=<format>   Output formats (supported: csv, json, jsonlines, xml) [default: jsonlines]
    --include=<sourcenames>         Include only sources that match the regular these expressions split by a comma.
    --exclude=<sourcenames>         Exclude the sources that match the regular these expressions split by a comma.
"""

from twisted.internet import reactor
from scrapy.crawler import Crawler
from scrapy import log, signals
from scrapy.utils.project import get_project_settings
import docopt

from FourmiCrawler.spider import FourmiSpider
from sourceloader import SourceLoader

def setup_crawler(searchable, settings, loader):
    spider = FourmiSpider(compound=searchable)
    spider.add_parsers(loader.sources)
    crawler = Crawler(settings)
    crawler.signals.connect(reactor.stop, signal=signals.spider_closed)
    crawler.configure()
    crawler.crawl(spider)
    crawler.start()

def scrapy_settings_manipulation(arguments):
    settings = get_project_settings()

    if arguments["--output"] != 'result.*format*':
        settings.overrides["FEED_URI"] = arguments["--output"]
    elif arguments["--format"] == "jsonlines":
        settings.overrides["FEED_URI"] = "results.json"
    elif arguments["--format"] is not None:
        settings.overrides["FEED_URI"] = "results." + arguments["--format"]

    if arguments["--format"] is not None:
        settings.overrides["FEED_FORMAT"] = arguments["--format"]

    return settings

def start_log(arguments):
    if arguments["--log"] is not None:
        if arguments["--verbose"]:
            log.start(logfile=arguments["--log"], logstdout=False, loglevel=log.DEBUG)
        else:
            log.start(logfile=arguments["--log"], logstdout=True, loglevel=log.WARNING)
    else:
        if arguments["--verbose"]:
            log.start(logstdout=False, loglevel=log.DEBUG)
        else:
            log.start(logstdout=True, loglevel=log.WARNING)

def search(arguments, loader):
    start_log(arguments)
    settings = scrapy_settings_manipulation(arguments)
    setup_crawler(arguments["<compound>"], settings, loader)
    reactor.run()

if __name__ == '__main__':
    arguments = docopt.docopt(__doc__, version='Fourmi - V0.1.0')
    loader = SourceLoader()

    if arguments["--include"]:
        loader.include(arguments["--include"].split(','))
    elif arguments["--exclude"]:
        loader.exclude(arguments["--exclude"].split(','))

    if arguments["search"]:
        search(arguments, loader)
    elif arguments["list"]:
        print "-== Available Sources ==-"
        print str(loader)
<|endoftext|>


<commit_before>#!/usr/bin/python

# Converts open-ephys .kwd files to raw binary, a format ready
# for spike sorting with the Phy spike sorting program

from __future__ import division, print_function, unicode_literals
import os.path
import numpy as np
from openephys.kwik import load, load_all
from openephys.klustafiles import _prm_file, _prb_group

# number of data points to write at a time, prevents excess memory usage
BUFFERSIZE = 131072

def write_binary(filename, data, channels=()):
    if not data.dtype == np.int16:
        raise TypeError("data should be of type int16")
    if not channels:
        channels = range(data.shape[1])
    n_rows = int(BUFFERSIZE / len(channels))
    with open(filename, "wb") as f:
        for i in range(int(data.shape[0] / n_rows) + 1):
            index = i * n_rows
            buffer = data[index:index + n_rows, channels]
            f.write(buffer.tobytes())

def main(filename, channels=(), recordings=()):
    if recordings:
        all_data = [load(filename, x) for x in recordings]
    else:
        all_data = load_all(filename)
    sampling_rate = all_data[0]["info"]["sample_rate"]
    if channels:
        n_channels = len(channels)
    else:
        n_channels = all_data[0]["data"].shape[1]
    root_filename = os.path.splitext(filename)[0]
    dat_filenames = []
    for group_i, data in enumerate(all_data):
        binary_filename = "{fname}_{i}.dat".format(fname=root_filename,
                                                   i=group_i)
        dat_filenames.append(binary_filename)
        write_binary(binary_filename, data["data"], channels)

    prb_filename = root_filename + ".prb"
    dummy_prb_file(prb_filename, n_channels)
    prm_filename = root_filename + ".prm"
    dummy_prm_file(prm_filename, prb_filename, dat_filenames, sampling_rate,
                   n_channels)

if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser(
        prog="kwik2dat",
        description="""Converts Open-Ephys .kwd files to raw binary
        for use with Phy or Neuroscope""")
    p.add_argument("kwdfile", help="input .kwd file")
    p.add_argument("-c",
                   "--channels",
                   help="Channels to extract (0-indexed), all by default",
                   nargs="+",
                   type=int)
    p.add_argument("-r",
                   "--recordings",
                   help="Recordings to extract, all by default",
                   nargs="+",
                   type=int)
    p.add_argument("--klustafiles",
                   help="creates template files for klustakwik/phy")
    options = p.parse_args()
    main(options.kwdfile, options.channels, options.recordings,
         options.klustafiles)
<commit_msg>fixed klustafiles flag<commit_after>#!/usr/bin/python

# Converts open-ephys .kwd files to raw binary, a format ready
# for spike sorting with the Phy spike sorting program

from __future__ import division, print_function, unicode_literals
import os.path
import numpy as np
from openephys.kwik import load, load_all
from openephys.klustafiles import _prm_file, _prb_group

# number of data points to write at a time, prevents excess memory usage
BUFFERSIZE = 131072

def write_binary(filename, data, channels=()):
    if not data.dtype == np.int16:
        raise TypeError("data should be of type int16")
    if not channels:
        channels = range(data.shape[1])
    n_rows = int(BUFFERSIZE / len(channels))
    with open(filename, "wb") as f:
        for i in range(int(data.shape[0] / n_rows) + 1):
            index = i * n_rows
            buffer = data[index:index + n_rows, channels]
            f.write(buffer.tobytes())

def main(filename, channels=(), recordings=(), klustafiles=False):
    if recordings:
        all_data = [load(filename, x) for x in recordings]
    else:
        all_data = load_all(filename)
    sampling_rate = all_data[0]["info"]["sample_rate"]
    if channels:
        n_channels = len(channels)
    else:
        n_channels = all_data[0]["data"].shape[1]
    root_filename = os.path.splitext(filename)[0]
    dat_filenames = []
    for group_i, data in enumerate(all_data):
        binary_filename = "{fname}_{i}.dat".format(fname=root_filename,
                                                   i=group_i)
        dat_filenames.append(binary_filename)
        write_binary(binary_filename, data["data"], channels)

    if klustafiles:
        prb_filename = root_filename + ".prb"
        dummy_prb_file(prb_filename, n_channels)
        prm_filename = root_filename + ".prm"
        dummy_prm_file(prm_filename, prb_filename, dat_filenames, sampling_rate,
                    n_channels)

if __name__ == "__main__":
    import argparse
    p = argparse.ArgumentParser(
        prog="kwik2dat",
        description="""Converts Open-Ephys .kwd files to raw binary
        for use with Phy or Neuroscope""")
    p.add_argument("kwdfile", help="input .kwd file")
    p.add_argument("-c",
                   "--channels",
                   help="Channels to extract (0-indexed), all by default",
                   nargs="+",
                   type=int)
    p.add_argument("-r",
                   "--recordings",
                   help="Recordings to extract, all by default",
                   nargs="+",
                   type=int)
    p.add_argument("--klustafiles",
                   help="creates template files for klustakwik/phy EXPERIMENTAL",
                   action="store_true")
    options = p.parse_args()
    main(options.kwdfile, options.channels, options.recordings,
         options.klustafiles)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#-*- coding: utf-8 -*-

#This software is distributed under the Creative Commons license (CC0) version 1.0. A copy of this license should have been distributed with this software.
#The license can also be read online: <https://creativecommons.org/publicdomain/zero/1.0/>. If this online license differs from the license provided with this software, the license provided with this software should be applied.

"""
Defines a dictionary-like data structure where all values are validated.

The purpose of this dictionary is to be used in configuration data structures
where the keys are more or less frozen after the application has completed its
start-up phase.
"""

import collections #For namedtuple.

import luna.plugins #To access the loggers.

class ValidatedDictionary(dict):
	"""
	A dictionary-like data structure where all values are validated.

	New keys can only be added by also providing a validator predicate. Adding
	new keys can't be done with the usual ``dictionary[key] = value``, but via a
	separate method ``add``. Adding new keys in the usual way would result in an
	exception stating that the key cannot be found.

	When setting a key to a new value, the predicate associated with the key
	will be executed to determine whether the value is allowed. If the value is
	not allowed, an exception will be raised.

	As an additional feature, items in this dictionary can be reset to their
	original value, that they were set to when they were defined. This is to
	further serve the task that this dictionary was intended for in
	configuration.

	For the rest, this class should behave like an ordinary dictionary.
	"""

	"""
	The metadata belonging to one item in the dictionary.

	This is a named tuple consisting of the following fields:
	* default: The default value for the item, meaning the original value. This
	can later be restored.
	* validator: A predicate that indicates whether a certain value is allowed
	in this item.
	"""
	_Metadata = collections.namedtuple("_Item", "default validator")

	def __init__(self):
		"""
		Initialises a new validated dictionary.

		It is initially empty.
		"""
		super().__init__()
		self._metadata = {} #The metadata for every key.

	def __setitem__(self, key, value):
		"""
		Changes the value of an item in the dictionary.

		The key must already exist at this point. If it doesn't exist, a
		``KeyError`` will be raised.

		Additionally, as per the main functionality of this class, the value is
		validated first. If the value is invalid, a ``ValueError`` will be
		raised.
		:param key: The key of the item to change.
		:param value: The new value for the item.
		:raises KeyError: The key does not exist in the dictionary.
		:raises ValueError: The value is invalid for the specified key.
		"""
		if not super().__contains__(key):
			raise KeyError("The key \"{key}\" is not defined in this validated dictionary.".format(key=key))

		if not self._metadata[key].validator(value):
			raise ValueError("The value for the \"{key}\" item is invalid: {value}".format(key=key, value=str(value)))

		super().__setitem__(key, value)

	def __delitem__(self, key):
		"""
		Removes an item from the dictionary.
		:param key: The key of the item to remove.
		"""
		super().__delitem(key)
		del self._metadata[key]

	def add(self, key, value, validator = lambda _: True):
		"""
		Add a new item to the dictionary.

		If no validator is specified, all values will be considered valid for
		this key.

		The key may not exist already. If it does, a ``KeyError`` will be
		raised.
		:param key: The key of the item to add to the dictionary.
		:param value: The default value of the item to add to the dictionary.
		:param validator: A predicate that indicates whether a given value is
		valid or not. The predicate must accept exactly 1 parameter (so don't
		make it a method) and must return a value that can be cast to a boolean
		value that indicates whether the value is valid or not.
		:raises KeyError: The key already exists.
		:raises ValueError: The specified function is not a predicate.
		"""
		if super().__contains__(key):
			raise KeyError("The key \"{key}\" already exists in this validated dictionary.".format(key=key))
		if validator.__code__.co_argcount != 1: #To catch programming mistakes early, check whether this is a predicate. The method vs. function is a common mistake.
			raise ValueError("The validator for key \"{key}\" is not a predicate. It has {argument_count} arguments instead of 1.".format(key=key, argument_count = validator.__code__.co_argcount))
		if not validator(value):
			raise ValueError("The default for key \"{key}\" is invalid: {value}".format(key=key, value=value))

		super().__setitem__(key, value)
		self._metadata[key] = self._Metadata(default=value, validator=validator)

	def reset(self, key):
		"""
		Resets the specified item to its default value.
		:param key: The key of the item to reset.
		"""
		super().__setitem__(key, self._metadata[key].default)<commit_msg>Remove unnecessary import<commit_after>#!/usr/bin/env python
#-*- coding: utf-8 -*-

#This software is distributed under the Creative Commons license (CC0) version 1.0. A copy of this license should have been distributed with this software.
#The license can also be read online: <https://creativecommons.org/publicdomain/zero/1.0/>. If this online license differs from the license provided with this software, the license provided with this software should be applied.

"""
Defines a dictionary-like data structure where all values are validated.

The purpose of this dictionary is to be used in configuration data structures
where the keys are more or less frozen after the application has completed its
start-up phase.
"""

import collections #For namedtuple.

class ValidatedDictionary(dict):
	"""
	A dictionary-like data structure where all values are validated.

	New keys can only be added by also providing a validator predicate. Adding
	new keys can't be done with the usual ``dictionary[key] = value``, but via a
	separate method ``add``. Adding new keys in the usual way would result in an
	exception stating that the key cannot be found.

	When setting a key to a new value, the predicate associated with the key
	will be executed to determine whether the value is allowed. If the value is
	not allowed, an exception will be raised.

	As an additional feature, items in this dictionary can be reset to their
	original value, that they were set to when they were defined. This is to
	further serve the task that this dictionary was intended for in
	configuration.

	For the rest, this class should behave like an ordinary dictionary.
	"""

	"""
	The metadata belonging to one item in the dictionary.

	This is a named tuple consisting of the following fields:
	* default: The default value for the item, meaning the original value. This
	can later be restored.
	* validator: A predicate that indicates whether a certain value is allowed
	in this item.
	"""
	_Metadata = collections.namedtuple("_Item", "default validator")

	def __init__(self):
		"""
		Initialises a new validated dictionary.

		It is initially empty.
		"""
		super().__init__()
		self._metadata = {} #The metadata for every key.

	def __setitem__(self, key, value):
		"""
		Changes the value of an item in the dictionary.

		The key must already exist at this point. If it doesn't exist, a
		``KeyError`` will be raised.

		Additionally, as per the main functionality of this class, the value is
		validated first. If the value is invalid, a ``ValueError`` will be
		raised.
		:param key: The key of the item to change.
		:param value: The new value for the item.
		:raises KeyError: The key does not exist in the dictionary.
		:raises ValueError: The value is invalid for the specified key.
		"""
		if not super().__contains__(key):
			raise KeyError("The key \"{key}\" is not defined in this validated dictionary.".format(key=key))

		if not self._metadata[key].validator(value):
			raise ValueError("The value for the \"{key}\" item is invalid: {value}".format(key=key, value=str(value)))

		super().__setitem__(key, value)

	def __delitem__(self, key):
		"""
		Removes an item from the dictionary.
		:param key: The key of the item to remove.
		"""
		super().__delitem(key)
		del self._metadata[key]

	def add(self, key, value, validator = lambda _: True):
		"""
		Add a new item to the dictionary.

		If no validator is specified, all values will be considered valid for
		this key.

		The key may not exist already. If it does, a ``KeyError`` will be
		raised.
		:param key: The key of the item to add to the dictionary.
		:param value: The default value of the item to add to the dictionary.
		:param validator: A predicate that indicates whether a given value is
		valid or not. The predicate must accept exactly 1 parameter (so don't
		make it a method) and must return a value that can be cast to a boolean
		value that indicates whether the value is valid or not.
		:raises KeyError: The key already exists.
		:raises ValueError: The specified function is not a predicate.
		"""
		if super().__contains__(key):
			raise KeyError("The key \"{key}\" already exists in this validated dictionary.".format(key=key))
		if validator.__code__.co_argcount != 1: #To catch programming mistakes early, check whether this is a predicate. The method vs. function is a common mistake.
			raise ValueError("The validator for key \"{key}\" is not a predicate. It has {argument_count} arguments instead of 1.".format(key=key, argument_count = validator.__code__.co_argcount))
		if not validator(value):
			raise ValueError("The default for key \"{key}\" is invalid: {value}".format(key=key, value=value))

		super().__setitem__(key, value)
		self._metadata[key] = self._Metadata(default=value, validator=validator)

	def reset(self, key):
		"""
		Resets the specified item to its default value.
		:param key: The key of the item to reset.
		"""
		super().__setitem__(key, self._metadata[key].default)<|endoftext|>


<commit_before>"""Crypto utilities."""
import contextlib
import logging
import re
import socket
import sys

import OpenSSL

from acme import errors

logger = logging.getLogger(__name__)

# TLSSNI01 certificate serving and probing is not affected by SSL
# vulnerabilities: prober needs to check certificate for expected
# contents anyway. Working SNI is the only thing that's necessary for
# the challenge and thus scoping down SSL/TLS method (version) would
# cause interoperability issues: TLSv1_METHOD is only compatible with
# TLSv1_METHOD, while SSLv23_METHOD is compatible with all other
# methods, including TLSv2_METHOD (read more at
# https://www.openssl.org/docs/ssl/SSLv23_method.html). _serve_sni
# should be changed to use "set_options" to disable SSLv2 and SSLv3,
# in case it's used for things other than probing/serving!
_DEFAULT_TLSSNI01_SSL_METHOD = OpenSSL.SSL.SSLv23_METHOD

class SSLSocket(object):  # pylint: disable=too-few-public-methods
    """SSL wrapper for sockets.

    :ivar socket sock: Original wrapped socket.
    :ivar dict certs: Mapping from domain names (`bytes`) to
        `OpenSSL.crypto.X509`.
    :ivar method: See `OpenSSL.SSL.Context` for allowed values.

    """
    def __init__(self, sock, certs, method=_DEFAULT_TLSSNI01_SSL_METHOD):
        self.sock = sock
        self.certs = certs
        self.method = method

    def __getattr__(self, name):
        return getattr(self.sock, name)

    def _pick_certificate_cb(self, connection):
        """SNI certificate callback.

        This method will set a new OpenSSL context object for this
        connection when an incoming connection provides an SNI name
        (in order to serve the appropriate certificate, if any).

        :param connection: The TLS connection object on which the SNI
            extension was received.
        :type connection: :class:`OpenSSL.Connection`

        """
        server_name = connection.get_servername()
        try:
            key, cert = self.certs[server_name]
        except KeyError:
            logger.debug("Server name (%s) not recognized, dropping SSL",
                         server_name)
            return
        new_context = OpenSSL.SSL.Context(self.method)
        new_context.use_privatekey(key)
        new_context.use_certificate(cert)
        connection.set_context(new_context)

    class FakeConnection(object):
        """Fake OpenSSL.SSL.Connection."""

        # pylint: disable=missing-docstring

        def __init__(self, connection):
            self._wrapped = connection

        def __getattr__(self, name):
            return getattr(self._wrapped, name)

        def shutdown(self, *unused_args):
            # OpenSSL.SSL.Connection.shutdown doesn't accept any args
            return self._wrapped.shutdown()

    def accept(self):  # pylint: disable=missing-docstring
        sock, addr = self.sock.accept()

        context = OpenSSL.SSL.Context(self.method)
        context.set_tlsext_servername_callback(self._pick_certificate_cb)

        ssl_sock = self.FakeConnection(OpenSSL.SSL.Connection(context, sock))
        ssl_sock.set_accept_state()

        logger.debug("Performing handshake with %s", addr)
        try:
            ssl_sock.do_handshake()
        except OpenSSL.SSL.Error as error:
            # _pick_certificate_cb might have returned without
            # creating SSL context (wrong server name)
            raise socket.error(error)

        return ssl_sock, addr

def probe_sni(name, host, port=443, timeout=300,
              method=_DEFAULT_TLSSNI01_SSL_METHOD, source_address=('0', 0)):
    """Probe SNI server for SSL certificate.

    :param bytes name: Byte string to send as the server name in the
        client hello message.
    :param bytes host: Host to connect to.
    :param int port: Port to connect to.
    :param int timeout: Timeout in seconds.
    :param method: See `OpenSSL.SSL.Context` for allowed values.
    :param tuple source_address: Enables multi-path probing (selection
        of source interface). See `socket.creation_connection` for more
        info. Available only in Python 2.7+.

    :raises acme.errors.Error: In case of any problems.

    :returns: SSL certificate presented by the server.
    :rtype: OpenSSL.crypto.X509

    """
    context = OpenSSL.SSL.Context(method)
    context.set_timeout(timeout)

    socket_kwargs = {} if sys.version_info < (2, 7) else {
        'source_address': source_address}

    try:
        # pylint: disable=star-args
        sock = socket.create_connection((host, port), **socket_kwargs)
    except socket.error as error:
        raise errors.Error(error)

    with contextlib.closing(sock) as client:
        client_ssl = OpenSSL.SSL.Connection(context, client)
        client_ssl.set_connect_state()
        client_ssl.set_tlsext_host_name(name)  # pyOpenSSL>=0.13
        try:
            client_ssl.do_handshake()
            client_ssl.shutdown()
        except OpenSSL.SSL.Error as error:
            raise errors.Error(error)
    return client_ssl.get_peer_certificate()

def _pyopenssl_cert_or_req_san(cert_or_req):
    """Get Subject Alternative Names from certificate or CSR using pyOpenSSL.

    .. todo:: Implement directly in PyOpenSSL!

    .. note:: Although this is `acme` internal API, it is used by
        `letsencrypt`.

    :param cert_or_req: Certificate or CSR.
    :type cert_or_req: `OpenSSL.crypto.X509` or `OpenSSL.crypto.X509Req`.

    :returns: A list of Subject Alternative Names.
    :rtype: `list` of `unicode`

    """
    # constants based on PyOpenSSL certificate/CSR text dump
    part_separator = ":"
    parts_separator = ", "
    prefix = "DNS" + part_separator

    if isinstance(cert_or_req, OpenSSL.crypto.X509):
        func = OpenSSL.crypto.dump_certificate
    else:
        func = OpenSSL.crypto.dump_certificate_request
    text = func(OpenSSL.crypto.FILETYPE_TEXT, cert_or_req).decode("utf-8")
    match = re.search(r"X509v3 Subject Alternative Name:\s*(.*)", text)
    sans_parts = [] if match is None else match.group(1).split(parts_separator)
    # WARNING: this function assumes that no SAN can include
    # parts_separator, hence the split!

    return [part.split(part_separator)[1]
            for part in sans_parts if part.startswith(prefix)]

def gen_ss_cert(key, domains, not_before=None,
                validity=(7 * 24 * 60 * 60), force_san=True):
    """Generate new self-signed certificate.

    :type domains: `list` of `unicode`
    :param OpenSSL.crypto.PKey key:
    :param bool force_san:

    If more than one domain is provided, all of the domains are put into
    ``subjectAltName`` X.509 extension and first domain is set as the
    subject CN. If only one domain is provided no ``subjectAltName``
    extension is used, unless `force_san` is ``True``.

    """
    assert domains, "Must provide one or more hostnames for the cert."
    cert = OpenSSL.crypto.X509()
    cert.set_serial_number(1337)
    cert.set_version(2)

    extensions = [
        OpenSSL.crypto.X509Extension(
            b"basicConstraints", True, b"CA:TRUE, pathlen:0"),
    ]

    cert.get_subject().CN = domains[0]
    # TODO: what to put into cert.get_subject()?
    cert.set_issuer(cert.get_subject())

    if force_san or len(domains) > 1:
        extensions.append(OpenSSL.crypto.X509Extension(
            b"subjectAltName",
            critical=False,
            value=b", ".join(b"DNS:" + d.encode() for d in domains)
        ))

    cert.add_extensions(extensions)

    cert.gmtime_adj_notBefore(0 if not_before is None else not_before)
    cert.gmtime_adj_notAfter(validity)

    cert.set_pubkey(key)
    cert.sign(key, "sha256")
    return cert
<commit_msg>Comment _pyopenssl_cert_or_req_san method<commit_after>"""Crypto utilities."""
import contextlib
import logging
import re
import socket
import sys

import OpenSSL

from acme import errors

logger = logging.getLogger(__name__)

# TLSSNI01 certificate serving and probing is not affected by SSL
# vulnerabilities: prober needs to check certificate for expected
# contents anyway. Working SNI is the only thing that's necessary for
# the challenge and thus scoping down SSL/TLS method (version) would
# cause interoperability issues: TLSv1_METHOD is only compatible with
# TLSv1_METHOD, while SSLv23_METHOD is compatible with all other
# methods, including TLSv2_METHOD (read more at
# https://www.openssl.org/docs/ssl/SSLv23_method.html). _serve_sni
# should be changed to use "set_options" to disable SSLv2 and SSLv3,
# in case it's used for things other than probing/serving!
_DEFAULT_TLSSNI01_SSL_METHOD = OpenSSL.SSL.SSLv23_METHOD

class SSLSocket(object):  # pylint: disable=too-few-public-methods
    """SSL wrapper for sockets.

    :ivar socket sock: Original wrapped socket.
    :ivar dict certs: Mapping from domain names (`bytes`) to
        `OpenSSL.crypto.X509`.
    :ivar method: See `OpenSSL.SSL.Context` for allowed values.

    """
    def __init__(self, sock, certs, method=_DEFAULT_TLSSNI01_SSL_METHOD):
        self.sock = sock
        self.certs = certs
        self.method = method

    def __getattr__(self, name):
        return getattr(self.sock, name)

    def _pick_certificate_cb(self, connection):
        """SNI certificate callback.

        This method will set a new OpenSSL context object for this
        connection when an incoming connection provides an SNI name
        (in order to serve the appropriate certificate, if any).

        :param connection: The TLS connection object on which the SNI
            extension was received.
        :type connection: :class:`OpenSSL.Connection`

        """
        server_name = connection.get_servername()
        try:
            key, cert = self.certs[server_name]
        except KeyError:
            logger.debug("Server name (%s) not recognized, dropping SSL",
                         server_name)
            return
        new_context = OpenSSL.SSL.Context(self.method)
        new_context.use_privatekey(key)
        new_context.use_certificate(cert)
        connection.set_context(new_context)

    class FakeConnection(object):
        """Fake OpenSSL.SSL.Connection."""

        # pylint: disable=missing-docstring

        def __init__(self, connection):
            self._wrapped = connection

        def __getattr__(self, name):
            return getattr(self._wrapped, name)

        def shutdown(self, *unused_args):
            # OpenSSL.SSL.Connection.shutdown doesn't accept any args
            return self._wrapped.shutdown()

    def accept(self):  # pylint: disable=missing-docstring
        sock, addr = self.sock.accept()

        context = OpenSSL.SSL.Context(self.method)
        context.set_tlsext_servername_callback(self._pick_certificate_cb)

        ssl_sock = self.FakeConnection(OpenSSL.SSL.Connection(context, sock))
        ssl_sock.set_accept_state()

        logger.debug("Performing handshake with %s", addr)
        try:
            ssl_sock.do_handshake()
        except OpenSSL.SSL.Error as error:
            # _pick_certificate_cb might have returned without
            # creating SSL context (wrong server name)
            raise socket.error(error)

        return ssl_sock, addr

def probe_sni(name, host, port=443, timeout=300,
              method=_DEFAULT_TLSSNI01_SSL_METHOD, source_address=('0', 0)):
    """Probe SNI server for SSL certificate.

    :param bytes name: Byte string to send as the server name in the
        client hello message.
    :param bytes host: Host to connect to.
    :param int port: Port to connect to.
    :param int timeout: Timeout in seconds.
    :param method: See `OpenSSL.SSL.Context` for allowed values.
    :param tuple source_address: Enables multi-path probing (selection
        of source interface). See `socket.creation_connection` for more
        info. Available only in Python 2.7+.

    :raises acme.errors.Error: In case of any problems.

    :returns: SSL certificate presented by the server.
    :rtype: OpenSSL.crypto.X509

    """
    context = OpenSSL.SSL.Context(method)
    context.set_timeout(timeout)

    socket_kwargs = {} if sys.version_info < (2, 7) else {
        'source_address': source_address}

    try:
        # pylint: disable=star-args
        sock = socket.create_connection((host, port), **socket_kwargs)
    except socket.error as error:
        raise errors.Error(error)

    with contextlib.closing(sock) as client:
        client_ssl = OpenSSL.SSL.Connection(context, client)
        client_ssl.set_connect_state()
        client_ssl.set_tlsext_host_name(name)  # pyOpenSSL>=0.13
        try:
            client_ssl.do_handshake()
            client_ssl.shutdown()
        except OpenSSL.SSL.Error as error:
            raise errors.Error(error)
    return client_ssl.get_peer_certificate()

def _pyopenssl_cert_or_req_san(cert_or_req):
    """Get Subject Alternative Names from certificate or CSR using pyOpenSSL.

    .. todo:: Implement directly in PyOpenSSL!

    .. note:: Although this is `acme` internal API, it is used by
        `letsencrypt`.

    :param cert_or_req: Certificate or CSR.
    :type cert_or_req: `OpenSSL.crypto.X509` or `OpenSSL.crypto.X509Req`.

    :returns: A list of Subject Alternative Names.
    :rtype: `list` of `unicode`

    """
    # constants based on PyOpenSSL certificate/CSR text dump
    part_separator = ":"
    parts_separator = ", "
    prefix = "DNS" + part_separator

    if isinstance(cert_or_req, OpenSSL.crypto.X509):
        func = OpenSSL.crypto.dump_certificate
    else:
        func = OpenSSL.crypto.dump_certificate_request

    # This method of finding SANs is used to support PyOpenSSL version 0.13.
    text = func(OpenSSL.crypto.FILETYPE_TEXT, cert_or_req).decode("utf-8")
    match = re.search(r"X509v3 Subject Alternative Name:\s*(.*)", text)

    # WARNING: this function assumes that no SAN can include
    # parts_separator, hence the split!
    sans_parts = [] if match is None else match.group(1).split(parts_separator)

    return [part.split(part_separator)[1]
            for part in sans_parts if part.startswith(prefix)]

def gen_ss_cert(key, domains, not_before=None,
                validity=(7 * 24 * 60 * 60), force_san=True):
    """Generate new self-signed certificate.

    :type domains: `list` of `unicode`
    :param OpenSSL.crypto.PKey key:
    :param bool force_san:

    If more than one domain is provided, all of the domains are put into
    ``subjectAltName`` X.509 extension and first domain is set as the
    subject CN. If only one domain is provided no ``subjectAltName``
    extension is used, unless `force_san` is ``True``.

    """
    assert domains, "Must provide one or more hostnames for the cert."
    cert = OpenSSL.crypto.X509()
    cert.set_serial_number(1337)
    cert.set_version(2)

    extensions = [
        OpenSSL.crypto.X509Extension(
            b"basicConstraints", True, b"CA:TRUE, pathlen:0"),
    ]

    cert.get_subject().CN = domains[0]
    # TODO: what to put into cert.get_subject()?
    cert.set_issuer(cert.get_subject())

    if force_san or len(domains) > 1:
        extensions.append(OpenSSL.crypto.X509Extension(
            b"subjectAltName",
            critical=False,
            value=b", ".join(b"DNS:" + d.encode() for d in domains)
        ))

    cert.add_extensions(extensions)

    cert.gmtime_adj_notBefore(0 if not_before is None else not_before)
    cert.gmtime_adj_notAfter(validity)

    cert.set_pubkey(key)
    cert.sign(key, "sha256")
    return cert
<|endoftext|>