<commit_before>from functools import wraps
import os
import json
from base64 import b64encode
import time as time_module
from copy import copy
import logging

from six.moves.urllib.parse import urlencode
from flask import request, session, redirect, url_for, g
from oauth2client.client import flow_from_clientsecrets, OAuth2WebServerFlow,\
    AccessTokenRefreshError
import httplib2
from itsdangerous import TimedJSONWebSignatureSerializer, SignatureExpired

__all__ = ['OpenIDConnect', 'MemoryCredentials']

logger = logging.getLogger(__name__)

class MemoryCredentials(dict):
    """
    Non-persistent local credentials store.
    Use this if you only have one app server, and don't mind making everyone
    log in again after a restart.
    """
    pass

class OpenIDConnect(object):
    """
    @see: https://developers.google.com/api-client-library/python/start/get_started
    @see: https://developers.google.com/api-client-library/python/samples/authorized_api_web_server_calendar.py
    """
    def __init__(self, app=None, credentials_store=None, http=None, time=None,
                 urandom=None):
        # set from app config in .init_app()
        self.callback_path = None
        self.flow = None
        self.cookie_serializer = None

        # optional, also set from app config
        self.google_apps_domain = None
        self.id_token_cookie_name = 'oidc_id_token'
        self.id_token_cookie_ttl = 7 * 86400  # one week
        # should ONLY be turned off for local debugging
        self.id_token_cookie_secure = True

        # stuff that we might want to override for tests
        self.http = http if http is not None else httplib2.Http()
        self.credentials_store = credentials_store\
            if credentials_store is not None\
            else MemoryCredentials()
        self.time = time if time is not None else time_module.time
        self.urandom = urandom if urandom is not None else os.urandom

        # get stuff from the app's config, which may override stuff set above
        if app is not None:
            self.init_app(app)

    def init_app(self, app):
        """
        Do setup that requires a Flask app.
        """
        # register callback route and cookie-setting decorator
        app.route('/oidc_callback')(self.oidc_callback)
        app.after_request(self.after_request)

        # load client_secrets.json
        self.flow = flow_from_clientsecrets(
            app.config['OIDC_CLIENT_SECRETS'],
            scope=['openid', 'email'])
        assert isinstance(self.flow, OAuth2WebServerFlow)

        # create a cookie signer using the Flask secret key
        self.cookie_serializer = TimedJSONWebSignatureSerializer(
            app.config['SECRET_KEY'])

        try:
            self.google_apps_domain = app.config['OIDC_GOOGLE_APPS_DOMAIN']
        except KeyError:
            pass

        try:
            self.id_token_cookie_name = app.config['OIDC_ID_TOKEN_COOKIE_NAME']
        except KeyError:
            pass

        try:
            self.id_token_cookie_ttl = app.config['OIDC_ID_TOKEN_COOKIE_TTL']
        except KeyError:
            pass

        try:
            self.id_token_cookie_secure =\
                app.config['OIDC_ID_TOKEN_COOKIE_SECURE']
        except KeyError:
            pass

        try:
            self.credentials_store = app.config['OIDC_CREDENTIALS_STORE']
        except KeyError:
            pass

    def get_cookie_id_token(self):
        try:
            id_token_cookie = request.cookies[self.id_token_cookie_name]
            return self.cookie_serializer.loads(id_token_cookie)
        except (KeyError, SignatureExpired):
            logger.debug("Missing or invalid ID token cookie", exc_info=True)
            return None

    def set_cookie_id_token(self, id_token):
        """
        Cooperates with @after_request to set a new ID token cookie.
        """
        g.oidc_id_token = id_token
        g.oidc_id_token_dirty = True

    def after_request(self, response):
        """
        Set a new ID token cookie if the ID token has changed.
        """
        if getattr(g, 'oidc_id_token_dirty', False):
            signed_id_token = self.cookie_serializer.dumps(g.oidc_id_token)
            response.set_cookie(
                self.id_token_cookie_name, signed_id_token,
                secure=self.id_token_cookie_secure,
                httponly=True,
                max_age=self.id_token_cookie_ttl)
        return response

    def authenticate_or_redirect(self):
        """
        Helper function suitable for @app.before_request and @check (below).
        Sets g.oidc_id_token to the ID token if the user has successfully
        authenticated, else returns a redirect object so they can go try
        to authenticate.
        :return: A redirect, or None if the user is authenticated.
        """
        # the auth callback and error pages don't need user to be authenticated
        if request.endpoint in frozenset(['oidc_callback', 'oidc_error']):
            return None

        # retrieve signed ID token cookie
        id_token = self.get_cookie_id_token()
        if id_token is None:
            return self.redirect_to_auth_server(request.url)

        # ID token expired
        # when Google is the IdP, this happens after one hour
        if self.time() >= id_token['exp']:
            # get credentials from store
            try:
                credentials = self.credentials_store[id_token['sub']]
            except KeyError:
                logger.debug("Expired ID token, credentials missing",
                             exc_info=True)
                return self.redirect_to_auth_server(request.url)

            # refresh and store credentials
            try:
                credentials.refresh(self.http)
                id_token = credentials.id_token
                self.credentials_store[id_token['sub']] = credentials
                self.set_cookie_id_token(id_token)
            except AccessTokenRefreshError:
                # Can't refresh. Wipe credentials and redirect user to IdP
                # for re-authentication.
                logger.debug("Expired ID token, can't refresh credentials",
                             exc_info=True)
                del self.credentials_store[id_token['sub']]
                return self.redirect_to_auth_server(request.url)

        # make ID token available to views
        g.oidc_id_token = id_token

        return None

    def check(self, view_func):
        """
        Use this to decorate view functions if only some of your app's views
        require authentication.
        """
        @wraps(view_func)
        def decorated(*args, **kwargs):
            response = self.authenticate_or_redirect()
            if response is not None:
                return response
            return view_func(*args, **kwargs)
        return decorated

    def flow_for_request(self):
        """
        Build a flow with the correct absolute callback URL for this request.
        :return:
        """
        flow = copy(self.flow)
        flow.redirect_uri = url_for('oidc_callback', _external=True)
        return flow

    def redirect_to_auth_server(self, destination):
        """
        Set a CSRF token in the session, and redirect to the IdP.
        :param destination: the page that the user was going to,
                            before we noticed they weren't logged in
        :return: a redirect response
        """
        csrf_token = b64encode(self.urandom(24)).decode('utf-8')
        session['oidc_csrf_token'] = csrf_token
        state = {
            'csrf_token': csrf_token,
            'destination': destination,
        }
        extra_params = {
            'state': json.dumps(state),
        }
        flow = self.flow_for_request()
        auth_url = '{url}&{extra_params}'.format(
            url=flow.step1_get_authorize_url(),
            extra_params=urlencode(extra_params))
        # if the user has an ID token, it's invalid, or we wouldn't be here
        self.set_cookie_id_token(None)
        return redirect(auth_url)

    def is_id_token_valid(self, id_token):
        """
        Check if `id_token` is a current ID token for this application,
        was issued by the Apps domain we expected,
        and that the email address has been verified.

        @see: http://openid.net/specs/openid-connect-core-1_0.html#IDTokenValidation
        """
        if not id_token:
            return False

        # TODO: step 2: check issuer

        if isinstance(id_token['aud'], list):
            # step 3 for audience list
            if self.flow.client_id not in id_token['aud']:
                return False
            # step 4
            if 'azp' not in id_token:
                return False
        else:
            # step 3 for single audience
            if id_token['aud'] != self.flow.client_id:
                return False

        # step 5
        if 'azp' in id_token and id_token['azp'] != self.flow.client_id:
            return False

        # steps 9, 10
        if not (id_token['iat'] <= self.time() < id_token['exp']):
            return False

        # (not required if using HTTPS?) step 11: check nonce

        # additional steps specific to our usage

        if id_token.get('hd') != self.google_apps_domain:
            return False

        if not id_token['email_verified']:
            return False

        return True

    WRONG_GOOGLE_APPS_DOMAIN = 'WRONG_GOOGLE_APPS_DOMAIN'

    def oidc_callback(self):
        """
        Exchange the auth code for actual credentials,
        then redirect to the originally requested page.
        """
        # retrieve session and callback variables
        try:
            session_csrf_token = session.pop('oidc_csrf_token')

            state = json.loads(request.args['state'])
            csrf_token = state['csrf_token']
            destination = state['destination']

            code = request.args['code']
        except (KeyError, ValueError):
            logger.debug("Can't retrieve CSRF token, state, or code",
                         exc_info=True)
            return self.oidc_error()

        # check callback CSRF token passed to IdP
        # against session CSRF token held by user
        if csrf_token != session_csrf_token:
            logger.debug("CSRF token mismatch")
            return self.oidc_error()

        # make a request to IdP to exchange the auth code for OAuth credentials
        flow = self.flow_for_request()
        credentials = flow.step2_exchange(code, http=self.http)
        id_token = credentials.id_token
        if not self.is_id_token_valid(id_token):
            logger.debug("Invalid ID token")
            # None will allow all domains otherwise limit down to the specific domain.
            if self.google_apps_domain is not None:
                if not isinstance(self.google_apps_domain, list):
                    self.google_apps_domain = [self.google_apps_domain]
                for domain in self.google_apps_domain:
                    if id_token.get('hd') != domain:
                        return self.oidc_error(
                            "You must log in with an account from the {0} domain."
                            .format(self.google_apps_domain),
                            self.WRONG_GOOGLE_APPS_DOMAIN)
            return self.oidc_error()

        # store credentials by subject
        # when Google is the IdP, the subject is their G+ account number
        self.credentials_store[id_token['sub']] = credentials

        # set a persistent signed cookie containing the ID token
        # and redirect to the final destination
        # TODO: validate redirect destination
        response = redirect_to(destination)
        self.set_cookie_id_token(id_token)
        return response

    def oidc_error(self, message='Not Authorized', code=None):
        return (message, 401, {
            'Content-Type': 'text/plain',
        })

    def redirect_to(self, destination, *args, **kwargs):
        path = request.args.get('next') or \
            request.referrer or \
            destination
        return redirect(path)
<commit_msg>adding profile to the scope<commit_after>from functools import wraps
import os
import json
from base64 import b64encode
import time as time_module
from copy import copy
import logging

from six.moves.urllib.parse import urlencode
from flask import request, session, redirect, url_for, g
from oauth2client.client import flow_from_clientsecrets, OAuth2WebServerFlow,\
    AccessTokenRefreshError
import httplib2
from itsdangerous import TimedJSONWebSignatureSerializer, SignatureExpired

__all__ = ['OpenIDConnect', 'MemoryCredentials']

logger = logging.getLogger(__name__)

class MemoryCredentials(dict):
    """
    Non-persistent local credentials store.
    Use this if you only have one app server, and don't mind making everyone
    log in again after a restart.
    """
    pass

class OpenIDConnect(object):
    """
    @see: https://developers.google.com/api-client-library/python/start/get_started
    @see: https://developers.google.com/api-client-library/python/samples/authorized_api_web_server_calendar.py
    """
    def __init__(self, app=None, credentials_store=None, http=None, time=None,
                 urandom=None):
        # set from app config in .init_app()
        self.callback_path = None
        self.flow = None
        self.cookie_serializer = None

        # optional, also set from app config
        self.google_apps_domain = None
        self.id_token_cookie_name = 'oidc_id_token'
        self.id_token_cookie_ttl = 7 * 86400  # one week
        # should ONLY be turned off for local debugging
        self.id_token_cookie_secure = True

        # stuff that we might want to override for tests
        self.http = http if http is not None else httplib2.Http()
        self.credentials_store = credentials_store\
            if credentials_store is not None\
            else MemoryCredentials()
        self.time = time if time is not None else time_module.time
        self.urandom = urandom if urandom is not None else os.urandom

        # get stuff from the app's config, which may override stuff set above
        if app is not None:
            self.init_app(app)

    def init_app(self, app):
        """
        Do setup that requires a Flask app.
        """
        # register callback route and cookie-setting decorator
        app.route('/oidc_callback')(self.oidc_callback)
        app.after_request(self.after_request)

        # load client_secrets.json
        self.flow = flow_from_clientsecrets(
            app.config['OIDC_CLIENT_SECRETS'],
            scope=['openid', 'profile', 'email'])
        assert isinstance(self.flow, OAuth2WebServerFlow)

        # create a cookie signer using the Flask secret key
        self.cookie_serializer = TimedJSONWebSignatureSerializer(
            app.config['SECRET_KEY'])

        try:
            self.google_apps_domain = app.config['OIDC_GOOGLE_APPS_DOMAIN']
        except KeyError:
            pass

        try:
            self.id_token_cookie_name = app.config['OIDC_ID_TOKEN_COOKIE_NAME']
        except KeyError:
            pass

        try:
            self.id_token_cookie_ttl = app.config['OIDC_ID_TOKEN_COOKIE_TTL']
        except KeyError:
            pass

        try:
            self.id_token_cookie_secure =\
                app.config['OIDC_ID_TOKEN_COOKIE_SECURE']
        except KeyError:
            pass

        try:
            self.credentials_store = app.config['OIDC_CREDENTIALS_STORE']
        except KeyError:
            pass

    def get_cookie_id_token(self):
        try:
            id_token_cookie = request.cookies[self.id_token_cookie_name]
            return self.cookie_serializer.loads(id_token_cookie)
        except (KeyError, SignatureExpired):
            logger.debug("Missing or invalid ID token cookie", exc_info=True)
            return None

    def set_cookie_id_token(self, id_token):
        """
        Cooperates with @after_request to set a new ID token cookie.
        """
        g.oidc_id_token = id_token
        g.oidc_id_token_dirty = True

    def after_request(self, response):
        """
        Set a new ID token cookie if the ID token has changed.
        """
        if getattr(g, 'oidc_id_token_dirty', False):
            signed_id_token = self.cookie_serializer.dumps(g.oidc_id_token)
            response.set_cookie(
                self.id_token_cookie_name, signed_id_token,
                secure=self.id_token_cookie_secure,
                httponly=True,
                max_age=self.id_token_cookie_ttl)
        return response

    def authenticate_or_redirect(self):
        """
        Helper function suitable for @app.before_request and @check (below).
        Sets g.oidc_id_token to the ID token if the user has successfully
        authenticated, else returns a redirect object so they can go try
        to authenticate.
        :return: A redirect, or None if the user is authenticated.
        """
        # the auth callback and error pages don't need user to be authenticated
        if request.endpoint in frozenset(['oidc_callback', 'oidc_error']):
            return None

        # retrieve signed ID token cookie
        id_token = self.get_cookie_id_token()
        if id_token is None:
            return self.redirect_to_auth_server(request.url)

        # ID token expired
        # when Google is the IdP, this happens after one hour
        if self.time() >= id_token['exp']:
            # get credentials from store
            try:
                credentials = self.credentials_store[id_token['sub']]
            except KeyError:
                logger.debug("Expired ID token, credentials missing",
                             exc_info=True)
                return self.redirect_to_auth_server(request.url)

            # refresh and store credentials
            try:
                credentials.refresh(self.http)
                id_token = credentials.id_token
                self.credentials_store[id_token['sub']] = credentials
                self.set_cookie_id_token(id_token)
            except AccessTokenRefreshError:
                # Can't refresh. Wipe credentials and redirect user to IdP
                # for re-authentication.
                logger.debug("Expired ID token, can't refresh credentials",
                             exc_info=True)
                del self.credentials_store[id_token['sub']]
                return self.redirect_to_auth_server(request.url)

        # make ID token available to views
        g.oidc_id_token = id_token

        return None

    def check(self, view_func):
        """
        Use this to decorate view functions if only some of your app's views
        require authentication.
        """
        @wraps(view_func)
        def decorated(*args, **kwargs):
            response = self.authenticate_or_redirect()
            if response is not None:
                return response
            return view_func(*args, **kwargs)
        return decorated

    def flow_for_request(self):
        """
        Build a flow with the correct absolute callback URL for this request.
        :return:
        """
        flow = copy(self.flow)
        flow.redirect_uri = url_for('oidc_callback', _external=True)
        return flow

    def redirect_to_auth_server(self, destination):
        """
        Set a CSRF token in the session, and redirect to the IdP.
        :param destination: the page that the user was going to,
                            before we noticed they weren't logged in
        :return: a redirect response
        """
        csrf_token = b64encode(self.urandom(24)).decode('utf-8')
        session['oidc_csrf_token'] = csrf_token
        state = {
            'csrf_token': csrf_token,
            'destination': destination,
        }
        extra_params = {
            'state': json.dumps(state),
        }
        flow = self.flow_for_request()
        auth_url = '{url}&{extra_params}'.format(
            url=flow.step1_get_authorize_url(),
            extra_params=urlencode(extra_params))
        # if the user has an ID token, it's invalid, or we wouldn't be here
        self.set_cookie_id_token(None)
        return redirect(auth_url)

    def is_id_token_valid(self, id_token):
        """
        Check if `id_token` is a current ID token for this application,
        was issued by the Apps domain we expected,
        and that the email address has been verified.

        @see: http://openid.net/specs/openid-connect-core-1_0.html#IDTokenValidation
        """
        if not id_token:
            return False

        # TODO: step 2: check issuer

        if isinstance(id_token['aud'], list):
            # step 3 for audience list
            if self.flow.client_id not in id_token['aud']:
                return False
            # step 4
            if 'azp' not in id_token:
                return False
        else:
            # step 3 for single audience
            if id_token['aud'] != self.flow.client_id:
                return False

        # step 5
        if 'azp' in id_token and id_token['azp'] != self.flow.client_id:
            return False

        # steps 9, 10
        if not (id_token['iat'] <= self.time() < id_token['exp']):
            return False

        # (not required if using HTTPS?) step 11: check nonce

        # additional steps specific to our usage

        if id_token.get('hd') != self.google_apps_domain:
            return False

        if not id_token['email_verified']:
            return False

        return True

    WRONG_GOOGLE_APPS_DOMAIN = 'WRONG_GOOGLE_APPS_DOMAIN'

    def oidc_callback(self):
        """
        Exchange the auth code for actual credentials,
        then redirect to the originally requested page.
        """
        # retrieve session and callback variables
        try:
            session_csrf_token = session.pop('oidc_csrf_token')

            state = json.loads(request.args['state'])
            csrf_token = state['csrf_token']
            destination = state['destination']

            code = request.args['code']
        except (KeyError, ValueError):
            logger.debug("Can't retrieve CSRF token, state, or code",
                         exc_info=True)
            return self.oidc_error()

        # check callback CSRF token passed to IdP
        # against session CSRF token held by user
        if csrf_token != session_csrf_token:
            logger.debug("CSRF token mismatch")
            return self.oidc_error()

        # make a request to IdP to exchange the auth code for OAuth credentials
        flow = self.flow_for_request()
        credentials = flow.step2_exchange(code, http=self.http)
        id_token = credentials.id_token
        if not self.is_id_token_valid(id_token):
            logger.debug("Invalid ID token")
            # None will allow all domains otherwise limit down to the specific domain.
            if self.google_apps_domain is not None:
                if not isinstance(self.google_apps_domain, list):
                    self.google_apps_domain = [self.google_apps_domain]
                for domain in self.google_apps_domain:
                    if id_token.get('hd') != domain:
                        return self.oidc_error(
                            "You must log in with an account from the {0} domain."
                            .format(self.google_apps_domain),
                            self.WRONG_GOOGLE_APPS_DOMAIN)
            return self.oidc_error()

        # store credentials by subject
        # when Google is the IdP, the subject is their G+ account number
        self.credentials_store[id_token['sub']] = credentials

        # set a persistent signed cookie containing the ID token
        # and redirect to the final destination
        # TODO: validate redirect destination
        response = redirect_to(destination)
        self.set_cookie_id_token(id_token)
        return response

    def oidc_error(self, message='Not Authorized', code=None):
        return (message, 401, {
            'Content-Type': 'text/plain',
        })

    def redirect_to(self, destination, *args, **kwargs):
        path = request.args.get('next') or \
            request.referrer or \
            destination
        return redirect(path)
<|endoftext|>


<commit_before>"""Base functionality for fancy validation."""
import __builtin__
import operator
from flatland.util import adict
from flatland.schema.util import find_i18n_function

N_ = lambda translatable: translatable

class Validator(object):
    """Base class for fancy validators."""

    _transform_finder = operator.attrgetter('ugettext')
    _tuple_transform_finder = operator.attrgetter('ungettext')

    def __init__(self, **kw):
        """Construct a validator.

        :param \*\*kw: override any extant class attribute on this instance.

        """
        cls = type(self)
        for attr, value in kw.iteritems():
            if hasattr(cls, attr):
                setattr(self, attr, value)
            else:
                raise TypeError("%s has no attribute %r, can not override." % (
                    cls.__name__, attr))

    def __call__(self, element, state):
        """Adapts Validator to the Element.validate callable interface."""
        return self.validate(element, state)

    def validate(self, element, state):
        """Validate an element returning True if valid.

        :param element:
          an :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by
          :meth:`Element.validate <flatland.schema.base.Element.validate>`.

        :returns: True if valid

        """
        return False

    def note_error(self, element, state, key=None, message=None, **info):
        """Record a validation error message on an element.

        :param element:
          An :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by :meth:`Element.validate
          <flatland.schema.base.Element.validate>`.

        :param key: semi-optional, default None.
          The name of a message-holding attribute on this instance.  Will be
          used to ``message = getattr(self, key)``.

        :param message: semi-optional, default None.
          A validation message.

        :param \*\*info: optional.
          Additional data to make available to validation message
          string formatting.

        :returns: False

        Either *key* or *message* is required.  The message will have
        formatting expanded by :meth:`expand_message` and be appended to
        :attr:`element.errors <flatland.schema.base.Element.errors>`.

        Always returns False.  This enables a convenient shorthand when
        writing validators:

        .. testcode::

          from flatland.valid import Validator

          class MyValidator(Validator):
              my_message = 'Oh noes!'

              def validate(self, element, state):
                  if not element.value:
                      return self.note_error(element, state, 'my_message')
                  else:
                      return True

        .. testcode:: :hide:

          from flatland import String
          el = String('x').create_element()
          v = MyValidator()
          assert not v.validate(el, None)
          assert el.errors == ['Oh noes!']
          el.set('foo')
          assert v.validate(el, None)
          assert el.errors == ['Oh noes!']

        """
        message = message or getattr(self, key)
        if message:
            element.add_error(
                self.expand_message(element, state, message, **info))
        return False

    def note_warning(self, element, state, key=None, message=None, **info):
        """Record a validation warning message on an element.

        :param element:
          An :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by :meth:`Element.validate
          <flatland.schema.base.Element.validate>`.

        :param key: semi-optional, default None.
          The name of a message-holding attribute on this instance.  Will be
          used to ``message = getattr(self, key)``.

        :param message: semi-optional, default None.
          A validation message.

        :param \*\*info: optional.
          Additional data to make available to validation message
          string formatting.

        :returns: False

        Either *key* or *message* is required.  The message will have
        formatting expanded by :meth:`expand_message` and be appended to
        :attr:`element.warnings <flatland.schema.base.Element.warnings>`.

        Always returns False.
        """
        message = message or getattr(self, key)
        if message:
            element.add_warning(
                self.expand_message(element, state, message, **info))
        return False

    def find_transformer(self, element, state, message, finder):
        return find_i18n_function(element, finder)

    def expand_message(self, element, state, message, **extra_format_args):
        """Apply formatting to a validation message.

        :param element:
          an :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by
          :meth:`Element.validate <flatland.schema.base.Element.validate>`.

        :param message: a string, 3-tuple or callable.
          If a 3-tuple, must be of the form ('single form', 'plural form',
          n_key).

          If callable, will be called with 2 positional arguments (*element*,
          *state*) and must return a string or 3-tuple.

        :param \*\*extra_format_args: optional.
          Additional data to make available to validation message
          string formatting.

        :returns: the formatted string

        See :ref:`validation_messaging` for full information on how messages
        are expanded.

        """
        if callable(message):
            message = message(element, state)

        message_transform = mapping_transform = self.find_transformer(
            element, state, message, self._transform_finder)

        if extra_format_args:
            extra_format_args = adict(extra_format_args)
            format_map = as_format_mapping(
                extra_format_args, element, self,
                transform=mapping_transform)
        else:
            format_map = as_format_mapping(
                element, self,
                transform=mapping_transform)

        if isinstance(message, tuple):
            # a transformer must be present if message is a tuple
            transform = self.find_transformer(
                element, state, message, self._tuple_transform_finder)

            single, plural, n_key = message
            try:
                n = format_map[n_key]
                try:
                    n = int(n)
                except TypeError:
                    pass
            except KeyError:
                n = n_key

            if transform:
                message = transform(single, plural, n)
            else:
                message = single if n == 1 else plural
        elif message_transform:
            message = message_transform(message)

        return message % format_map

class as_format_mapping(object):
    """A unified, optionally transformed, mapping view over multiple instances.

    Allows regular instance attributes to be accessed by "%(attrname)s" in
    string formats.  Optionally passes values through a ``transform`` (such
    ``gettext``) before returning.

    """

    __slots__ = 'targets', 'transform'

    def __init__(self, *targets, **kw):
        self.targets = targets
        self.transform = kw.pop('transform', None)
        if kw:
            raise TypeError('unexpected keyword argument')

    def __getitem__(self, item):
        for target in self.targets:
            try:
                value = getattr(target, item)
            except AttributeError:
                pass
            else:
                if self.transform:
                    return self.transform(value)
                else:
                    return value
        raise KeyError(item)

    def __contains__(self, item):
        try:
            self[item]
            return True
        except KeyError:
            return False

    def __iter__(self):
        keys = set()
        for target in self.targets:
            keys |= set(dir(target))
        return iter(keys)

<commit_msg>Simplified ugettext finder + documented.<commit_after>"""Base functionality for fancy validation."""
import __builtin__
import operator

from flatland.util import adict, Unspecified
from flatland.schema.util import find_i18n_function

N_ = lambda translatable: translatable
_ugettext_finder = operator.attrgetter('ugettext')
_ungettext_finder = operator.attrgetter('ungettext')

class Validator(object):
    """Base class for fancy validators."""

    def __init__(self, **kw):
        """Construct a validator.

        :param \*\*kw: override any extant class attribute on this instance.

        """
        cls = type(self)
        for attr, value in kw.iteritems():
            if hasattr(cls, attr):
                setattr(self, attr, value)
            else:
                raise TypeError("%s has no attribute %r, can not override." % (
                    cls.__name__, attr))

    def __call__(self, element, state):
        """Adapts Validator to the Element.validate callable interface."""
        return self.validate(element, state)

    def validate(self, element, state):
        """Validate an element returning True if valid.

        Abstract.

        :param element:
          an :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by
          :meth:`Element.validate <flatland.schema.base.Element.validate>`.

        :returns: True if valid

        """
        return False

    def note_error(self, element, state, key=None, message=None, **info):
        """Record a validation error message on an element.

        :param element:
          An :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by :meth:`Element.validate
          <flatland.schema.base.Element.validate>`.

        :param key: semi-optional, default None.
          The name of a message-holding attribute on this instance.  Will be
          used to ``message = getattr(self, key)``.

        :param message: semi-optional, default None.  A validation
          message.  Use to provide a specific message rather than look
          one up by *key*.

        :param \*\*info: optional.
          Additional data to make available to validation message
          string formatting.

        :returns: False

        Either *key* or *message* is required.  The message will have
        formatting expanded by :meth:`expand_message` and be appended to
        :attr:`element.errors <flatland.schema.base.Element.errors>`.

        Always returns False.  This enables a convenient shorthand when
        writing validators:

        .. testcode::

          from flatland.valid import Validator

          class MyValidator(Validator):
              my_message = 'Oh noes!'

              def validate(self, element, state):
                  if not element.value:
                      return self.note_error(element, state, 'my_message')
                  else:
                      return True

        .. testcode:: :hide:

          from flatland import String
          el = String('x').create_element()
          v = MyValidator()
          assert not v.validate(el, None)
          assert el.errors == ['Oh noes!']
          el.set('foo')
          assert v.validate(el, None)
          assert el.errors == ['Oh noes!']

        """
        message = message or getattr(self, key)
        if message:
            element.add_error(
                self.expand_message(element, state, message, **info))
        return False

    def note_warning(self, element, state, key=None, message=None, **info):
        """Record a validation warning message on an element.

        :param element:
          An :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by :meth:`Element.validate
          <flatland.schema.base.Element.validate>`.

        :param key: semi-optional, default None.
          The name of a message-holding attribute on this instance.  Will be
          used to ``message = getattr(self, key)``.

        :param message: semi-optional, default None.  A validation
          message.  Use to provide a specific message rather than look
          one up by *key*.

        :param \*\*info: optional.
          Additional data to make available to validation message
          string formatting.

        :returns: False

        Either *key* or *message* is required.  The message will have
        formatting expanded by :meth:`expand_message` and be appended to
        :attr:`element.warnings <flatland.schema.base.Element.warnings>`.

        Always returns False.
        """
        message = message or getattr(self, key)
        if message:
            element.add_warning(
                self.expand_message(element, state, message, **info))
        return False

    def find_transformer(self, type, element, state, message):
        """Locate a message-transforming function, such as ugettext.

        Returns None or a callable.  The callable must return a
        message.  The call signature of the callable is expected to
        match ``ugettext`` or ``ungettext``:

        - If *type* is 'ugettext', the callable should take a message
          as a positional argument.

        - If *type* is 'ungettext', the callable should take three
          positional arguments: a message for the singular form, a
          message for the plural form, and an integer.

        Subclasses may override this method to provide advanced
        message transformation and translation functionality, on a
        per-element or per-message granularity if desired.

        The default implementation uses the following logic to locate
        a transformer:

        1.  If *state* has an attribute or item named *type*, return that.

        2.  If the *element* or any of its parents have an attribute
            named *type*, return that.

        3.  If the schema of *element* or the schema of any of its
            parents have an attribute named *type*, return that.

        4.  If *type* is in ``__builtin__``, return that.

        5.  Otherwise return ``None``.

        """
        if hasattr(state, type):
            return getattr(state, type)
        if hasattr(state, '__getitem__'):
            try:
                return state[type]
            except KeyError:
                pass

        if type == 'ugettext':
            finder = _ugettext_finder
        elif type == 'ungettext':
            finder = _ungettext_finder
        else:
            raise RuntimeError("Unknown transformation %r" % type)
        return find_i18n_function(element, finder)

    def expand_message(self, element, state, message, **extra_format_args):
        """Apply formatting to a validation message.

        :param element:
          an :class:`~flatland.schema.base.Element` instance.

        :param state:
          an arbitrary object.  Supplied by
          :meth:`Element.validate <flatland.schema.base.Element.validate>`.

        :param message: a string, 3-tuple or callable.
          If a 3-tuple, must be of the form ('single form', 'plural form',
          n_key).

          If callable, will be called with 2 positional arguments (*element*,
          *state*) and must return a string or 3-tuple.

        :param \*\*extra_format_args: optional.
          Additional data to make available to validation message
          string formatting.

        :returns: the formatted string

        See :ref:`Message Templating`, :ref:`Message Pluralization` and
        :ref:`Message Internationalization` for full information on how
        messages are expanded.

        """
        if callable(message):
            message = message(element, state)

        ugettext = self.find_transformer('ugettext', element, state, message)

        format_map = as_format_mapping(
            extra_format_args, state, self, element,
            transform=ugettext)

        if isinstance(message, tuple):
            ungettext = self.find_transformer(
                'ungettext', element, state, message)

            single, plural, n_key = message
            try:
                n = format_map[n_key]
                try:
                    n = int(n)
                except TypeError:
                    pass
            except KeyError:
                n = n_key

            if ungettext:
                message = ungettext(single, plural, n)
            else:
                if ugettext:
                    single = ugettext(single)
                    plural = ugettext(plural)
                message = single if n == 1 else plural
        elif ugettext:
            message = ugettext(message)

        return message % format_map

class as_format_mapping(object):
    """A unified, optionally transformed, mapping view over multiple instances.

    Allows regular instance attributes to be accessed by "%(attrname)s" in
    string formats.  Dictionaries may be included as well.  Optionally
    passes values through a ``transform`` (such ``gettext``) before
    returning.

    """

    __slots__ = 'targets', 'transform'

    def __init__(self, *targets, **kw):
        self.targets = [t for t in targets if t is not None]
        self.transform = kw.pop('transform', None)
        if kw:
            raise TypeError('unexpected keyword argument')

    def __getitem__(self, item):
        for target in self.targets:
            # try target[item] first
            if hasattr(target, '__getitem__'):
                try:
                    value = target[item]
                    break
                except (LookupError, TypeError):
                    pass
            # then target.item
            try:
                value = getattr(target, item)
                break
            except AttributeError:
                pass
        else:
            # not found on any target
            raise KeyError(item)

        if self.transform:
            return self.transform(value)
        else:
            return value

    def __contains__(self, item):
        try:
            self[item]
            return True
        except KeyError:
            return False

    def __iter__(self):
        keys = set()
        for target in self.targets:
            keys |= set(dir(target))
        return iter(keys)

<|endoftext|>


<commit_before>from array import array
from string import join
from struct import pack, unpack

#Functions to convert between different encodings
_DECODE = lambda x, e: list(array('B', x.decode(e)))
_ENCODE = lambda x, e: join([chr(i) for i in x], '').encode(e)
hex_to_bytes = lambda x: _DECODE(x, 'hex')
txt_to_bytes = lambda x: hex_to_bytes(x.encode('hex'))
bytes_to_hex = lambda x: _ENCODE(x, 'hex')
bytes_to_txt = lambda x: bytes_to_hex(x).decode('hex')
#Step1- The message is padded(extended) such that its length is congruent to 448 modulo 512.
#		If the message is already congruent tp 448 modulo 512, then 512 more bits are added. Padding- Single '1' and rest '0's
#Step2- 64 bit representation of the length of the message is appended to the padded message
def pad_message(msg):
	n = len(msg)
	bit_len = n * 8
	index = (bit_len >> 3) & 0x3fL
	pad_len = 120 - index
	if index < 56:
		pad_len = 56 - index
	padding = '\x80' + '\x00'*63
	padded_msg = msg + padding[:pad_len] + pack('<Q', bit_len)
	return padded_msg
#We define the following functions to be used in Step4
def _left_rotate(n, b):
	return ((n << b) | ((n & 0xffffffff) >> (32 - b))) & 0xffffffff
def _f(x, y, z): return x & y | ~x & z
def _g(x, y, z): return x & y | x & z | y & z
def _h(x, y, z): return x ^ y ^ z

def _f1(a, b, c, d, k, s, X): return _left_rotate(a + _f(b, c, d) + X[k], s)
def _f2(a, b, c, d, k, s, X): return _left_rotate(a + _g(b, c, d) + X[k] + 0x5a827999, s)
def _f3(a, b, c, d, k, s, X): return _left_rotate(a + _h(b, c, d) + X[k] + 0x6ed9eba1, s)
class MD4:
	#Step3- We use four 32 bit word buffers to compute the message digest. We initialize each of the buffers to the following hex values.
	def __init__(self):
		self.A = 0x01234567
		self.B = 0x89abcdef
		self.C = 0xfedcba98
		self.D = 0x76543210
	def update(self, message_string):
		msg_bytes = txt_to_bytes(pad_message(message_string))
		for i in range(0, len(msg_bytes), 64):
			self._compress(msg_bytes[i:i+64])
	#Step4- Process message in 16 bit word blocks. We use the previously defined functions in three rounds of compression(hashing).
	# 		We use sqrt(2) as round 2 constant and sqrt(3) as round three constant.
	def _compress(self, block):
		#We save the initial state of the buffers
		a, b, c, d = self.A, self.B, self.C, self.D
		x = []
		for i in range(0, 64, 4):
			x.append(unpack('<I', bytes_to_txt(block[i:i+4]))[0])
		#Round 1- Do following 16 operations
		a = _f1(a,b,c,d, 0, 3, x)
		d = _f1(d,a,b,c, 1, 7, x)
		c = _f1(c,d,a,b, 2,11, x)
		b = _f1(b,c,d,a, 3,19, x)
		a = _f1(a,b,c,d, 4, 3, x)
		d = _f1(d,a,b,c, 5, 7, x)
		c = _f1(c,d,a,b, 6,11, x)
		b = _f1(b,c,d,a, 7,19, x)
		a = _f1(a,b,c,d, 8, 3, x)
		d = _f1(d,a,b,c, 9, 7, x)
		c = _f1(c,d,a,b,10,11, x)
		b = _f1(b,c,d,a,11,19, x)
		a = _f1(a,b,c,d,12, 3, x)
		d = _f1(d,a,b,c,13, 7, x)
		c = _f1(c,d,a,b,14,11, x)
		b = _f1(b,c,d,a,15,19, x)
		#Round 2- Do following 16 operations
		a = _f2(a,b,c,d, 0, 3, x)
		d = _f2(d,a,b,c, 4, 5, x)
		c = _f2(c,d,a,b, 8, 9, x)
		b = _f2(b,c,d,a,12,13, x)
		a = _f2(a,b,c,d, 1, 3, x)
		d = _f2(d,a,b,c, 5, 5, x)
		c = _f2(c,d,a,b, 9, 9, x)
		b = _f2(b,c,d,a,13,13, x)
		a = _f2(a,b,c,d, 2, 3, x)
		d = _f2(d,a,b,c, 6, 5, x)
		c = _f2(c,d,a,b,10, 9, x)
		b = _f2(b,c,d,a,14,13, x)
		a = _f2(a,b,c,d, 3, 3, x)
		d = _f2(d,a,b,c, 7, 5, x)
		c = _f2(c,d,a,b,11, 9, x)
		b = _f2(b,c,d,a,15,13, x)
		#Round 3- Do following 16 operations
		a = _f3(a,b,c,d, 0, 3, x)
		d = _f3(d,a,b,c, 8, 9, x)
		c = _f3(c,d,a,b, 4,11, x)
		b = _f3(b,c,d,a,12,15, x)
		a = _f3(a,b,c,d, 2, 3, x)
		d = _f3(d,a,b,c,10, 9, x)
		c = _f3(c,d,a,b, 6,11, x)
		b = _f3(b,c,d,a,14,15, x)
		a = _f3(a,b,c,d, 1, 3, x)
		d = _f3(d,a,b,c, 9, 9, x)
		c = _f3(c,d,a,b, 5,11, x)
		b = _f3(b,c,d,a,13,15, x)
		a = _f3(a,b,c,d, 3, 3, x)
		d = _f3(d,a,b,c,11, 9, x)
		c = _f3(c,d,a,b, 7,11, x)
		b = _f3(b,c,d,a,15,15, x)

		# We add the initial states of the buffers back to the respective buffers
		self.A = (self.A + a) & 0xffffffff
		self.B = (self.B + b) & 0xffffffff
		self.C = (self.C + c) & 0xffffffff
		self.D = (self.D + d) & 0xffffffff
		#Step5- The message digest produced as output is ABCD(the buffers), starting with lower order byte of A and ending with higher order byte of D
	def digest(self):
		return bytes_to_hex(txt_to_bytes(pack('<IIII', self.A, self.B, self.C, self.D)))
#driver function
if __name__ == '__main__':

	def md4(msg):
		m = MD4()
		m.update(msg)
		print m.digest()
	print "Hash of blank message- ",
	md4("")
	print "Plaintext- a\nHash- ",
	md4("a")
	print "Plaintext- abc\nHash- ",
	md4("abc")
	print "Plaintext- abcdefghijklmnopqrstuvwxyz\nHash- ",
	md4("abcdefghijklmnopqrstuvwxyz")
	print "Plaintext- ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789\nHash- ",
	md4("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789")
	print "Plaintext- NITK\nHash- ",
	md4("NITK")
	print "Plaintext- md4\nHash- ",
	md4("md4")
	message=raw_input("Enter text to be encoded- ")
	print "Hashed message- ",
	md4(message)<commit_msg>working<commit_after>from array import array
from string import join
from struct import pack, unpack

#Functions to convert between different encodings
_DECODE = lambda x, e: list(array('B', x.decode(e)))
_ENCODE = lambda x, e: join([chr(i) for i in x], '').encode(e)
hex_to_bytes = lambda x: _DECODE(x, 'hex')
txt_to_bytes = lambda x: hex_to_bytes(x.encode('hex'))
bytes_to_hex = lambda x: _ENCODE(x, 'hex')
bytes_to_txt = lambda x: bytes_to_hex(x).decode('hex')
#Step1- The message is padded(extended) such that its length is congruent to 448 modulo 512.
#		If the message is already congruent tp 448 modulo 512, then 512 more bits are added. Padding- Single '1' and rest '0's
#Step2- 64 bit representation of the length of the message is appended to the padded message
def pad_message(msg):
	n = len(msg)
	bit_len = n * 8
	index = (bit_len >> 3) & 0x3fL
	pad_len = 120 - index
	if index < 56:
		pad_len = 56 - index
	padding = '\x80' + '\x00'*63
	padded_msg = msg + padding[:pad_len] + pack('<Q', bit_len)
	return padded_msg
#We define the following functions to be used in Step4
def _left_rotate(n, b):
	return ((n << b) | ((n & 0xffffffff) >> (32 - b))) & 0xffffffff
def _f(x, y, z): return x & y | ~x & z
def _g(x, y, z): return x & y | x & z | y & z
def _h(x, y, z): return x ^ y ^ z

def _f1(a, b, c, d, k, s, X): return _left_rotate(a + _f(b, c, d) + X[k], s)
def _f2(a, b, c, d, k, s, X): return _left_rotate(a + _g(b, c, d) + X[k] + 0x5a827999, s)
def _f3(a, b, c, d, k, s, X): return _left_rotate(a + _h(b, c, d) + X[k] + 0x6ed9eba1, s)
class MD4:
	#Step3- We use four 32 bit word buffers to compute the message digest. We initialize each of the buffers to the following hex values.
	def __init__(self):
		self.A = 0x67452301
		self.B = 0xefcdab89
		self.C = 0x98badcfe
		self.D = 0x10325476
	def update(self, message_string):
		msg_bytes = txt_to_bytes(pad_message(message_string))
		for i in range(0, len(msg_bytes), 64):
			self._compress(msg_bytes[i:i+64])
	#Step4- Process message in 16 bit word blocks. We use the previously defined functions in three rounds of compression(hashing).
	# 		We use sqrt(2) as round 2 constant and sqrt(3) as round three constant.
	def _compress(self, block):
		#We save the initial state of the buffers
		a, b, c, d = self.A, self.B, self.C, self.D
		x = []
		for i in range(0, 64, 4):
			x.append(unpack('<I', bytes_to_txt(block[i:i+4]))[0])
		#Round 1- Do following 16 operations
		a = _f1(a,b,c,d, 0, 3, x)
		d = _f1(d,a,b,c, 1, 7, x)
		c = _f1(c,d,a,b, 2,11, x)
		b = _f1(b,c,d,a, 3,19, x)
		a = _f1(a,b,c,d, 4, 3, x)
		d = _f1(d,a,b,c, 5, 7, x)
		c = _f1(c,d,a,b, 6,11, x)
		b = _f1(b,c,d,a, 7,19, x)
		a = _f1(a,b,c,d, 8, 3, x)
		d = _f1(d,a,b,c, 9, 7, x)
		c = _f1(c,d,a,b,10,11, x)
		b = _f1(b,c,d,a,11,19, x)
		a = _f1(a,b,c,d,12, 3, x)
		d = _f1(d,a,b,c,13, 7, x)
		c = _f1(c,d,a,b,14,11, x)
		b = _f1(b,c,d,a,15,19, x)
		#Round 2- Do following 16 operations
		a = _f2(a,b,c,d, 0, 3, x)
		d = _f2(d,a,b,c, 4, 5, x)
		c = _f2(c,d,a,b, 8, 9, x)
		b = _f2(b,c,d,a,12,13, x)
		a = _f2(a,b,c,d, 1, 3, x)
		d = _f2(d,a,b,c, 5, 5, x)
		c = _f2(c,d,a,b, 9, 9, x)
		b = _f2(b,c,d,a,13,13, x)
		a = _f2(a,b,c,d, 2, 3, x)
		d = _f2(d,a,b,c, 6, 5, x)
		c = _f2(c,d,a,b,10, 9, x)
		b = _f2(b,c,d,a,14,13, x)
		a = _f2(a,b,c,d, 3, 3, x)
		d = _f2(d,a,b,c, 7, 5, x)
		c = _f2(c,d,a,b,11, 9, x)
		b = _f2(b,c,d,a,15,13, x)
		#Round 3- Do following 16 operations
		a = _f3(a,b,c,d, 0, 3, x)
		d = _f3(d,a,b,c, 8, 9, x)
		c = _f3(c,d,a,b, 4,11, x)
		b = _f3(b,c,d,a,12,15, x)
		a = _f3(a,b,c,d, 2, 3, x)
		d = _f3(d,a,b,c,10, 9, x)
		c = _f3(c,d,a,b, 6,11, x)
		b = _f3(b,c,d,a,14,15, x)
		a = _f3(a,b,c,d, 1, 3, x)
		d = _f3(d,a,b,c, 9, 9, x)
		c = _f3(c,d,a,b, 5,11, x)
		b = _f3(b,c,d,a,13,15, x)
		a = _f3(a,b,c,d, 3, 3, x)
		d = _f3(d,a,b,c,11, 9, x)
		c = _f3(c,d,a,b, 7,11, x)
		b = _f3(b,c,d,a,15,15, x)

		# We add the initial states of the buffers back to the respective buffers
		self.A = (self.A + a) & 0xffffffff
		self.B = (self.B + b) & 0xffffffff
		self.C = (self.C + c) & 0xffffffff
		self.D = (self.D + d) & 0xffffffff
		#Step5- The message digest produced as output is ABCD(the buffers), starting with lower order byte of A and ending with higher order byte of D
	def digest(self):
		return bytes_to_hex(txt_to_bytes(pack('<IIII', self.A, self.B, self.C, self.D)))
#driver function
if __name__ == '__main__':

	def md4(msg):
		m = MD4()
		m.update(msg)
		print m.digest()
	print "Hash of blank message- ",
	md4("")
	print "Plaintext- a\nHash- ",
	md4("a")
	print "Plaintext- abc\nHash- ",
	md4("abc")
	print "Plaintext- abcdefghijklmnopqrstuvwxyz\nHash- ",
	md4("abcdefghijklmnopqrstuvwxyz")
	print "Plaintext- ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789\nHash- ",
	md4("ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789")
	print "Plaintext- NITK\nHash- ",
	md4("NITK")
	print "Plaintext- md4\nHash- ",
	md4("md4")
	message=raw_input("Enter text to be encoded- ")
	print "Hashed message- ",
	md4(message)<|endoftext|>


<commit_before>import asyncio
import json
import logging

from flist.chat import opcode as opcode

logger = logging.getLogger(__name__)

class FChatProtocol(object):
    """Websocket protocol with included ping handler
    connect method: return deferred when the protocol is established.

    add_op_callback: accepts a method which receives a dict object corresponding to the decoded JSON data.
    remove_op_callback

    add_message_handler: accepts a method which receives the opcode and dict of the decoded JSON data.
    remove_message_handler
    """

    def __init__(self, transport, loop=None):
        self.on_close = lambda *args: None
        self.on_open = lambda *args: None

        self.callbacks = {}
        self.handlers = []
        self.transport = transport

        self.pinger = None
        self.loop = loop or asyncio.get_event_loop()

        self.add_op_callback(opcode.PING, self._ping_handler)

    def connect(self):
        self.transport.fchat_on_message = lambda *args: self.on_message(*args)
        self.transport.fchat_on_close = lambda *args: self.on_close(*args)
        self.transport.fchat_on_open = lambda *args: self.on_open(*args)
        self.transport.connect()

    def close(self):
        self.transport.close()

    def _ping_handler(self, message):
        self.message(opcode.PING)

    @staticmethod
    def _load_json(message):
        try:
            j = json.loads(message[4:])
        except ValueError:
            j = None
        return j

    def on_message(self, message):
        op = message[:3]
        callbacks = self.callbacks.get(op, [])
        j = self._load_json(message)

        for f in callbacks.copy():
            # noinspection PyBroadException
            try:
                f(j)
            except BrokenPipeError:
                callbacks.remove(f)  # Caused by a closed provider.
            except Exception:
                callbacks.remove(f)
                logger.exception("While processing callbacks another exception"
                                 " occurred, callback function has been removed")

        for h in self.handlers:
            h(op, j)
        logger.getChild(op).info("<-- %s", message)

    def _write(self, message):
        if self.transport:
            logger.getChild(message[:3]).info("--> %s", message)
            self.transport.send_message(message)
        else:
            logger.error("Attempt to write message to Missing client.")

    def message(self, op, di=None):
        if di:
            self._write("%s %s" % (op, json.dumps(di)))
        else:
            self._write(op)

    def add_op_callback(self, op, callback):
        """Callbacks take the form of f(message)"""
        self.callbacks.setdefault(op, []).append(callback)

    def remove_op_callback(self, op, callback):
        self.callbacks.setdefault(op, []).remove(callback)

    def add_message_handler(self, handler):
        """A message handler takes the form of f(opcode, message)"""
        self.handlers.append(handler)

    def remove_message_handler(self, handler):
        self.handlers.remove(handler)

    def callback(self, status):
        def deco(func):
            self.add_op_callback(status, func)
            return func

        return deco
<commit_msg>Allow a callback to be an awaitable.<commit_after>import asyncio
import json
import logging
from inspect import isawaitable

from flist.chat import opcode as opcode

logger = logging.getLogger(__name__)

class FChatProtocol(object):
    """Websocket protocol with included ping handler
    connect method: return deferred when the protocol is established.

    add_op_callback: accepts a method which receives a dict object corresponding to the decoded JSON data.
    remove_op_callback

    add_message_handler: accepts a method which receives the opcode and dict of the decoded JSON data.
    remove_message_handler
    """

    def __init__(self, transport, loop=None):
        self.on_close = lambda *args: None
        self.on_open = lambda *args: None

        self.callbacks = {}
        self.handlers = []
        self.transport = transport

        self.pinger = None
        self.loop = loop or asyncio.get_event_loop()

        self.add_op_callback(opcode.PING, self._ping_handler)

    def connect(self):
        self.transport.fchat_on_message = lambda *args: self.on_message(*args)
        self.transport.fchat_on_close = lambda *args: self.on_close(*args)
        self.transport.fchat_on_open = lambda *args: self.on_open(*args)
        self.transport.connect()

    def close(self):
        self.transport.close()

    def _ping_handler(self, message):
        self.message(opcode.PING)

    @staticmethod
    def _load_json(message):
        try:
            j = json.loads(message[4:])
        except ValueError:
            j = None
        return j

    def on_message(self, message):
        op = message[:3]
        callbacks = self.callbacks.get(op, [])
        j = self._load_json(message)

        for f in callbacks.copy():
            # noinspection PyBroadException
            try:
                r = f(j)
                if isawaitable(r):
                    asyncio.ensure_future(r, loop=self.loop)
            except BrokenPipeError:
                callbacks.remove(f)  # Caused by a closed provider.
            except Exception:
                callbacks.remove(f)
                logger.exception("While processing callbacks another exception"
                                 " occurred, callback function has been removed")

        for h in self.handlers:
            h(op, j)
        logger.getChild(op).info("<-- %s", message)

    def _write(self, message):
        if self.transport:
            logger.getChild(message[:3]).info("--> %s", message)
            self.transport.send_message(message)
        else:
            logger.error("Attempt to write message to Missing client.")

    def message(self, op, di=None):
        if di:
            self._write("%s %s" % (op, json.dumps(di)))
        else:
            self._write(op)

    def add_op_callback(self, op, callback):
        """Callbacks take the form of f(message)"""
        self.callbacks.setdefault(op, []).append(callback)

    def remove_op_callback(self, op, callback):
        self.callbacks.setdefault(op, []).remove(callback)

    def add_message_handler(self, handler):
        """A message handler takes the form of f(opcode, message)"""
        self.handlers.append(handler)

    def remove_message_handler(self, handler):
        self.handlers.remove(handler)

    def callback(self, status):
        def deco(func):
            self.add_op_callback(status, func)
            return func

        return deco
<|endoftext|>


<commit_before># Copyright Hybrid Logic Ltd.  See LICENSE file for details.
# -*- test-case-name: flocker.node.test.test_script,flocker.node.functional.test_script -*- # noqa

"""
The command-line ``flocker-*-agent`` tools.
"""

from socket import socket
from os import getpid

import yaml

from jsonschema import FormatChecker, Draft4Validator

from pyrsistent import PRecord, field

from zope.interface import implementer

from twisted.python.filepath import FilePath
from twisted.python.usage import Options
from twisted.internet.ssl import Certificate

from ..volume.service import (
    ICommandLineVolumeScript, VolumeScript,
)

from ..volume.script import flocker_volume_options
from ..common.script import (
    ICommandLineScript,
    flocker_standard_options, FlockerScriptRunner, main_for_service)
from . import P2PManifestationDeployer, ApplicationNodeDeployer
from ._loop import AgentLoopService
from .agents.blockdevice import LoopbackBlockDeviceAPI, BlockDeviceDeployer
from ..ca import ControlServicePolicy, NodeCredential

__all__ = [
    "flocker_dataset_agent_main",
]

def _get_external_ip(host, port):
    """
    Get an external IP address for this node that can in theory connect to
    the given host and port.

    See https://clusterhq.atlassian.net/browse/FLOC-1751 for better solution.
    :param host: A host to connect to.
    :param port: The port to connect to.

    :return unicode: IP address of external interface on this node.
    """
    sock = socket()
    try:
        sock.setblocking(False)
        sock.connect_ex((host, port))
        return unicode(sock.getsockname()[0], "ascii")
    finally:
        sock.close()

class _TLSContext(PRecord):
    """
    Information extracted from the TLS certificates for this node.

    :ivar context_factory: A TLS
        context factory will validate the control service and present
        the node's certificate to the control service.

    :ivar NodeCredential node_credential: The node's certificate information.
    """
    context_factory = field(mandatory=True)
    node_credential = field(mandatory=True)

def _context_factory_and_credential(path, host, port):
    """
    Load a TLS context factory for the AMP client from the path where
    configuration and certificates live.

    The CA certificate and node private key and certificate are expected
    to be siblings of the configuration file.

    :param FilePath path: Path to directory where configuration lives.
    :param bytes host: The host we will be connecting to.
    :param int port: The port we will be connecting to.

    :return: ``_TLSContext`` instance.
    """
    ca = Certificate.loadPEM(path.child(b"cluster.crt").getContent())
    # This is a hack; from_path should be more
    # flexible. https://clusterhq.atlassian.net/browse/FLOC-1865
    node_credential = NodeCredential.from_path(path, b"node")
    policy = ControlServicePolicy(
        ca_certificate=ca, client_credential=node_credential.credential)
    return _TLSContext(context_factory=policy.creatorForNetloc(host, port),
                       node_credential=node_credential)

def validate_configuration(configuration):
    """
    Validate a provided configuration.

    :param dict configuration: A desired configuration for an agent.

    :raises: jsonschema.ValidationError if the configuration is invalid.
    """
    # XXX Create a function which loads and validates, and also setting
    # defaults. FLOC-1925.
    schema = {
        "$schema": "http://json-schema.org/draft-04/schema#",
        "type": "object",
        "required": ["version", "control-service", "dataset"],
        "properties": {
            "version": {
                "type": "number",
                "maximum": 1,
                "minimum": 1,
            },
            "control-service": {
                "type": "object",
                "required": ["hostname"],
                "properties": {
                    "hostname": {
                        "type": "string",
                        "format": "hostname",
                    },
                    "port": {"type": "integer"},
                }
            },
            "dataset": {
                "type": "object",
                "oneOf": [
                    {
                        "required": ["backend"],
                        "properties": {
                            "backend": {
                                "type": "string",
                                "pattern": "zfs",
                            },
                            "pool": {
                                "type": "string",
                            },
                        }
                    },
                    {
                        "required": ["backend"],
                        "properties": {
                            "backend": {
                                "type": "string",
                                "pattern": "loopback",
                            },
                            "pool": {
                                "type": "string",
                            },
                        }

                    },
                ]
            }
        }
    }

    v = Draft4Validator(schema, format_checker=FormatChecker())
    v.validate(configuration)

@flocker_standard_options
@flocker_volume_options
class _AgentOptions(Options):
    """
    Command line options for agents.
    """
    # Use as basis for subclass' synopsis:
    synopsis = "Usage: {} [OPTIONS]"

    optParameters = [
        ["agent-config", "c", "/etc/flocker/agent.yml",
         "The configuration file to set the control service."],
    ]

    def postOptions(self):
        self['agent-config'] = FilePath(self['agent-config'])

class DatasetAgentOptions(_AgentOptions):
    """
    Command line options for ``flocker-dataset-agent``.
    """
    longdesc = """\
    flocker-dataset-agent runs a dataset convergence agent on a node.
    """

    synopsis = _AgentOptions.synopsis.format("flocker-dataset-agent")

class ContainerAgentOptions(_AgentOptions):
    """
    Command line options for ``flocker-container-agent``.
    """
    longdesc = """\
    flocker-container-agent runs a container convergence agent on a node.
    """

    synopsis = _AgentOptions.synopsis.format("flocker-container-agent")

@implementer(ICommandLineScript)
class AgentScript(PRecord):
    """
    XXX This is temporarily not used for the ``flocker-dataset-agent`` script.
    See FLOC-1924.

    Implement top-level logic for the ``flocker-dataset-agent`` and
    ``flocker-container-agent`` scripts.

    :ivar service_factory: A two-argument callable that returns an ``IService``
        provider that will get run when this script is run.  The arguments
        passed to it are the reactor being used and a ``AgentOptions``
        instance which has parsed any command line options that were given.
    """
    service_factory = field(mandatory=True)

    def main(self, reactor, options):
        return main_for_service(
            reactor,
            self.service_factory(reactor, options)
        )

class AgentServiceFactory(PRecord):
    """
    Implement general agent setup in a way that's usable by
    ``AgentScript`` but also easily testable.

    Possibly ``ICommandLineScript`` should be replaced by something that is
    inherently more easily tested so that this separation isn't required.

    :ivar deployer_factory: A two-argument callable to create an
        ``IDeployer`` provider for this script.  The arguments are a
        ``hostname`` keyword argument, a ``cluster_uuid`` keyword and a
        ``node_uuid`` keyword argument. They must be passed by keyword.
    """
    # This should have an explicit interface:
    # https://clusterhq.atlassian.net/browse/FLOC-1929
    deployer_factory = field(mandatory=True)

    def get_service(self, reactor, options):
        """
        Create an ``AgentLoopService`` instance.

        :param reactor: The reactor to give to the service so it can schedule
            timed events and make network connections.

        :param AgentOptions options: The command-line options to use to
            configure the loop and the loop's deployer.

        :param context_factory: TLS context factory to pass to service.

        :param NodeCredential node_credential: The node credential.

        :return: The ``AgentLoopService`` instance.
        """
        agent_config = options[u'agent-config']
        configuration = yaml.safe_load(agent_config.getContent())

        validate_configuration(configuration=configuration)

        host = configuration['control-service']['hostname']
        port = configuration['control-service'].get('port', 4524)
        ip = _get_external_ip(host, port)
        tls_info = _context_factory_and_credential(
            options["agent-config"].parent(), host, port)
        return AgentLoopService(
            reactor=reactor,
            deployer=self.deployer_factory(
                node_uuid=tls_info.node_credential.uuid, hostname=ip,
                cluster_uuid=tls_info.node_credential.cluster_uuid),
            host=host, port=port,
            context_factory=tls_info.context_factory,
        )

def zfs_dataset_deployer(volume_service):
    """
    Create a deployer factory for a ZFS backend.

    :param VolumeService dataset_configuration: An already started volume
        service.

    :return: A callable which can be called with a node UUID and hostname to
        create a ZFS deployer.
    """
    def deployer_factory(hostname, node_uuid, cluster_uuid):
        return P2PManifestationDeployer(hostname=hostname, node_uuid=node_uuid,
                                        volume_service=volume_service)
    return deployer_factory

def loopback_dataset_deployer(volume_service):
    """
    Create a deployer factory for a loopback backend.

    :param VolumeService dataset_configuration: An already started volume
        service.

    :return: A callable which can be called with a node UUID and hostname to
        create a loopback deployer.
    """
    def deployer_factory(hostname, node_uuid, cluster_uuid):
        # In FLOC-1925 deployer_factory might also be called with the config
        # file, allowing for alteration of what is created and how it is
        # initialized.  That code will also want to pass cluster_uuid in
        # to relevant backend APIs, e.g. Cinder and EBS both want it.
        api = LoopbackBlockDeviceAPI.from_path(
            b"/var/lib/flocker/loopback",
            # Make up a new value every time this script starts.  This
            # will ensure different instances of the script using this
            # backend always appear to be running on different nodes (as
            # far as attachment is concerned).  This is a good thing since
            # it makes it easy to simulate a multi-node cluster by running
            # multiple instances of the script.  Similar effect could be
            # achieved by making this id a command line argument but that
            # would be harder to implement and harder to use.
            compute_instance_id=bytes(getpid()).decode('utf-8'),
        )
        return BlockDeviceDeployer(block_device_api=api, hostname=hostname,
                                   node_uuid=node_uuid)
    return deployer_factory

def dataset_deployer_from_configuration(dataset_configuration, volume_service):
    """
    Given a dataset configuration, return a dataset deployer factory.

    :param dict dataset_configuration: Desired configuration for a dataset
        deployer.
    :param VolumeService dataset_configuration: An already started volume
        service.

    :return: A callable which can be called with a node UUID and hostname to
        create a dataset deployer.
    """
    backend_to_deployer_factory = {
        'zfs': zfs_dataset_deployer,
        'loopback': loopback_dataset_deployer,
    }
    backend = dataset_configuration['backend']
    deployer_factory = backend_to_deployer_factory[backend]
    return deployer_factory(
        volume_service=volume_service
    )

@implementer(ICommandLineVolumeScript)
class GenericAgentScript(PRecord):
    """
    Implement top-level logic for the ``flocker-dataset-agent`` script.

    This is a temporary script, until the volume service can be created in
    ``zfs_dataset_deployer``. The majority of this script will be in
    ``flocker_dataset_agent_main`` and ``AgentScript``. See FLOC-1924.
    """
    def main(self, reactor, options, volume_service):
        agent_config = options[u'agent-config']
        configuration = yaml.safe_load(agent_config.getContent())

        validate_configuration(configuration=configuration)

        deployer_factory = dataset_deployer_from_configuration(
            dataset_configuration=configuration['dataset'],
            volume_service=volume_service
        )

        service_factory = AgentServiceFactory(
            deployer_factory=deployer_factory
        ).get_service

        service = service_factory(reactor, options)

        if configuration['dataset']['backend'] == 'zfs':
            # XXX This should not be a special case,
            # see https://clusterhq.atlassian.net/browse/FLOC-1924.
            volume_service.setServiceParent(service)

        return main_for_service(
            reactor=reactor,
            service=service,
        )

def flocker_dataset_agent_main():
    """
    Implementation of the ``flocker-dataset-agent`` command line script.

    This starts a dataset convergence agent.  It currently supports only the
    loopback block device backend.  Later it will be capable of starting a
    dataset agent using any of the support dataset backends.
    """
    options = DatasetAgentOptions()

    return FlockerScriptRunner(
        script=VolumeScript(GenericAgentScript()),
        options=options,
    ).main()

def flocker_container_agent_main():
    """
    Implementation of the ``flocker-container-agent`` command line script.

    This starts a Docker-based container convergence agent.
    """
    def deployer_factory(cluster_uuid, **kwargs):
        return ApplicationNodeDeployer(**kwargs)
    service_factory = AgentServiceFactory(
        deployer_factory=deployer_factory
    ).get_service
    agent_script = AgentScript(service_factory=service_factory)
    return FlockerScriptRunner(
        script=agent_script,
        options=ContainerAgentOptions()
    ).main()
<commit_msg>Sketch refactoring of node script<commit_after># Copyright Hybrid Logic Ltd.  See LICENSE file for details.
# -*- test-case-name: flocker.node.test.test_script,flocker.node.functional.test_script -*- # noqa

"""
The command-line ``flocker-*-agent`` tools.
"""

from socket import socket
from os import getpid

import yaml

from jsonschema import FormatChecker, Draft4Validator

from pyrsistent import PRecord, field

from zope.interface import implementer

from twisted.python.filepath import FilePath
from twisted.python.usage import Options
from twisted.internet.ssl import Certificate

from ..volume.service import (
    ICommandLineVolumeScript, VolumeScript,
)

from ..volume.script import flocker_volume_options
from ..common.script import (
    ICommandLineScript,
    flocker_standard_options, FlockerScriptRunner, main_for_service)
from . import P2PManifestationDeployer, ApplicationNodeDeployer
from ._loop import AgentLoopService
from .agents.blockdevice import LoopbackBlockDeviceAPI, BlockDeviceDeployer
from ..ca import ControlServicePolicy, NodeCredential

__all__ = [
    "flocker_dataset_agent_main",
]

def _get_external_ip(host, port):
    """
    Get an external IP address for this node that can in theory connect to
    the given host and port.

    See https://clusterhq.atlassian.net/browse/FLOC-1751 for better solution.
    :param host: A host to connect to.
    :param port: The port to connect to.

    :return unicode: IP address of external interface on this node.
    """
    sock = socket()
    try:
        sock.setblocking(False)
        sock.connect_ex((host, port))
        return unicode(sock.getsockname()[0], "ascii")
    finally:
        sock.close()

class _TLSContext(PRecord):
    """
    Information extracted from the TLS certificates for this node.

    :ivar context_factory: A TLS
        context factory will validate the control service and present
        the node's certificate to the control service.

    :ivar NodeCredential node_credential: The node's certificate information.
    """
    context_factory = field(mandatory=True)
    node_credential = field(mandatory=True)

def _context_factory_and_credential(path, host, port):
    """
    Load a TLS context factory for the AMP client from the path where
    configuration and certificates live.

    The CA certificate and node private key and certificate are expected
    to be siblings of the configuration file.

    :param FilePath path: Path to directory where configuration lives.
    :param bytes host: The host we will be connecting to.
    :param int port: The port we will be connecting to.

    :return: ``_TLSContext`` instance.
    """
    ca = Certificate.loadPEM(path.child(b"cluster.crt").getContent())
    # This is a hack; from_path should be more
    # flexible. https://clusterhq.atlassian.net/browse/FLOC-1865
    node_credential = NodeCredential.from_path(path, b"node")
    policy = ControlServicePolicy(
        ca_certificate=ca, client_credential=node_credential.credential)
    return _TLSContext(context_factory=policy.creatorForNetloc(host, port),
                       node_credential=node_credential)

def validate_configuration(configuration):
    """
    Validate a provided configuration.

    :param dict configuration: A desired configuration for an agent.

    :raises: jsonschema.ValidationError if the configuration is invalid.
    """
    # XXX Create a function which loads and validates, and also setting
    # defaults. FLOC-1925.
    schema = {
        "$schema": "http://json-schema.org/draft-04/schema#",
        "type": "object",
        "required": ["version", "control-service", "dataset"],
        "properties": {
            "version": {
                "type": "number",
                "maximum": 1,
                "minimum": 1,
            },
            "control-service": {
                "type": "object",
                "required": ["hostname"],
                "properties": {
                    "hostname": {
                        "type": "string",
                        "format": "hostname",
                    },
                    "port": {"type": "integer"},
                }
            },
            "dataset": {
                "type": "object",
                "oneOf": [
                    {
                        "required": ["backend"],
                        "properties": {
                            "backend": {
                                "type": "string",
                                "pattern": "zfs",
                            },
                            "pool": {
                                "type": "string",
                            },
                        }
                    },
                    {
                        "required": ["backend"],
                        "properties": {
                            "backend": {
                                "type": "string",
                                "pattern": "loopback",
                            },
                            "pool": {
                                "type": "string",
                            },
                        }

                    },
                ]
            }
        }
    }

    v = Draft4Validator(schema, format_checker=FormatChecker())
    v.validate(configuration)

@flocker_standard_options
@flocker_volume_options
class _AgentOptions(Options):
    """
    Command line options for agents.
    """
    # Use as basis for subclass' synopsis:
    synopsis = "Usage: {} [OPTIONS]"

    optParameters = [
        ["agent-config", "c", "/etc/flocker/agent.yml",
         "The configuration file to set the control service."],
    ]

    def postOptions(self):
        self['agent-config'] = FilePath(self['agent-config'])

class DatasetAgentOptions(_AgentOptions):
    """
    Command line options for ``flocker-dataset-agent``.
    """
    longdesc = """\
    flocker-dataset-agent runs a dataset convergence agent on a node.
    """

    synopsis = _AgentOptions.synopsis.format("flocker-dataset-agent")

class ContainerAgentOptions(_AgentOptions):
    """
    Command line options for ``flocker-container-agent``.
    """
    longdesc = """\
    flocker-container-agent runs a container convergence agent on a node.
    """

    synopsis = _AgentOptions.synopsis.format("flocker-container-agent")

@implementer(ICommandLineScript)
class AgentScript(PRecord):
    """
    XXX This is temporarily not used for the ``flocker-dataset-agent`` script.
    See FLOC-1924.

    Implement top-level logic for the ``flocker-dataset-agent`` and
    ``flocker-container-agent`` scripts.

    :ivar service_factory: A two-argument callable that returns an ``IService``
        provider that will get run when this script is run.  The arguments
        passed to it are the reactor being used and a ``AgentOptions``
        instance which has parsed any command line options that were given.
    """
    service_factory = field(mandatory=True)

    def main(self, reactor, options):
        return main_for_service(
            reactor,
            self.service_factory(reactor, options)
        )

class AgentServiceFactory(PRecord):
    """
    Implement general agent setup in a way that's usable by
    ``AgentScript`` but also easily testable.

    Possibly ``ICommandLineScript`` should be replaced by something that is
    inherently more easily tested so that this separation isn't required.

    :ivar deployer_factory: A two-argument callable to create an
        ``IDeployer`` provider for this script.  The arguments are a
        ``hostname`` keyword argument, a ``cluster_uuid`` keyword and a
        ``node_uuid`` keyword argument. They must be passed by keyword.
    """
    # This should have an explicit interface:
    # https://clusterhq.atlassian.net/browse/FLOC-1929
    deployer_factory = field(mandatory=True)

    def get_service(self, reactor, options):
        """
        Create an ``AgentLoopService`` instance.

        :param reactor: The reactor to give to the service so it can schedule
            timed events and make network connections.

        :param AgentOptions options: The command-line options to use to
            configure the loop and the loop's deployer.

        :param context_factory: TLS context factory to pass to service.

        :param NodeCredential node_credential: The node credential.

        :return: The ``AgentLoopService`` instance.
        """
        agent_config = options[u'agent-config']
        configuration = yaml.safe_load(agent_config.getContent())

        validate_configuration(configuration=configuration)

        host = configuration['control-service']['hostname']
        port = configuration['control-service'].get('port', 4524)
        ip = _get_external_ip(host, port)
        tls_info = _context_factory_and_credential(
            options["agent-config"].parent(), host, port)
        return AgentLoopService(
            reactor=reactor,
            deployer=self.deployer_factory(
                node_uuid=tls_info.node_credential.uuid, hostname=ip,
                cluster_uuid=tls_info.node_credential.cluster_uuid),
            host=host, port=port,
            context_factory=tls_info.context_factory,
        )

def zfs_dataset_deployer(volume_service):
    """
    Create a deployer factory for a ZFS backend.

    :param VolumeService dataset_configuration: An already started volume
        service.

    :return: A callable which can be called with a node UUID and hostname to
        create a ZFS deployer.
    """
    def deployer_factory(hostname, node_uuid, cluster_uuid):
        return P2PManifestationDeployer(hostname=hostname, node_uuid=node_uuid,
                                        volume_service=volume_service)
    return deployer_factory

def loopback_dataset_deployer(volume_service):
    """
    Create a deployer factory for a loopback backend.

    :param VolumeService dataset_configuration: An already started volume
        service.

    :return: A callable which can be called with a node UUID and hostname to
        create a loopback deployer.
    """
    def deployer_factory(hostname, node_uuid, cluster_uuid):
        # In FLOC-1925 deployer_factory might also be called with the config
        # file, allowing for alteration of what is created and how it is
        # initialized.  That code will also want to pass cluster_uuid in
        # to relevant backend APIs, e.g. Cinder and EBS both want it.
        api = LoopbackBlockDeviceAPI.from_path(
            b"/var/lib/flocker/loopback",
            # Make up a new value every time this script starts.  This
            # will ensure different instances of the script using this
            # backend always appear to be running on different nodes (as
            # far as attachment is concerned).  This is a good thing since
            # it makes it easy to simulate a multi-node cluster by running
            # multiple instances of the script.  Similar effect could be
            # achieved by making this id a command line argument but that
            # would be harder to implement and harder to use.
            compute_instance_id=bytes(getpid()).decode('utf-8'),
        )
        return BlockDeviceDeployer(block_device_api=api, hostname=hostname,
                                   node_uuid=node_uuid)
    return deployer_factory

def dataset_deployer_from_configuration(dataset_configuration, volume_service):
    """
    Given a dataset configuration, return a dataset deployer factory.

    :param dict dataset_configuration: Desired configuration for a dataset
        deployer.
    :param VolumeService dataset_configuration: An already started volume
        service.

    :return: A callable which can be called with a node UUID and hostname to
        create a dataset deployer.
    """
    backend_to_deployer_factory = {
        'zfs': zfs_dataset_deployer,
        'loopback': loopback_dataset_deployer,
    }
    backend = dataset_configuration['backend']
    deployer_factory = backend_to_deployer_factory[backend]
    return deployer_factory(
        volume_service=volume_service
    )

@implementer(ICommandLineVolumeScript)
class GenericAgentScript(PRecord):
    """
    Implement top-level logic for the ``flocker-dataset-agent`` script.

    This is a temporary script, until the volume service can be created in
    ``zfs_dataset_deployer``. The majority of this script will be in
    ``flocker_dataset_agent_main`` and ``AgentScript``. See FLOC-1924.
    """
    def main(self, reactor, options, volume_service):
        agent_config = options[u'agent-config']
        configuration = yaml.safe_load(agent_config.getContent())

        validate_configuration(configuration=configuration)

        deployer_factory = dataset_deployer_from_configuration(
            dataset_configuration=configuration['dataset'],
            volume_service=volume_service
        )

        service_factory = AgentServiceFactory(
            deployer_factory=deployer_factory
        ).get_service

        service = service_factory(reactor, options)

        if configuration['dataset']['backend'] == 'zfs':
            # XXX This should not be a special case,
            # see https://clusterhq.atlassian.net/browse/FLOC-1924.
            volume_service.setServiceParent(service)

        return main_for_service(
            reactor=reactor,
            service=service,
        )

def flocker_dataset_agent_main():
    """
    Implementation of the ``flocker-dataset-agent`` command line script.

    This starts a dataset convergence agent.  It currently supports only the
    loopback block device backend.  Later it will be capable of starting a
    dataset agent using any of the support dataset backends.
    """
    options = DatasetAgentOptions()

    return FlockerScriptRunner(
        script=VolumeScript(GenericAgentScript()),
        options=options,
    ).main()

def flocker_container_agent_main():
    """
    Implementation of the ``flocker-container-agent`` command line script.

    This starts a Docker-based container convergence agent.
    """
    def deployer_factory(cluster_uuid, **kwargs):
        return ApplicationNodeDeployer(**kwargs)
    service_factory = AgentServiceFactory(
        deployer_factory=deployer_factory
    ).get_service
    agent_script = AgentScript(service_factory=service_factory)
    return FlockerScriptRunner(
        script=agent_script,
        options=ContainerAgentOptions()
    ).main()

def get_configuration(options):
    agent_config = options[u'agent-config']
    configuration = yaml.safe_load(agent_config.getContent())

    validate_configuration(configuration=configuration)

    path = agent_config.parent()
    ca = Certificate.loadPEM(path.child(b"cluster.crt").getContent())
    # This is a hack; from_path should be more
    # flexible. https://clusterhq.atlassian.net/browse/FLOC-1865
    configuration['node-credential'] = NodeCredential.from_path(path, b"node")

    return configuration

from flocker.volume import zfs

# These structures should be created dynamically to handle plug-ins
default_backends = {
    'zfs': (zfs.StoragePool, {}),
    'loopback': (LoopbackBlockDeviceAPI.from_path, {
        'root_path': b"/var/lib/flocker/loopback",
        'compute_instance_id': bytes(getpid()).decode('utf-8'),
        })
    # 'openstack': (CinderBlockDeviceAPI.from_path, {
    #     'root_path': b"/var/lib/flocker/loopback",
    #     'compute_instance_id': bytes(getpid()).decode('utf-8'),
    #     })
}
default_deployers = {
    'zfs': P2PManifestationDeployer,
    'loopback': BlockDeviceDeployer,
    # 'openstack': BlockDeviceDeployer,
}

class AgentService:

    def __init__(
            self, configuration, backends=default_backends,
            deployers=default_deployers):
        self._configuration = configuration
        self._backends = backends
        self._deployers = deployers
        self._host = configuration['control-service']['hostname']
        self._port = configuration['control-service'].get('port', 4524)
        node_credential = configuration['node-credential']
        policy = ControlServicePolicy(
            ca_certificate=ca, client_credential=node_credential.credential)
        self._tls_context = _TLSContext(
            context_factory=policy.creatorForNetloc(self._host, self._port),
            node_credential=node_credential)
        self._backend = self._configuration['dataset']['backend']

    def get_api(self):
        (api_factory, api_args) = self._backends[self._backend]
        return api_factory(**api_args)

    def get_deployer(self, api):
        deployer_factory = self._deployers[self._backend]
        hostname = _get_external_ip(self._host, self._port)
        node_uuid = self._tls_context.node_credential.uuid
        # cluster_uuid = self._tls_context.node_credential.cluster_uuid
        return deployer_factory(
            block_device_api=api, hostname=hostname, node_uuid=node_uuid)

    def get_loop_service(self, deployer, reactor):
        return AgentLoopService(
            reactor=reactor,
            deployer=deployer,
            host=self._host, port=self._port,
            context_factory=self._tls_context.context_factory,
        )

class NewAgentScript(PRecord):

    agent_service_factory = field(type=object, initial=AgentService)
    configuration_factory = field(type=callable, initial=get_configuration)
    api_factory = field(type=callable, initial=get_api)
    deployer_factory = field(type=callable, initial=get_deployer)
    service_factory = field(type=callable, initial=get_loop_service)
    post_loop = field(type=callable, initial=lambda loop: None)

    def main(self, reactor, options):

        configuration = self.configuration_factory(options)

        agent_service = agent_service_factory(configuration)

        api = agent_service.get_api()

        host = configuration['control-service']['hostname']
        port = configuration['control-service'].get('port', 4524)
        tls_info = _context_factory_and_credential(
            options["agent-config"].parent(), host, port)

        deployer = self.deployer_factory(configuration, api, tls_info)

        loop_service = self.service_factory(deployer)

        api.set_service(loop_service)

        return main_for_service(
            reactor,
            loop_service
        )

def new_flocker_dataset_agent_main():
    """
    Implementation of the ``flocker-dataset-agent`` command line script.

    This starts a dataset convergence agent.  It currently supports only the
    loopback block device backend.  Later it will be capable of starting a
    dataset agent using any of the supported dataset backends.
    """
    options = DatasetAgentOptions()

    return FlockerScriptRunner(
        script=NewAgentScript(),
        options=options,
    ).main()
<|endoftext|>