<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from collections import namedtuple

import requests
import urllib
from urlparse import urljoin, urlparse

from django.conf import settings
from django.contrib.auth import login as django_login
from django.contrib.auth import logout as django_logout
from django.contrib.auth.backends import ModelBackend
from django.contrib.auth.decorators import login_required
from django.http import (
    HttpResponseBadRequest,
    HttpResponseRedirect,
    HttpResponseRedirectBase
)
from django.utils import simplejson
from django.http import HttpResponse
from django.views.generic.base import View
from django.utils.decorators import method_decorator

from itsdangerous import URLSafeTimedSerializer

from lizard_auth_client.client import construct_user

# used so we can login User objects we instantiated ourselves
BACKEND = ModelBackend()

class HttpResponseServiceUnavailable(HttpResponse):
    status_code = 503

class TestHomeView(View):
    '''
    Test view, only used in standalone mode.
    Display a minimal HTML page with some URLS.
    '''
    def get(self, request, *args, **kwargs):
        user = request.user
        return HttpResponse(
            'home | protected'
            '| logout | user={} | home @ client'
            ''.format(user)
        )

class TestProtectedView(View):
    '''
    Test view, only used in standalone mode.
    Display a minimal HTML page with some URLS, but is "protected" by
    a login_required decorator.
    '''
    @method_decorator(login_required)
    def dispatch(self, request, *args, **kwargs):
        return super(TestProtectedView, self).dispatch(
            request, *args, **kwargs)

    def get(self, request, *args, **kwargs):
        user = request.user
        return HttpResponse(
            'home | protected'
            '| logout '
            '| user={} | protected @ client'
            ''.format(user)
        )

class LoginApiView(View):
    '''
    Login API, which can optionally be included in the urls of the
    root Django app. This allows a non-webbrowser client, like
    JavaScripts XmlHttpRequest, to implement HTTP redirects in their
    own custom way.
    '''
    def get(self, request, *args, **kwargs):
        # Redirect to the webclient after the SSO server dance
        request.session['sso_after_login_next'] = settings.WEBCLIENT

        # Get the login url with the token
        wrapped_response = get_request_token_and_determine_response()

        if issubclass(wrapped_response.http_response,
                      HttpResponseRedirectBase):
            # The response is a redirect (302) to the SSO server.
            # Wrap it in a normal HttpResponse, and have client-side code
            # handle the actual redirect.
            response_class = HttpResponse
            content_dict = {'login_url': wrapped_response.url}
        else:
            # Response is something else, like an error message.
            # Use the response class as-is and wrap the message
            # in JSON.
            response_class = wrapped_response.http_response
            content_dict = {'message': wrapped_response.message}

        content = simplejson.dumps(content_dict)
        return response_class(content=content, content_type='application/json')

class LogoutApiView(View):
    '''
    Logout API, which can optionally be included in the urls of the
    root Django app. This allows a non-webbrowser client, like
    JavaScripts XmlHttpRequest, to implement HTTP redirects in their
    own custom way.
    '''
    def get(self, request, *args, **kwargs):
        # Redirect to the webclient after the SSO server dance
        request.session['sso_after_logout_next'] = settings.WEBCLIENT

        # Simple wrap the logout url in a JSON dict
        logout_url = build_sso_portal_action_url('logout')
        content_dict = {'logout_url': logout_url}
        content = simplejson.dumps(content_dict)
        return HttpResponse(content=content, content_type='application/json')

class LoginView(View):
    '''
    View that redirects the user to the SSO server.

    Requests a Request Token and then redirects the User to the the SSO Server.
    '''
    def get(self, request, *args, **kwargs):
        # store the GET parameter named 'next', so we can redirect the user
        # to the requested page, after SSO login.
        next = get_next(request)
        request.session['sso_after_login_next'] = next

        wrapped_response = get_request_token_and_determine_response()

        if isinstance(wrapped_response.http_response,
                      HttpResponseRedirect):
            return wrapped_response.http_response(
                wrapped_response.url
            )
        else:
            return wrapped_response.http_response(
                wrapped_response.message
            )

class LocalLoginView(View):
    '''
    Verifies the user token with the SSO server, and logs the user in.
    '''
    def get(self, request, *args, **kwargs):
        # verify the authentication token and
        # retrieve the User instance from the SSO server
        user = verify_auth_token(request.GET['message'])
        if not user:
            return HttpResponseBadRequest('Verification failed')

        # link the user instance to the default database backend
        # and call django-auth's login function
        user.backend = "%s.%s" % (BACKEND.__module__,
                                  BACKEND.__class__.__name__)
        django_login(request, user)

        # redirect the user to the stored "next" url, which is probably a
        # protected page
        sso_after_login_next = request.session.get('sso_after_login_next', '/')
        request.session.delete('sso_after_login_next')
        return HttpResponseRedirect(sso_after_login_next)

class LogoutView(View):
    '''
    Redirect user to SSO server, to log out there.
    '''
    def get(self, request, *args, **kwargs):
        # store the 'next' parameter in the session so we can
        # redirect the user afterwards
        next = get_next(request)
        request.session['sso_after_logout_next'] = next

        url = build_sso_portal_action_url('logout')
        # send the redirect response
        return HttpResponseRedirect(url)

class LocalLogoutView(View):
    '''
    Log out locally. Mostly the user is redirected here by the SSO server.
    '''
    def get(self, request, *args, **kwargs):
        # get the redirect url
        next = request.session.get('sso_after_logout_next', '/')

        # django_logout also calls session.flush()
        django_logout(request)

        return HttpResponseRedirect(next)

def get_request_token():
    '''
    Requests a Request Token from the SSO Server. Returns False if the request
    failed.
    '''
    # construct a signed message containing the portal key
    params = {
        'key': settings.SSO_KEY
    }

    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    url = urljoin(settings.SSO_SERVER_PRIVATE_URL,
                  'sso/api/request_token') + '/'

    # send the message to the SSO server
    response = requests.get(
        url,
        params={
            'key': settings.SSO_KEY,
            'message': message
        },
        timeout=10
    )
    if response.status_code != 200:
        return False

    # grab the token from the response
    data = URLSafeTimedSerializer(settings.SSO_SECRET).loads(
        response.content, max_age=300)
    return data['request_token']

def verify_auth_token(untrusted_message):
    '''
    Verifies a Auth Token. Returns a
    django.contrib.auth.models.User instance if successful or False.
    '''
    # decrypt the message
    untrusted = URLSafeTimedSerializer(settings.SSO_SECRET).loads(
        untrusted_message, max_age=300)

    # do some extra validation
    if 'auth_token' not in untrusted:
        return False
    if 'request_token' not in untrusted:
        return False

    # call the SSO server to verify the token
    params = {
        'auth_token': untrusted['auth_token'],
        'key': settings.SSO_KEY
    }
    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    url = urljoin(settings.SSO_SERVER_PRIVATE_URL, 'sso/api/verify') + '/'
    response = requests.get(
        url,
        params={
            'key': settings.SSO_KEY,
            'message': message
        },
        timeout=10
    )

    # ensure the response is sane
    if response.status_code != 200:
        return False

    # build a User object from the message
    data = URLSafeTimedSerializer(settings.SSO_SECRET).loads(
        response.content, max_age=300)
    user_data = simplejson.loads(data['user'])
    return construct_user(user_data)

def get_next(request):
    '''
    Given a request, returns the URL where a user should be redirected to
    after login. Defaults to '/'.
    '''
    next = request.GET.get('next', None)
    if not next:
        return '/'
    netloc = urlparse(next)[1]

    # security check -- don't allow redirection to a different host
    # taken from django.contrib.auth.views.login
    if netloc and netloc != request.get_host():
        return '/'
    return next

def get_request_token_and_determine_response():
    '''
    Retrieve a Request token from the SSO server, and determine the proper
    HttpResponse to send to the user.

    When logging using via the REST API, the response is wrapped in JSON,
    because the redirection takes place in a client-side script.
    '''
    WrappedResponse = namedtuple('WrappedResponse', 'http_response, message, redirect_url')

    # get a request token, which is used by the SSO server to verify
    # that the user is allowed to make a login request
    request_token = get_request_token()
    if not request_token:
        # Status code 503 service (sso server) unavailable
        return WrappedResponse(HttpResponseServiceUnavailable, 'Unable to obtain token', None)

    # construct a (signed) set of GET parameters which are used to
    # redirect the user to the SSO server
    params = {
        'request_token': request_token,
        'key': settings.SSO_KEY
    }
    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    query_string = urllib.urlencode([('message', message),
                                     ('key', settings.SSO_KEY)])
    # build an absolute URL pointing to the SSO server out of it
    url = urljoin(settings.SSO_SERVER_PUBLIC_URL, 'sso/authorize') + '/'
    url = '%s?%s' % (url, query_string)

    return WrappedResponse(HttpResponseRedirect, 'OK', url)

def build_sso_portal_action_url(action):
    '''
    Constructs and signs a message containing the specified action parameter,
    and returns a URL which can be used to redirect the user.

    For example, with action='logout', this can be used to logout the user
    on the SSO server.
    '''
    params = {
        'action': action,
        'key': settings.SSO_KEY
    }
    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    query_string = urllib.urlencode([('message', message),
                                     ('key', settings.SSO_KEY)])
    url = urljoin(settings.SSO_SERVER_PUBLIC_URL, 'sso/portal_action') + '/'
    url = '%s?%s' % (url, query_string)
    return url
<commit_msg>oops, forgot to commit the rest<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals

from collections import namedtuple

import requests
import urllib
from urlparse import urljoin, urlparse

from django.conf import settings
from django.contrib.auth import login as django_login
from django.contrib.auth import logout as django_logout
from django.contrib.auth.backends import ModelBackend
from django.contrib.auth.decorators import login_required
from django.http import (
    HttpResponseBadRequest,
    HttpResponseRedirect,
    HttpResponseRedirectBase
)
from django.utils import simplejson
from django.http import HttpResponse
from django.views.generic.base import View
from django.utils.decorators import method_decorator

from itsdangerous import URLSafeTimedSerializer

from lizard_auth_client.client import construct_user

# used so we can login User objects we instantiated ourselves
BACKEND = ModelBackend()

class HttpResponseServiceUnavailable(HttpResponse):
    status_code = 503

class TestHomeView(View):
    '''
    Test view, only used in standalone mode.
    Display a minimal HTML page with some URLS.
    '''
    def get(self, request, *args, **kwargs):
        user = request.user
        return HttpResponse(
            'home | protected'
            '| logout | user={} | home @ client'
            ''.format(user)
        )

class TestProtectedView(View):
    '''
    Test view, only used in standalone mode.
    Display a minimal HTML page with some URLS, but is "protected" by
    a login_required decorator.
    '''
    @method_decorator(login_required)
    def dispatch(self, request, *args, **kwargs):
        return super(TestProtectedView, self).dispatch(
            request, *args, **kwargs)

    def get(self, request, *args, **kwargs):
        user = request.user
        return HttpResponse(
            'home | protected'
            '| logout '
            '| user={} | protected @ client'
            ''.format(user)
        )

class LoginApiView(View):
    '''
    Login API, which can optionally be included in the urls of the
    root Django app. This allows a non-webbrowser client, like
    JavaScripts XmlHttpRequest, to implement HTTP redirects in their
    own custom way.
    '''
    def get(self, request, *args, **kwargs):
        # Redirect to the webclient after the SSO server dance
        request.session['sso_after_login_next'] = settings.WEBCLIENT

        # Get the login url with the token
        wrapped_response = get_request_token_and_determine_response()

        if issubclass(wrapped_response.http_response,
                      HttpResponseRedirectBase):
            # The response is a redirect (302) to the SSO server.
            # Wrap it in a normal HttpResponse, and have client-side code
            # handle the actual redirect.
            response_class = HttpResponse
            content_dict = {'login_url': wrapped_response.redirect_url}
        else:
            # Response is something else, like an error message.
            # Use the response class as-is and wrap the message
            # in JSON.
            response_class = wrapped_response.http_response
            content_dict = {'message': wrapped_response.message}

        content = simplejson.dumps(content_dict)
        return response_class(content=content, content_type='application/json')

class LogoutApiView(View):
    '''
    Logout API, which can optionally be included in the urls of the
    root Django app. This allows a non-webbrowser client, like
    JavaScripts XmlHttpRequest, to implement HTTP redirects in their
    own custom way.
    '''
    def get(self, request, *args, **kwargs):
        # Redirect to the webclient after the SSO server dance
        request.session['sso_after_logout_next'] = settings.WEBCLIENT

        # Simple wrap the logout url in a JSON dict
        logout_url = build_sso_portal_action_url('logout')
        content_dict = {'logout_url': logout_url}
        content = simplejson.dumps(content_dict)
        return HttpResponse(content=content, content_type='application/json')

class LoginView(View):
    '''
    View that redirects the user to the SSO server.

    Requests a Request Token and then redirects the User to the the SSO Server.
    '''
    def get(self, request, *args, **kwargs):
        # store the GET parameter named 'next', so we can redirect the user
        # to the requested page, after SSO login.
        next = get_next(request)
        request.session['sso_after_login_next'] = next

        wrapped_response = get_request_token_and_determine_response()

        if issubclass(wrapped_response.http_response,
                      HttpResponseRedirectBase):
            return wrapped_response.http_response(
                wrapped_response.redirect_url
            )
        else:
            return wrapped_response.http_response(
                wrapped_response.message
            )

class LocalLoginView(View):
    '''
    Verifies the user token with the SSO server, and logs the user in.
    '''
    def get(self, request, *args, **kwargs):
        # verify the authentication token and
        # retrieve the User instance from the SSO server
        user = verify_auth_token(request.GET['message'])
        if not user:
            return HttpResponseBadRequest('Verification failed')

        # link the user instance to the default database backend
        # and call django-auth's login function
        user.backend = "%s.%s" % (BACKEND.__module__,
                                  BACKEND.__class__.__name__)
        django_login(request, user)

        # redirect the user to the stored "next" url, which is probably a
        # protected page
        sso_after_login_next = request.session.get('sso_after_login_next', '/')
        request.session.delete('sso_after_login_next')
        return HttpResponseRedirect(sso_after_login_next)

class LogoutView(View):
    '''
    Redirect user to SSO server, to log out there.
    '''
    def get(self, request, *args, **kwargs):
        # store the 'next' parameter in the session so we can
        # redirect the user afterwards
        next = get_next(request)
        request.session['sso_after_logout_next'] = next

        url = build_sso_portal_action_url('logout')
        # send the redirect response
        return HttpResponseRedirect(url)

class LocalLogoutView(View):
    '''
    Log out locally. Mostly the user is redirected here by the SSO server.
    '''
    def get(self, request, *args, **kwargs):
        # get the redirect url
        next = request.session.get('sso_after_logout_next', '/')

        # django_logout also calls session.flush()
        django_logout(request)

        return HttpResponseRedirect(next)

def get_request_token():
    '''
    Requests a Request Token from the SSO Server. Returns False if the request
    failed.
    '''
    # construct a signed message containing the portal key
    params = {
        'key': settings.SSO_KEY
    }

    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    url = urljoin(settings.SSO_SERVER_PRIVATE_URL,
                  'sso/api/request_token') + '/'

    # send the message to the SSO server
    response = requests.get(
        url,
        params={
            'key': settings.SSO_KEY,
            'message': message
        },
        timeout=10
    )
    if response.status_code != 200:
        return False

    # grab the token from the response
    data = URLSafeTimedSerializer(settings.SSO_SECRET).loads(
        response.content, max_age=300)
    return data['request_token']

def verify_auth_token(untrusted_message):
    '''
    Verifies a Auth Token. Returns a
    django.contrib.auth.models.User instance if successful or False.
    '''
    # decrypt the message
    untrusted = URLSafeTimedSerializer(settings.SSO_SECRET).loads(
        untrusted_message, max_age=300)

    # do some extra validation
    if 'auth_token' not in untrusted:
        return False
    if 'request_token' not in untrusted:
        return False

    # call the SSO server to verify the token
    params = {
        'auth_token': untrusted['auth_token'],
        'key': settings.SSO_KEY
    }
    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    url = urljoin(settings.SSO_SERVER_PRIVATE_URL, 'sso/api/verify') + '/'
    response = requests.get(
        url,
        params={
            'key': settings.SSO_KEY,
            'message': message
        },
        timeout=10
    )

    # ensure the response is sane
    if response.status_code != 200:
        return False

    # build a User object from the message
    data = URLSafeTimedSerializer(settings.SSO_SECRET).loads(
        response.content, max_age=300)
    user_data = simplejson.loads(data['user'])
    return construct_user(user_data)

def get_next(request):
    '''
    Given a request, returns the URL where a user should be redirected to
    after login. Defaults to '/'.
    '''
    next = request.GET.get('next', None)
    if not next:
        return '/'
    netloc = urlparse(next)[1]

    # security check -- don't allow redirection to a different host
    # taken from django.contrib.auth.views.login
    if netloc and netloc != request.get_host():
        return '/'
    return next

def get_request_token_and_determine_response():
    '''
    Retrieve a Request token from the SSO server, and determine the proper
    HttpResponse to send to the user.

    When logging using via the REST API, the response is wrapped in JSON,
    because the redirection takes place in a client-side script.
    '''
    WrappedResponse = namedtuple('WrappedResponse', 'http_response, message, redirect_url')

    # get a request token, which is used by the SSO server to verify
    # that the user is allowed to make a login request
    request_token = get_request_token()
    if not request_token:
        # Status code 503 service (sso server) unavailable
        return WrappedResponse(HttpResponseServiceUnavailable, 'Unable to obtain token', None)

    # construct a (signed) set of GET parameters which are used to
    # redirect the user to the SSO server
    params = {
        'request_token': request_token,
        'key': settings.SSO_KEY
    }
    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    query_string = urllib.urlencode([('message', message),
                                     ('key', settings.SSO_KEY)])
    # build an absolute URL pointing to the SSO server out of it
    url = urljoin(settings.SSO_SERVER_PUBLIC_URL, 'sso/authorize') + '/'
    url = '%s?%s' % (url, query_string)

    return WrappedResponse(HttpResponseRedirect, 'OK', url)

def build_sso_portal_action_url(action):
    '''
    Constructs and signs a message containing the specified action parameter,
    and returns a URL which can be used to redirect the user.

    For example, with action='logout', this can be used to logout the user
    on the SSO server.
    '''
    params = {
        'action': action,
        'key': settings.SSO_KEY
    }
    message = URLSafeTimedSerializer(settings.SSO_SECRET).dumps(params)
    query_string = urllib.urlencode([('message', message),
                                     ('key', settings.SSO_KEY)])
    url = urljoin(settings.SSO_SERVER_PUBLIC_URL, 'sso/portal_action') + '/'
    url = '%s?%s' % (url, query_string)
    return url
<|endoftext|>


<commit_before>#
# Copyright (c) 2004 Specifix, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program is distributed in the hope that it will be useful, but
# without any waranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#
"""
Command-line option handling
"""

import log
import sys

(NO_PARAM,   # arg may occur, no parameter
 ONE_PARAM,  # arg may occur once, req'd parameter
 OPT_PARAM,  # arg may occur once, optional parameter
 MULT_PARAM, # arg may occur N times, w/ parameter
 ) = range(0,4)

class OptionError(Exception):
    def __init__(self, val):
        self.val = val

def processArgs(argDef, cfgMap, cfg, usage, argv=sys.argv):
    otherArgs = [ argv[0] ]
    argSet = {}
    # don't mangle the command line
    argv = argv[:]

    for arg in cfgMap.keys():
	argDef[arg] = 1

    i = 1
    while i < len(argv):
	if argv[i][:2] != "--":
	    otherArgs.append(argv[i])
        # stop processing args after --
        elif argv[i] == '--':
                otherArgs.extend(argv[i+1:])
                break
	else:
            arg = argv[i][2:]
            arg_parts = arg.split('=', 1)
            if len(arg_parts) > 1:
                arg = arg_parts[0]
                # don't allow --foo=bar arg if foo doesn't exist
                # or doesn't take an arg.
                if not argDef.has_key(arg) and argDef[arg] != NO_PARAM:
                    raise OptionError(usage())
                argv[i] = arg
                argv.insert(i+1, arg_parts[1])
	    if not argDef.has_key(arg): raise OptionError(usage())

	    if argDef[arg] == NO_PARAM:
		argSet[arg] = True
	    elif argDef[arg] == OPT_PARAM:
		# max one setting
		if argSet.has_key(arg): raise OptionError(usage())
		if i >= len(argv):
		    argSet[arg] = True
                else:
                    next_arg = argv[i+1]
                    if next_arg == '':
                        argSet[arg] = True
                        i = i + 1
                    elif next_arg[0:2] == '--':
                        argSet[arg] = True
                    else:
                        argSet[arg] = next_arg
                        i = i + 1
	    else:
		# the argument takes a parameter
		i = i + 1
		if i >= len(argv): raise OptionError(usage())

		if argDef[arg] == ONE_PARAM:
		    # exactly one parameter is allowd
		    if argSet.has_key(arg): raise OptionError(usage())
		    argSet[arg] = argv[i]
		else:
		    # multiple parameters may occur
		    if argSet.has_key(arg):
			argSet[arg].append(argv[i])
		    else:
			argSet[arg] = [argv[i]]
	i = i + 1

    if 'config-file' in argSet:
	cfg.read(argSet['config-file'])
	del argSet['config-file']

    if argSet.has_key('config'):
	for param in argSet['config']:
	    cfg.configLine(param)

	del argSet['config']

    for (arg, name) in cfgMap.items():
	if argSet.has_key(arg):
	    cfg.configLine("%s %s" % (name, argSet[arg]))
	    del argSet[arg]

    if argSet.has_key('debug'):
	del argSet['debug']
	import pdb
	pdb.set_trace()

    if '-v' in otherArgs:
	otherArgs.remove('-v')
	log.setVerbosity(1)
    else:
	log.setVerbosity(0)

    return argSet, otherArgs
<commit_msg>use epdb<commit_after>#
# Copyright (c) 2004 Specifix, Inc.
#
# This program is distributed under the terms of the Common Public License,
# version 1.0. A copy of this license should have been distributed with this
# source file in a file called LICENSE. If it is not present, the license
# is always available at http://www.opensource.org/licenses/cpl.php.
#
# This program is distributed in the hope that it will be useful, but
# without any waranty; without even the implied warranty of merchantability
# or fitness for a particular purpose. See the Common Public License for
# full details.
#
"""
Command-line option handling
"""

import log
import sys

(NO_PARAM,   # arg may occur, no parameter
 ONE_PARAM,  # arg may occur once, req'd parameter
 OPT_PARAM,  # arg may occur once, optional parameter
 MULT_PARAM, # arg may occur N times, w/ parameter
 ) = range(0,4)

class OptionError(Exception):
    def __init__(self, val):
        self.val = val

def processArgs(argDef, cfgMap, cfg, usage, argv=sys.argv):
    otherArgs = [ argv[0] ]
    argSet = {}
    # don't mangle the command line
    argv = argv[:]

    for arg in cfgMap.keys():
	argDef[arg] = 1

    i = 1
    while i < len(argv):
	if argv[i][:2] != "--":
	    otherArgs.append(argv[i])
        # stop processing args after --
        elif argv[i] == '--':
                otherArgs.extend(argv[i+1:])
                break
	else:
            arg = argv[i][2:]
            arg_parts = arg.split('=', 1)
            if len(arg_parts) > 1:
                arg = arg_parts[0]
                # don't allow --foo=bar arg if foo doesn't exist
                # or doesn't take an arg.
                if not argDef.has_key(arg) and argDef[arg] != NO_PARAM:
                    raise OptionError(usage())
                argv[i] = arg
                argv.insert(i+1, arg_parts[1])
	    if not argDef.has_key(arg): raise OptionError(usage())

	    if argDef[arg] == NO_PARAM:
		argSet[arg] = True
	    elif argDef[arg] == OPT_PARAM:
		# max one setting
		if argSet.has_key(arg): raise OptionError(usage())
		if i >= len(argv):
		    argSet[arg] = True
                else:
                    next_arg = argv[i+1]
                    if next_arg == '':
                        argSet[arg] = True
                        i = i + 1
                    elif next_arg[0:2] == '--':
                        argSet[arg] = True
                    else:
                        argSet[arg] = next_arg
                        i = i + 1
	    else:
		# the argument takes a parameter
		i = i + 1
		if i >= len(argv): raise OptionError(usage())

		if argDef[arg] == ONE_PARAM:
		    # exactly one parameter is allowd
		    if argSet.has_key(arg): raise OptionError(usage())
		    argSet[arg] = argv[i]
		else:
		    # multiple parameters may occur
		    if argSet.has_key(arg):
			argSet[arg].append(argv[i])
		    else:
			argSet[arg] = [argv[i]]
	i = i + 1

    if 'config-file' in argSet:
	cfg.read(argSet['config-file'])
	del argSet['config-file']

    if argSet.has_key('config'):
	for param in argSet['config']:
	    cfg.configLine(param)

	del argSet['config']

    for (arg, name) in cfgMap.items():
	if argSet.has_key(arg):
	    cfg.configLine("%s %s" % (name, argSet[arg]))
	    del argSet[arg]

    if argSet.has_key('debug'):
	del argSet['debug']
	import epdb
	epdb.set_trace()

    if '-v' in otherArgs:
	otherArgs.remove('-v')
	log.setVerbosity(1)
    else:
	log.setVerbosity(0)

    return argSet, otherArgs
<|endoftext|>


<commit_before>import lib.util
import lib.menu
import lib.term
import os
import subprocess
import getpass
import fnmatch

git_env = os.environ.copy()
git_env["GIT_SSH"] = os.path.join(os.getcwd(),"ssh_wrapper.sh")

def cache_path():
    '''Return relative path to repo cache'''
    return os.path.join('.cache','repos')

class project:

    def __init__(self,name,url,location,remote_user=getpass.getuser(),remote_user_key=None):
        self.name     = name
        self.url      = url
        self.location = location
        self.user     = remote_user
        self.key      = remote_user_key
        self.fetched  = False

    def clone(self):
        '''Clones project into cache, unless it's already there'''
        lib.util.mkdir_p(cache_path())
        if not os.path.isdir( os.path.join( self.get_cache_dir(), '.git' ) ):
            lib.term.print_c("Cloning...\n",lib.term.BLUE)
            subprocess.call(["git","clone",self.url],cwd=cache_path(),env=git_env)

    def fetch(self):
        '''Fetches all remote data'''
        self.clone()
        if not self.fetched:
            lib.term.print_c("Fetching....\n",lib.term.BLUE)
            self.fetched = True
            subprocess.call(["git","fetch","-v","--all"],
                            cwd=self.get_cache_dir(),
                            env=git_env)

    def branches(self):
        '''Lists all current remote braches'''
        self.fetch()
        subprocess.call(["git","remote","prune","origin"],cwd=self.get_cache_dir())
        out = str(subprocess.check_output(["git","branch","-r"],
                                          cwd=self.get_cache_dir(),
                                          env=git_env),'utf8')
        branches = []
        for branch in out.split("\n"):
            if not "HEAD" in branch:
                try:
                    branches.append(branch.split("/")[1])
                except Exception: pass
        return branches

    def current_branch(self):
        '''Returns the name of the current branch'''
        cwd = self.get_cache_dir()
        branch = str(subprocess.check_output(["git","rev-parse","--abbrev-ref","HEAD"],cwd=cwd),'utf8')
        return branch.rstrip()

    def current_commit(self):
        '''Returns the name of the current commit'''
        cwd = self.get_cache_dir()
        branch = str(subprocess.check_output(["git","rev-parse","HEAD"],cwd=cwd),'utf8')
        return branch.rstrip()

    def current_commit_message(self):
        '''Returns the current commit message'''
        cwd = self.get_cache_dir()
        branch = str(subprocess.check_output(["git","log","-1","--pretty=%B"],cwd=cwd),'utf8')
        return branch.rstrip()

    def checkout(self,branch):
        '''Checks out the given branch in the local cache repo, does a hard reset'''
        self.fetch()
        cwd = self.get_cache_dir()
        lib.term.print_c("Checking out....\n",lib.term.BLUE)
        with open(os.devnull) as null:
            subprocess.call(["git","branch","-f",branch],stdout=null,stderr=null,cwd=cwd)
            subprocess.call(["git","checkout",branch],stdout=null,stderr=null,cwd=cwd)
            subprocess.call(["git","reset","--hard","origin/"+branch],cwd=cwd)
            subprocess.call(["git","clean","-f","-d"],cwd=cwd)

    def tag(self,tagname):
        '''Tags whatever commit the project is on and attempts to push that to the remote repo'''
        cwd = self.get_cache_dir()
        subprocess.call(["git","tag",tagname],cwd=cwd)
        subprocess.call(["git","push","--tags"],cwd=cwd)

    def get_cache_dir(self):
        '''Return relative path to project's repo cache'''
        return os.path.join(cache_path(),self.name)

    def get_snap_dir(self):
        '''Returns full path to project's snap directory'''
        return os.path.join(self.get_cache_dir(),"snap")

    def choose_and_checkout_branch(self):
        '''Gives user list of branches to choose from, and checks out the chosen one'''
        branches = {}
        for b in self.branches():
            branches[b] = b
        branch = lib.menu.navigate("Choose a branch from {0}".format(self.name),branches)
        self.checkout(branch)
        return branch

    def snap_script(self,script):
        '''Runs a script in the snap directory'''
        cache_path = self.get_cache_dir()
        snap_path  = self.get_snap_dir()
        fn_abs = os.path.join( snap_path, script )
        fn_rel = os.path.join( "snap",    script )
        if os.path.isfile(fn_abs):
            os.system("chmod +x {0}".format(fn_abs))
            lib.util.command_check_stderr([fn_rel], cwd=cache_path)

    def get_nosnap(self):
        '''Returns a nosnap message that exists in the root of the project, or None'''
        cache_path = self.get_cache_dir()
        root_contents = os.listdir(cache_path)
        for f in fnmatch.filter(root_contents,'*.nosnap'):
            fullf = os.path.join(cache_path,f)
            if os.path.isfile(fullf):
                s = ""
                with open(fullf,'r') as fh:
                    for l in fh: s += l
                return s
        return None
<commit_msg>made it possible to revert back to a tag<commit_after>import lib.util
import lib.menu
import lib.term
import os
import subprocess
import getpass
import fnmatch

git_env = os.environ.copy()
git_env["GIT_SSH"] = os.path.join(os.getcwd(),"ssh_wrapper.sh")

def cache_path():
    '''Return relative path to repo cache'''
    return os.path.join('.cache','repos')

class project:

    def __init__(self,name,url,location,remote_user=getpass.getuser(),remote_user_key=None):
        self.name     = name
        self.url      = url
        self.location = location
        self.user     = remote_user
        self.key      = remote_user_key
        self.fetched  = False

    def clone(self):
        '''Clones project into cache, unless it's already there'''
        lib.util.mkdir_p(cache_path())
        if not os.path.isdir( os.path.join( self.get_cache_dir(), '.git' ) ):
            lib.term.print_c("Cloning...\n",lib.term.BLUE)
            subprocess.call(["git","clone",self.url],cwd=cache_path(),env=git_env)

    def fetch(self):
        '''Fetches all remote data'''
        self.clone()
        if not self.fetched:
            lib.term.print_c("Fetching....\n",lib.term.BLUE)
            self.fetched = True
            subprocess.call(["git","fetch","-pv","--all"],
                            cwd=self.get_cache_dir(),
                            env=git_env)

    def branches(self):
        '''Lists all current remote braches'''
        self.fetch()
        subprocess.call(["git","remote","prune","origin"],cwd=self.get_cache_dir())
        out = str(subprocess.check_output(["git","branch","-r"],
                                          cwd=self.get_cache_dir(),
                                          env=git_env),'utf8')
        branches = []
        for branch in out.split("\n"):
            if not "HEAD" in branch:
                try:
                    branches.append(branch.split("/")[1])
                except Exception: pass
        return branches

    def tags(self):
        '''List all current tags for the project'''
        subprocess.call(["git","fetch","-pvt","--all"],
                        cwd=self.get_cache_dir(),
                        env=git_env)
        out = str(subprocess.check_output(["git","tag","-l"],cwd=self.get_cache_dir()),'utf8')
        return out.split("\n")

    def current_branch(self):
        '''Returns the name of the current branch'''
        cwd = self.get_cache_dir()
        branch = str(subprocess.check_output(["git","rev-parse","--abbrev-ref","HEAD"],cwd=cwd),'utf8')
        return branch.rstrip()

    def current_commit(self):
        '''Returns the name of the current commit'''
        cwd = self.get_cache_dir()
        branch = str(subprocess.check_output(["git","rev-parse","HEAD"],cwd=cwd),'utf8')
        return branch.rstrip()

    def current_commit_message(self):
        '''Returns the current commit message'''
        cwd = self.get_cache_dir()
        branch = str(subprocess.check_output(["git","log","-1","--pretty=%B"],cwd=cwd),'utf8')
        return branch.rstrip()

    def checkout(self,branch):
        '''Checks out the given branch in the local cache repo, does a hard reset'''
        self.fetch()
        cwd = self.get_cache_dir()
        lib.term.print_c("Checking out....\n",lib.term.BLUE)
        with open(os.devnull) as null:
            subprocess.call(["git","checkout",branch],stdout=null,stderr=null,cwd=cwd)
            subprocess.call(["git","reset","--hard",branch],cwd=cwd)
            subprocess.call(["git","clean","-f","-d","-x"],cwd=cwd)

    def tag(self,tagname):
        '''Tags whatever commit the project is on and attempts to push that to the remote repo'''
        cwd = self.get_cache_dir()
        subprocess.call(["git","tag",tagname],cwd=cwd)
        subprocess.call(["git","push","--tags"],cwd=cwd)

    def get_cache_dir(self):
        '''Return relative path to project's repo cache'''
        return os.path.join(cache_path(),self.name)

    def get_snap_dir(self):
        '''Returns full path to project's snap directory'''
        return os.path.join(self.get_cache_dir(),"snap")

    def choose_and_checkout_branch(self):
        '''Gives user list of branches to choose from, and checks out the chosen one'''
        branches = {}
        for b in self.branches():
            branches[b] = b
        branches["# Choose a tag instead #"] = True
        branch = lib.menu.navigate("Choose a branch from {0}".format(self.name),branches)

        if branch is True:
            tags = {}
            for t in self.tags():
                if t: tags[t] = t

            if not tags:
                lib.term.big_error("There are no tags associated with this repo\n")
                lib.term.big_error("Press enter to go back\n")
                lib.term.readline()
                return self.choose_and_checkout_branch()

            branch = lib.menu.navigate("Choose a tag from {0}".format(self.name),tags)
        else:
            branch = "origin/"+branch

        self.checkout(branch)
        return branch

    def snap_script(self,script):
        '''Runs a script in the snap directory'''
        cache_path = self.get_cache_dir()
        snap_path  = self.get_snap_dir()
        fn_abs = os.path.join( snap_path, script )
        fn_rel = os.path.join( "snap",    script )
        if os.path.isfile(fn_abs):
            os.system("chmod +x {0}".format(fn_abs))
            lib.util.command_check_stderr([fn_rel], cwd=cache_path)

    def get_nosnap(self):
        '''Returns a nosnap message that exists in the root of the project, or None'''
        cache_path = self.get_cache_dir()
        root_contents = os.listdir(cache_path)
        for f in fnmatch.filter(root_contents,'*.nosnap'):
            fullf = os.path.join(cache_path,f)
            if os.path.isfile(fullf):
                s = ""
                with open(fullf,'r') as fh:
                    for l in fh: s += l
                return s
        return None
<|endoftext|>


<commit_before>import json
import re
import scrapy
import logging
from locations.items import GeojsonPointItem

def process_hours(hours):
    if 'Hours' not in hours:
        return None

    opening_hours = ''
    days = hours['Hours']
    for day in days:
        shortname = day['ShortName']
        if '-' in shortname:
            startday, endday = shortname.split('-')
            if startday == "M": startday = "Mo"
            if endday == "M": endday = "Mo"
            shortname = '-'.join([startday, endday])
        timeperiod = day['TimePeriod']
        starttime = timeperiod['BeginTime'][:-3]
        endtime = timeperiod['ThruTime'][:-3]
        time_together = '-'.join([starttime, endtime])
        opening_hours += "{shortname} {time_together};".format(
            shortname=shortname,
            time_together=time_together
        )
    return opening_hours[:-1]

class TargetSpider(scrapy.Spider):
    ''' The Target api allows us a maximum search radius of 625 miles, so we have
        to use multiple api requests with zip codes from different states to
        search the whole country.

        NOTE: There is a "key" variable used in the url for the requests.
        It seems that this does not vary between natural requests (those sent
        from a browser), but it might be something that expires and renders requests
        invalid. So, if this spider stops working, that's probably why.
    '''
    zips = ['97062', '90021', '73301', '82633', '29401', '60007', '04019', '87101', '59601', '39056', '14201', '55111', '94203', '64030', '33601']
    name = "target"
    allowed_domains = ["target.com"]
    start_urls = tuple([
        'https://api.target.com/v2/store?nearby={zipcode}&range=625&limit=999999&locale=en-US&key=eb2551e4accc14f38cc42d32fbc2b2ea&callback=jQuery21409262105278815946_1512699093216&_=1512699093219'.format(zipcode=z) for z in zips
    ])

    def parse(self, response):
        data = json.loads(response.body_as_unicode()[41:-1])
        stores = data['Locations']['Location']
        for store in stores:
            loc_info = store['Address']
            properties = {
                'ref': store['ID'],
                'name': store['Name'],
                'addr_full': loc_info['AddressLine1'],
                'city': loc_info['City'],
                'state': loc_info['Subdivision'],
                'country': loc_info['CountryName'],
                'postcode': loc_info['PostalCode'],
                'lat': loc_info['Latitude'],
                'lon': loc_info['Longitude'],
            }

            phones = store['TelephoneNumber']
            if type(phones) == "list":
                for i in phones:
                    if i['FunctionalTypeDescription'] == 'Main':
                        properties['phone'] = i['PhoneNumber']
            elif type(phones) == 'dict':
                if phones['FunctionalTypeDescription'] == 'Main':
                    properties['phone'] = phones['PhoneNumber']

            if 'OperatingHours' in store:
                processed_hours = process_hours(store['OperatingHours'])
                if processed_hours:
                    properties['opening_hours'] = processed_hours

            yield GeojsonPointItem(**properties)

<commit_msg>Added a few zip codes so that it scrapes all stores<commit_after>import json
import re
import scrapy
import logging
from locations.items import GeojsonPointItem

def process_hours(hours):
    if 'Hours' not in hours:
        return None

    opening_hours = ''
    days = hours['Hours']
    for day in days:
        shortname = day['ShortName']
        if '-' in shortname:
            startday, endday = shortname.split('-')
            if startday == "M": startday = "Mo"
            if endday == "M": endday = "Mo"
            shortname = '-'.join([startday, endday])
        timeperiod = day['TimePeriod']
        starttime = timeperiod['BeginTime'][:-3]
        endtime = timeperiod['ThruTime'][:-3]
        time_together = '-'.join([starttime, endtime])
        opening_hours += "{shortname} {time_together};".format(
            shortname=shortname,
            time_together=time_together
        )
    return opening_hours[:-1]

class TargetSpider(scrapy.Spider):
    ''' The Target api allows us a maximum search radius of 625 miles, so we have
        to use multiple api requests with zip codes from different states to
        search the whole country.

        NOTE: There is a "key" variable used in the url for the requests.
        It seems that this does not vary between natural requests (those sent
        from a browser), but it might be something that expires and renders requests
        invalid. So, if this spider stops working, that's probably why.
    '''
    zips = ['97062', '90021', '73301', '82633', '29401', '60007', '04019', '87101', '59601', '39056', '14201', '55111', '94203', '64030', '33601', '10001', '98101', '37011', '58501', '78501', '97501']
    name = "target"
    allowed_domains = ["target.com"]
    start_urls = tuple([
        'https://api.target.com/v2/store?nearby={zipcode}&range=625&limit=999999&locale=en-US&key=eb2551e4accc14f38cc42d32fbc2b2ea&callback=jQuery21409262105278815946_1512699093216&_=1512699093219'.format(zipcode=z) for z in zips
    ])

    def parse(self, response):
        data = json.loads(response.body_as_unicode()[41:-1])
        stores = data['Locations']['Location']
        for store in stores:
            loc_info = store['Address']
            properties = {
                'ref': store['ID'],
                'name': store['Name'],
                'addr_full': loc_info['AddressLine1'],
                'city': loc_info['City'],
                'state': loc_info['Subdivision'],
                'country': loc_info['CountryName'],
                'postcode': loc_info['PostalCode'],
                'lat': loc_info['Latitude'],
                'lon': loc_info['Longitude'],
            }

            phones = store['TelephoneNumber']
            if type(phones) == "list":
                for i in phones:
                    if i['FunctionalTypeDescription'] == 'Main':
                        properties['phone'] = i['PhoneNumber']
            elif type(phones) == 'dict':
                if phones['FunctionalTypeDescription'] == 'Main':
                    properties['phone'] = phones['PhoneNumber']

            if 'OperatingHours' in store:
                processed_hours = process_hours(store['OperatingHours'])
                if processed_hours:
                    properties['opening_hours'] = processed_hours

            yield GeojsonPointItem(**properties)

<|endoftext|>


<commit_before>"""
Estimate nonlinear normalization parameters from Freesurfer conformed space
to FSL's nonlinear MNI152 target using FLIRT and FNIRT.

See the docstring for the create_normalization_workflow function for more
information about the processing.

"""
import os
import os.path as op
from moss.mosaic import Mosaic

from nipype import (IdentityInterface,
                    SelectFiles, DataSink,
                    Node, Workflow)
from nipype.interfaces import fsl, freesurfer as fs
from nipype.interfaces.base import (BaseInterface,
                                    BaseInterfaceInputSpec,
                                    TraitedSpec, File)

from lyman.tools import submit_cmdline

def create_fsl_workflow(data_dir=None, subjects=None, name="fslwarp"):
    """Set up the anatomical normalzation workflow using FNIRT.

    Your anatomical data must have been processed in Freesurfer.
    Unlike most lyman workflows, the DataGrabber and DataSink
    nodes are hardwired within the returned workflow, as this
    tightly integrates with the Freesurfer subjects directory
    structure.

    Parameters
    ----------
    data_dir : path
        top level of data hierarchy/FS subjects directory
    subjects : list of strings
        list of subject IDs
    name : alphanumeric string, optional
        workflow name

    """
    if data_dir is None:
        data_dir = os.environ["SUBJECTS_DIR"]
    if subjects is None:
        subjects = []

    # Get target images
    target_brain = fsl.Info.standard_image("avg152T1_brain.nii.gz")
    target_head = fsl.Info.standard_image("avg152T1.nii.gz")
    hires_head = fsl.Info.standard_image("MNI152_T1_1mm.nii.gz")
    target_mask = fsl.Info.standard_image(
        "MNI152_T1_2mm_brain_mask_dil.nii.gz")
    fnirt_cfg = os.path.join(
        os.environ["FSLDIR"], "etc/flirtsch/T1_2_MNI152_2mm.cnf")

    # Subject source node
    subjectsource = Node(IdentityInterface(fields=["subject_id"]),
                         iterables=("subject_id", subjects),
                         name="subjectsource")

    # Grab recon-all outputs
    templates = dict(aseg="{subject_id}/mri/aparc+aseg.mgz",
                     head="{subject_id}/mri/T1.mgz")
    datasource = Node(SelectFiles(templates,
                                  base_directory=data_dir),
                      "datasource")

    # Convert images to nifti storage and float representation
    cvtaseg = Node(fs.MRIConvert(out_type="niigz"), "convertaseg")

    cvthead = Node(fs.MRIConvert(out_type="niigz",
                                 out_datatype="float"),
                   "converthead")

    # Turn the aparc+aseg into a brainmask
    makemask = Node(fs.Binarize(dilate=1, min=0.5), "makemask")

    # Extract the brain from the orig.mgz using the mask
    skullstrip = Node(fsl.ApplyMask(), "skullstrip")

    # FLIRT brain to MNI152_brain
    flirt = Node(fsl.FLIRT(reference=target_brain), "flirt")

    sw = [-180, 180]
    for dim in ["x", "y", "z"]:
        setattr(flirt.inputs, "searchr_%s" % dim, sw)

    # FNIRT head to MNI152
    fnirt = Node(fsl.FNIRT(ref_file=target_head,
                           refmask_file=target_mask,
                           config_file=fnirt_cfg,
                           fieldcoeff_file=True),
                 "fnirt")

    # Warp and rename the images
    warpbrain = Node(fsl.ApplyWarp(ref_file=target_head,
                                   interp="spline",
                                   out_file="brain_warp.nii.gz"),
                     "warpbrain")

    warpbrainhr = Node(fsl.ApplyWarp(ref_file=hires_head,
                                     interp="spline",
                                     out_file="brain_warp_hires.nii.gz"),
                       "warpbrainhr")

    # Generate a png summarizing the registration
    warpreport = Node(WarpReport(), "warpreport")

    # Save relevant files to the data directory
    datasink = Node(DataSink(base_directory=data_dir,
                             parameterization=False,
                             substitutions=[
                                 ("orig_out_masked_flirt.mat", "affine.mat"),
                                 ("orig_out_fieldwarp", "warpfield"),
                                 ("orig_out_masked", "brain"),
                                 ("orig_out", "T1")]),
                    "datasink")

    # Define and connect the workflow
    # -------------------------------

    normalize = Workflow(name=name)

    normalize.connect([
        (subjectsource, datasource,
            [("subject_id", "subject_id")]),
        (datasource, cvtaseg,
            [("aseg", "in_file")]),
        (datasource, cvthead,
            [("head", "in_file")]),
        (cvtaseg, makemask,
            [("out_file", "in_file")]),
        (cvthead, skullstrip,
            [("out_file", "in_file")]),
        (makemask, skullstrip,
            [("binary_file", "mask_file")]),
        (skullstrip, flirt,
            [("out_file", "in_file")]),
        (flirt, fnirt,
            [("out_matrix_file", "affine_file")]),
        (cvthead, fnirt,
            [("out_file", "in_file")]),
        (skullstrip, warpbrain,
            [("out_file", "in_file")]),
        (fnirt, warpbrain,
            [("fieldcoeff_file", "field_file")]),
        (skullstrip, warpbrainhr,
            [("out_file", "in_file")]),
        (fnirt, warpbrainhr,
            [("fieldcoeff_file", "field_file")]),
        (warpbrain, warpreport,
            [("out_file", "in_file")]),
        (subjectsource, datasink,
            [("subject_id", "container")]),
        (skullstrip, datasink,
            [("out_file", "normalization.@brain")]),
        (cvthead, datasink,
            [("out_file", "normalization.@t1")]),
        (flirt, datasink,
            [("out_file", "normalization.@brain_flirted")]),
        (flirt, datasink,
            [("out_matrix_file", "normalization.@affine")]),
        (warpbrain, datasink,
            [("out_file", "normalization.@brain_warped")]),
        (warpbrainhr, datasink,
            [("out_file", "normalization.@brain_hires")]),
        (fnirt, datasink,
            [("fieldcoeff_file", "normalization.@warpfield")]),
        (warpreport, datasink,
            [("out_file", "normalization.@report")]),
    ])

    return normalize

def create_ants_workflow(data_dir=None, subjects=None, name="antswarp"):
    """Set up the anatomical normalzation workflow using ANTS.

    Your anatomical data must have been processed in Freesurfer.
    Unlike most lyman workflows, the DataGrabber and DataSink
    nodes are hardwired within the returned workflow, as this
    tightly integrates with the Freesurfer subjects directory
    structure.

    Parameters
    ----------
    data_dir : path
        top level of data hierarchy/FS subjects directory
    subjects : list of strings
        list of subject IDs
    name : alphanumeric string, optional
        workflow name

    """
    if data_dir is None:
        data_dir = os.environ["SUBJECTS_DIR"]
    if subjects is None:
        subjects = []

    # Subject source node
    subjectsource = Node(IdentityInterface(fields=["subject_id"]),
                         iterables=("subject_id", subjects),
                         name="subjectsource")

    # Grab recon-all outputs
    templates = dict(aseg="{subject_id}/mri/aparc+aseg.mgz",
                     head="{subject_id}/mri/brain.mgz")
    datasource = Node(SelectFiles(templates,
                                  base_directory=data_dir),
                      "datasource")

    # Convert images to nifti storage and float representation
    cvtaseg = Node(fs.MRIConvert(out_type="niigz"), "convertaseg")

    cvtbrain = Node(fs.MRIConvert(out_type="niigz",
                                  out_datatype="float"),
                    "convertbrain")

    # Turn the aparc+aseg into a brainmask
    makemask = Node(fs.Binarize(dilate=4, erode=3, min=0.5), "makemask")

    # Extract the brain from the orig.mgz using the mask
    skullstrip = Node(fsl.ApplyMask(), "skullstrip")

    # Normalize using ANTS
    antswarp = Node(ANTSIntroduction(), "antswarp")

    # Generate a png summarizing the registration
    warpreport = Node(WarpReport(), "warpreport")

    # Save relevant files to the data directory
    datasink = Node(DataSink(base_directory=data_dir,
                             parameterization=False),
                    name="datasink")

    # Define and connect the workflow
    # -------------------------------

    normalize = Workflow(name=name)

    normalize.connect([
        (subjectsource, datasource,
            [("subject_id", "subject_id")]),
        (datasource, cvtaseg,
            [("aseg", "in_file")]),
        (datasource, cvtbrain,
            [("head", "in_file")]),
        (cvtaseg, makemask,
            [("out_file", "in_file")]),
        (cvtbrain, skullstrip,
            [("out_file", "in_file")]),
        (makemask, skullstrip,
            [("binary_file", "mask_file")]),
        (skullstrip, antswarp,
            [("out_file", "in_file")]),
        (antswarp, warpreport,
            [("brain_file", "in_file")]),
        (subjectsource, datasink,
            [("subject_id", "container")]),
        (antswarp, datasink,
            [("warp_file", "normalization.@warpfield"),
             ("inv_warp_file", "normalization.@inverse_warpfield"),
             ("affine_file", "normalization.@affine"),
             ("brain_file", "normalization.@brain")]),
        (warpreport, datasink,
            [("out_file", "normalization.@report")]),
    ])

    return normalize

class ANTSIntroductionInputSpec(BaseInterfaceInputSpec):

    in_file = File(exists=True)

class ANTSIntroductionOutputSpec(TraitedSpec):

    brain_file = File(exitsts=True)
    warp_file = File(exitsts=True)
    inv_warp_file = File(exitsts=True)
    affine_file = File(exitsts=True)

class ANTSIntroduction(BaseInterface):

    input_spec = ANTSIntroductionInputSpec
    output_spec = ANTSIntroductionOutputSpec

    def _run_interface(self, runtime):

        target_brain = fsl.Info.standard_image("avg152T1_brain.nii.gz")
        cmdline = " ".join(["antsIntroduction.sh",
                            "-d", "3",
                            "-r", target_brain,
                            "-i", self.inputs.in_file,
                            "-o", "ants_"])

        runtime.environ["ITK_NUM_THREADS"] = "1"
        runtime = submit_cmdline(runtime, cmdline)

        os.rename("ants_affine.txt", "affine.txt")
        os.rename("ants_Warp.nii.gz", "warpfield.nii.gz")
        os.rename("ants_deformed.nii.gz", "brain_warp.nii.gz")
        os.rename("ants_InverseWarp.nii.gz", "inverse_warpfield.nii.gz")

        return runtime

    def _list_outputs(self):

        outputs = self._outputs().get()
        outputs["affine_file"] = op.abspath("affine.txt")
        outputs["warp_file"] = op.abspath("warpfield.nii.gz")
        outputs["brain_file"] = op.abspath("brain_warp.nii.gz")
        outputs["inv_warp_file"] = op.abspath("inverse_warpfield.nii.gz")

        return outputs

class WarpReportInput(BaseInterfaceInputSpec):

    in_file = File(exists=True)

class WarpReportOutput(TraitedSpec):

    out_file = File(exitsts=True)

class WarpReport(BaseInterface):

    input_spec = WarpReportInput
    output_spec = WarpReportOutput

    def _run_interface(self, runtime):

        target_brain = fsl.Info.standard_image("avg152T1_brain.nii.gz")
        m = Mosaic(self.inputs.in_file, target_brain)
        m.plot_contours("Reds", 2)
        m.savefig("warp_report.png")
        m.close()

        return runtime

    def _list_outputs(self):

        outputs = self._outputs().get()
        outputs["out_file"] = op.realpath("warp_report.png")
        return outputs
<commit_msg>Fix ants output name<commit_after>"""
Estimate nonlinear normalization parameters from Freesurfer conformed space
to FSL's nonlinear MNI152 target using FLIRT and FNIRT.

See the docstring for the create_normalization_workflow function for more
information about the processing.

"""
import os
import os.path as op
from moss.mosaic import Mosaic

from nipype import (IdentityInterface,
                    SelectFiles, DataSink,
                    Node, Workflow)
from nipype.interfaces import fsl, freesurfer as fs
from nipype.interfaces.base import (BaseInterface,
                                    BaseInterfaceInputSpec,
                                    TraitedSpec, File)

from lyman.tools import submit_cmdline

def create_fsl_workflow(data_dir=None, subjects=None, name="fslwarp"):
    """Set up the anatomical normalzation workflow using FNIRT.

    Your anatomical data must have been processed in Freesurfer.
    Unlike most lyman workflows, the DataGrabber and DataSink
    nodes are hardwired within the returned workflow, as this
    tightly integrates with the Freesurfer subjects directory
    structure.

    Parameters
    ----------
    data_dir : path
        top level of data hierarchy/FS subjects directory
    subjects : list of strings
        list of subject IDs
    name : alphanumeric string, optional
        workflow name

    """
    if data_dir is None:
        data_dir = os.environ["SUBJECTS_DIR"]
    if subjects is None:
        subjects = []

    # Get target images
    target_brain = fsl.Info.standard_image("avg152T1_brain.nii.gz")
    target_head = fsl.Info.standard_image("avg152T1.nii.gz")
    hires_head = fsl.Info.standard_image("MNI152_T1_1mm.nii.gz")
    target_mask = fsl.Info.standard_image(
        "MNI152_T1_2mm_brain_mask_dil.nii.gz")
    fnirt_cfg = os.path.join(
        os.environ["FSLDIR"], "etc/flirtsch/T1_2_MNI152_2mm.cnf")

    # Subject source node
    subjectsource = Node(IdentityInterface(fields=["subject_id"]),
                         iterables=("subject_id", subjects),
                         name="subjectsource")

    # Grab recon-all outputs
    templates = dict(aseg="{subject_id}/mri/aparc+aseg.mgz",
                     head="{subject_id}/mri/T1.mgz")
    datasource = Node(SelectFiles(templates,
                                  base_directory=data_dir),
                      "datasource")

    # Convert images to nifti storage and float representation
    cvtaseg = Node(fs.MRIConvert(out_type="niigz"), "convertaseg")

    cvthead = Node(fs.MRIConvert(out_type="niigz",
                                 out_datatype="float"),
                   "converthead")

    # Turn the aparc+aseg into a brainmask
    makemask = Node(fs.Binarize(dilate=1, min=0.5), "makemask")

    # Extract the brain from the orig.mgz using the mask
    skullstrip = Node(fsl.ApplyMask(), "skullstrip")

    # FLIRT brain to MNI152_brain
    flirt = Node(fsl.FLIRT(reference=target_brain), "flirt")

    sw = [-180, 180]
    for dim in ["x", "y", "z"]:
        setattr(flirt.inputs, "searchr_%s" % dim, sw)

    # FNIRT head to MNI152
    fnirt = Node(fsl.FNIRT(ref_file=target_head,
                           refmask_file=target_mask,
                           config_file=fnirt_cfg,
                           fieldcoeff_file=True),
                 "fnirt")

    # Warp and rename the images
    warpbrain = Node(fsl.ApplyWarp(ref_file=target_head,
                                   interp="spline",
                                   out_file="brain_warp.nii.gz"),
                     "warpbrain")

    warpbrainhr = Node(fsl.ApplyWarp(ref_file=hires_head,
                                     interp="spline",
                                     out_file="brain_warp_hires.nii.gz"),
                       "warpbrainhr")

    # Generate a png summarizing the registration
    warpreport = Node(WarpReport(), "warpreport")

    # Save relevant files to the data directory
    datasink = Node(DataSink(base_directory=data_dir,
                             parameterization=False,
                             substitutions=[
                                 ("orig_out_masked_flirt.mat", "affine.mat"),
                                 ("orig_out_fieldwarp", "warpfield"),
                                 ("orig_out_masked", "brain"),
                                 ("orig_out", "T1")]),
                    "datasink")

    # Define and connect the workflow
    # -------------------------------

    normalize = Workflow(name=name)

    normalize.connect([
        (subjectsource, datasource,
            [("subject_id", "subject_id")]),
        (datasource, cvtaseg,
            [("aseg", "in_file")]),
        (datasource, cvthead,
            [("head", "in_file")]),
        (cvtaseg, makemask,
            [("out_file", "in_file")]),
        (cvthead, skullstrip,
            [("out_file", "in_file")]),
        (makemask, skullstrip,
            [("binary_file", "mask_file")]),
        (skullstrip, flirt,
            [("out_file", "in_file")]),
        (flirt, fnirt,
            [("out_matrix_file", "affine_file")]),
        (cvthead, fnirt,
            [("out_file", "in_file")]),
        (skullstrip, warpbrain,
            [("out_file", "in_file")]),
        (fnirt, warpbrain,
            [("fieldcoeff_file", "field_file")]),
        (skullstrip, warpbrainhr,
            [("out_file", "in_file")]),
        (fnirt, warpbrainhr,
            [("fieldcoeff_file", "field_file")]),
        (warpbrain, warpreport,
            [("out_file", "in_file")]),
        (subjectsource, datasink,
            [("subject_id", "container")]),
        (skullstrip, datasink,
            [("out_file", "normalization.@brain")]),
        (cvthead, datasink,
            [("out_file", "normalization.@t1")]),
        (flirt, datasink,
            [("out_file", "normalization.@brain_flirted")]),
        (flirt, datasink,
            [("out_matrix_file", "normalization.@affine")]),
        (warpbrain, datasink,
            [("out_file", "normalization.@brain_warped")]),
        (warpbrainhr, datasink,
            [("out_file", "normalization.@brain_hires")]),
        (fnirt, datasink,
            [("fieldcoeff_file", "normalization.@warpfield")]),
        (warpreport, datasink,
            [("out_file", "normalization.@report")]),
    ])

    return normalize

def create_ants_workflow(data_dir=None, subjects=None, name="antswarp"):
    """Set up the anatomical normalzation workflow using ANTS.

    Your anatomical data must have been processed in Freesurfer.
    Unlike most lyman workflows, the DataGrabber and DataSink
    nodes are hardwired within the returned workflow, as this
    tightly integrates with the Freesurfer subjects directory
    structure.

    Parameters
    ----------
    data_dir : path
        top level of data hierarchy/FS subjects directory
    subjects : list of strings
        list of subject IDs
    name : alphanumeric string, optional
        workflow name

    """
    if data_dir is None:
        data_dir = os.environ["SUBJECTS_DIR"]
    if subjects is None:
        subjects = []

    # Subject source node
    subjectsource = Node(IdentityInterface(fields=["subject_id"]),
                         iterables=("subject_id", subjects),
                         name="subjectsource")

    # Grab recon-all outputs
    templates = dict(aseg="{subject_id}/mri/aparc+aseg.mgz",
                     head="{subject_id}/mri/brain.mgz")
    datasource = Node(SelectFiles(templates,
                                  base_directory=data_dir),
                      "datasource")

    # Convert images to nifti storage and float representation
    cvtaseg = Node(fs.MRIConvert(out_type="niigz"), "convertaseg")

    cvtbrain = Node(fs.MRIConvert(out_type="niigz",
                                  out_datatype="float"),
                    "convertbrain")

    # Turn the aparc+aseg into a brainmask
    makemask = Node(fs.Binarize(dilate=4, erode=3, min=0.5), "makemask")

    # Extract the brain from the orig.mgz using the mask
    skullstrip = Node(fsl.ApplyMask(), "skullstrip")

    # Normalize using ANTS
    antswarp = Node(ANTSIntroduction(), "antswarp")

    # Generate a png summarizing the registration
    warpreport = Node(WarpReport(), "warpreport")

    # Save relevant files to the data directory
    datasink = Node(DataSink(base_directory=data_dir,
                             parameterization=False),
                    name="datasink")

    # Define and connect the workflow
    # -------------------------------

    normalize = Workflow(name=name)

    normalize.connect([
        (subjectsource, datasource,
            [("subject_id", "subject_id")]),
        (datasource, cvtaseg,
            [("aseg", "in_file")]),
        (datasource, cvtbrain,
            [("head", "in_file")]),
        (cvtaseg, makemask,
            [("out_file", "in_file")]),
        (cvtbrain, skullstrip,
            [("out_file", "in_file")]),
        (makemask, skullstrip,
            [("binary_file", "mask_file")]),
        (skullstrip, antswarp,
            [("out_file", "in_file")]),
        (antswarp, warpreport,
            [("brain_file", "in_file")]),
        (subjectsource, datasink,
            [("subject_id", "container")]),
        (antswarp, datasink,
            [("warp_file", "normalization.@warpfield"),
             ("inv_warp_file", "normalization.@inverse_warpfield"),
             ("affine_file", "normalization.@affine"),
             ("brain_file", "normalization.@brain")]),
        (warpreport, datasink,
            [("out_file", "normalization.@report")]),
    ])

    return normalize

class ANTSIntroductionInputSpec(BaseInterfaceInputSpec):

    in_file = File(exists=True)

class ANTSIntroductionOutputSpec(TraitedSpec):

    brain_file = File(exitsts=True)
    warp_file = File(exitsts=True)
    inv_warp_file = File(exitsts=True)
    affine_file = File(exitsts=True)

class ANTSIntroduction(BaseInterface):

    input_spec = ANTSIntroductionInputSpec
    output_spec = ANTSIntroductionOutputSpec

    def _run_interface(self, runtime):

        target_brain = fsl.Info.standard_image("avg152T1_brain.nii.gz")
        cmdline = " ".join(["antsIntroduction.sh",
                            "-d", "3",
                            "-r", target_brain,
                            "-i", self.inputs.in_file,
                            "-o", "ants_"])

        runtime.environ["ITK_NUM_THREADS"] = "1"
        runtime = submit_cmdline(runtime, cmdline)

        os.rename("ants_Affine.txt", "affine.txt")
        os.rename("ants_Warp.nii.gz", "warpfield.nii.gz")
        os.rename("ants_deformed.nii.gz", "brain_warp.nii.gz")
        os.rename("ants_InverseWarp.nii.gz", "inverse_warpfield.nii.gz")

        return runtime

    def _list_outputs(self):

        outputs = self._outputs().get()
        outputs["affine_file"] = op.abspath("affine.txt")
        outputs["warp_file"] = op.abspath("warpfield.nii.gz")
        outputs["brain_file"] = op.abspath("brain_warp.nii.gz")
        outputs["inv_warp_file"] = op.abspath("inverse_warpfield.nii.gz")

        return outputs

class WarpReportInput(BaseInterfaceInputSpec):

    in_file = File(exists=True)

class WarpReportOutput(TraitedSpec):

    out_file = File(exitsts=True)

class WarpReport(BaseInterface):

    input_spec = WarpReportInput
    output_spec = WarpReportOutput

    def _run_interface(self, runtime):

        target_brain = fsl.Info.standard_image("avg152T1_brain.nii.gz")
        m = Mosaic(self.inputs.in_file, target_brain)
        m.plot_contours("Reds", 2)
        m.savefig("warp_report.png")
        m.close()

        return runtime

    def _list_outputs(self):

        outputs = self._outputs().get()
        outputs["out_file"] = op.realpath("warp_report.png")
        return outputs
<|endoftext|>


<commit_before>"""Objects for MAASs."""

__all__ = [
    "MAAS",
]

import enum
import re
from typing import (
    Optional,
    Sequence,
)

from . import (
    Object,
    ObjectType,
)
from ..bones import CallError
from ..utils.typecheck import typed

class DescriptiveEnum(enum.Enum):
    """An enum where each member is a `(parameter, description)` tuple."""

    @property
    def parameter(self):
        return self.value[0]

    @property
    def description(self):
        return self.value[1]

    @classmethod
    def lookup(cls, parameter):
        for member in cls:
            if member.parameter == parameter:
                return member
        else:
            raise KeyError(
                "%s value %r not recognised."
                % (cls.__name__, parameter))

class MAASType(ObjectType):
    """Metaclass for `MAAS`."""

    @typed
    async def get_name(cls) -> str:
        """The name of the MAAS instance."""
        return await cls.get_config("maas_name")

    @typed
    async def set_name(cls, name: str):
        """See `get_name`."""
        return await cls.set_config("maas_name", name)

    @typed
    async def get_main_archive(cls) -> str:
        """Main archive URL.

        Archive used by nodes to retrieve packages for Intel architectures,
        e.g. http://archive.ubuntu.com/ubuntu.
        """
        return await cls.get_config("main_archive")

    @typed
    async def set_main_archive(cls, url: str):
        """See `get_main_archive`."""
        await cls.set_config("main_archive", url)

    @typed
    async def get_ports_archive(cls) -> str:
        """Ports archive.

        Archive used by nodes to retrieve packages for non-Intel
        architectures. E.g. http://ports.ubuntu.com/ubuntu-ports.
        """
        return await cls.get_config("ports_archive")

    @typed
    async def set_ports_archive(cls, series: str):
        """See `get_ports_archive`."""
        await cls.set_config("ports_archive", series)

    @typed
    async def get_default_os(cls) -> str:
        """Default OS used for deployment."""
        return await cls.get_config("default_osystem")

    @typed
    async def set_default_os(cls, series: str):
        """See `get_default_os`."""
        await cls.set_config("default_osystem", series)

    @typed
    async def get_default_distro_series(cls) -> str:
        """Default OS release used for deployment."""
        return await cls.get_config("default_distro_series")

    @typed
    async def set_default_distro_series(cls, series: str):
        """See `get_default_distro_series`."""
        await cls.set_config("default_distro_series", series)

    @typed
    async def get_commissioning_distro_series(cls) -> str:
        """Default Ubuntu release used for commissioning."""
        return await cls.get_config("commissioning_distro_series")

    @typed
    async def set_commissioning_distro_series(cls, series: str):
        """See `get_commissioning_distro_series`."""
        await cls.set_config("commissioning_distro_series", series)

    @typed
    async def get_http_proxy(cls) -> Optional[str]:
        """Proxy for APT and HTTP/HTTPS.

        This will be passed onto provisioned nodes to use as a proxy for APT
        traffic. MAAS also uses the proxy for downloading boot images. If no
        URL is provided, the built-in MAAS proxy will be used.
        """
        data = cls.get_config("http_proxy")
        return None if data is None or data == "" else data

    @typed
    async def set_http_proxy(cls, url: Optional[str]):
        """See `get_http_proxy`."""
        await cls.set_config("http_proxy", "" if url is None else url)

    @typed
    async def get_enable_http_proxy(cls) -> bool:
        """Enable the use of an APT and HTTP/HTTPS proxy.

        Provision nodes to use the built-in HTTP proxy (or user specified
        proxy) for APT. MAAS also uses the proxy for downloading boot images.
        """
        return await cls.get_config("enable_http_proxy")

    @typed
    async def set_enable_http_proxy(cls, enabled: bool):
        """See `get_enable_http_proxy`."""
        await cls.set_config("enable_http_proxy", "1" if enabled else "0")

    @typed
    async def get_curtin_verbose(cls) -> bool:
        """Should `curtin` log with high verbosity?

        Run the fast-path installer with higher verbosity. This provides more
        detail in the installation logs.
        """
        return await cls.get_config("curtin_verbose")

    @typed
    async def set_curtin_verbose(cls, verbose: bool):
        """See `get_curtin_verbose`."""
        await cls.set_config("curtin_verbose", "1" if verbose else "0")

    @typed
    async def get_kernel_options(cls) -> Optional[str]:
        """Kernel options.

        Boot parameters to pass to the kernel by default.
        """
        data = cls.get_config("kernel_opts")
        return None if data is None or data == "" else data

    @typed
    async def set_kernel_options(cls, options: Optional[str]):
        """See `get_kernel_options`."""
        await cls.set_config("kernel_opts", "" if options is None else options)

    @typed
    async def get_upstream_dns(cls) -> list:
        """Upstream DNS server addresses.

        Upstream DNS servers used to resolve domains not managed by this MAAS
        (space-separated IP addresses). Only used when MAAS is running its own
        DNS server. This value is used as the value of 'forwarders' in the DNS
        server config.
        """
        data = cls.get_config("upstream_dns")
        return [] if data is None else re.split("[,\s]+", data)

    @typed
    async def set_upstream_dns(cls, addresses: Optional[Sequence[str]]):
        """See `get_upstream_dns`."""
        await cls.set_config("upstream_dns", (
            "" if addresses is None else",".join(addresses)))

    class DNSSEC(DescriptiveEnum):
        """DNSSEC validation settings.

        See `get_dnssec_validation` and `set_dnssec_validation`.
        """

        AUTO = "auto", "Automatic (use default root key)"
        YES = "yes", "Yes (manually configured root key)"
        NO = "no", "No (disable DNSSEC)"

    @typed
    async def get_dnssec_validation(cls) -> DNSSEC:
        """Enable DNSSEC validation of upstream zones.

        Only used when MAAS is running its own DNS server. This value is used
        as the value of 'dnssec_validation' in the DNS server config.
        """
        data = cls.get_config("dnssec_validation")
        return cls.DNSSEC.lookup(data)

    @typed
    async def set_dnssec_validation(cls, validation: DNSSEC):
        """See `get_dnssec_validation`."""
        await cls.set_config("dnssec_validation", validation.parameter)

    @typed
    async def get_default_dns_ttl(cls) -> int:
        """Default Time-To-Live for DNS records.

        If no TTL value is specified at a more specific point this is how long
        DNS responses are valid, in seconds.
        """
        return int(cls.get_config("default_dns_ttl"))

    @typed
    async def set_default_dns_ttl(cls, ttl: int):
        """See `get_default_dns_ttl`."""
        await cls.set_config("default_dns_ttl", str(ttl))

    @typed
    async def get_enable_disk_erasing_on_release(cls) -> bool:
        """Should nodes' disks be erased prior to releasing."""
        return await cls.get_config("enable_disk_erasing_on_release")

    @typed
    async def set_enable_disk_erasing_on_release(cls, erase: bool):
        """Should nodes' disks be erased prior to releasing."""
        await cls.set_config(
            "enable_disk_erasing_on_release", "1" if erase else "0")

    @typed
    async def get_windows_kms_host(cls) -> Optional[str]:
        """Windows KMS activation host.

        FQDN or IP address of the host that provides the KMS Windows
        activation service. (Only needed for Windows deployments using KMS
        activation.)
        """
        data = cls.get_config("windows_kms_host")
        return None if data is None or data == "" else data

    @typed
    async def set_windows_kms_host(cls, host: Optional[str]):
        """See `get_windows_kms_host`."""
        await cls.set_config("windows_kms_host", "" if host is None else host)

    @typed
    async def get_boot_images_auto_import(cls) -> bool:
        """Automatically import/refresh the boot images every 60 minutes."""
        return await cls.get_config("boot_images_auto_import")

    @typed
    async def set_boot_images_auto_import(cls, auto: bool):
        """See `get_boot_images_auto_import`."""
        await cls.set_config("boot_images_auto_import", "1" if auto else "0")

    @typed
    async def get_ntp_server(cls) -> str:
        """Address of NTP server.

        NTP server address passed to nodes via a DHCP response. e.g.
        ntp.ubuntu.com.
        """
        return await cls.get_config("ntp_server")

    @typed
    async def set_ntp_server(cls, server: str):
        """See `get_ntp_server`."""
        await cls.set_config("ntp_server", server)

    class StorageLayout(DescriptiveEnum):

        FLAT = "flat", "Flat layout"
        LVM = "lvm", "LVM layout"
        BCACHE = "bcache", "Bcache layout"

    @typed
    async def get_default_storage_layout(cls) -> StorageLayout:
        """Default storage layout.

        Storage layout that is applied to a node when it is deployed.
        """
        data = cls.get_config("default_storage_layout")
        return cls.StorageLayout.lookup(data)

    @typed
    async def set_default_storage_layout(cls, series: StorageLayout):
        """See `get_default_storage_layout`."""
        await cls.set_config("default_storage_layout", series.parameter)

    @typed
    async def get_default_min_hwe_kernel(cls) -> Optional[str]:
        """Default minimum kernel version.

        The minimum kernel version used on new and commissioned nodes.
        """
        data = cls.get_config("default_min_hwe_kernel")
        return None if data is None or data == "" else data

    @typed
    async def set_default_min_hwe_kernel(cls, version: Optional[str]):
        """See `get_default_min_hwe_kernel`."""
        await cls.set_config(
            "default_min_hwe_kernel", "" if version is None else version)

    @typed
    async def get_enable_third_party_drivers(cls) -> bool:
        """Enable the installation of proprietary drivers, e.g. HPVSA."""
        return await cls.get_config("enable_third_party_drivers")

    @typed
    async def set_enable_third_party_drivers(cls, enabled: bool):
        """See `get_enable_third_party_drivers`."""
        await cls.set_config(
            "enable_third_party_drivers", "1" if enabled else "0")

    async def get_config(cls, name: str):
        """Get a configuration value from MAAS.

        Consult your MAAS server for recognised settings. Alternatively, use
        the pre-canned functions also defined on this object.
        """
        return await cls._handler.get_config(name=[name])

    async def set_config(cls, name: str, value):
        """Set a configuration value in MAAS.

        Consult your MAAS server for recognised settings. Alternatively, use
        the pre-canned functions also defined on this object.
        """
        return await cls._handler.set_config(name=[name], value=[value])

    async def _roundtrip(cls):
        """Testing helper: gets each value and sets it again."""
        getters = {
            name[4:]: getattr(cls, name) for name in dir(cls)
            if name.startswith("get_") and name != "get_config"
        }
        setters = {
            name[4:]: getattr(cls, name) for name in dir(cls)
            if name.startswith("set_") and name != "set_config"
        }

        for name, getter in getters.items():
            print(">>>", name)
            value = await getter()
            print(" ->", repr(value))
            print("   ", type(value))
            setter = setters[name]
            try:
                await setter(value)
            except CallError as error:
                print(error)
                print(error.content.decode("utf-8", "replace"))
            else:
                value2 = getter()
                if value2 != value:
                    print(
                        "!!! Round-trip failed:", repr(value),
                        "-->", repr(value2))

        getters_without_setters = set(getters).difference(setters)
        if getters_without_setters:
            print(
                "!!! Getters without setters:",
                " ".join(getters_without_setters))

        setters_without_getters = set(setters).difference(getters)
        if setters_without_getters:
            print(
                "!!! Setters without getters:",
                " ".join(setters_without_getters))

class MAAS(Object, metaclass=MAASType):
    """The current MAAS."""
<commit_msg>Get MAAS config setting and getting working.<commit_after>"""Objects for MAASs."""

__all__ = [
    "MAAS",
]

import enum
import re
from typing import (
    Optional,
    Sequence,
)

from . import (
    Object,
    ObjectType,
)
from ..bones import CallError
from ..utils.typecheck import typed

def _django_boolean(boolean):
    """Render a string suitable for use in a Django form.

    Django's `BooleanField` understands "true", "false", "1", "0", "t", and
    "f" as valid form values, but the widget `CheckboxInput` that is used by
    `BooleanField` by default gets a shot at converting the form input first.
    It inexplicably has different rules. See in `value_from_datadict`:

      values = {'true': True, 'false': False}
      if isinstance(value, six.string_types):
          value = values.get(value.lower(), value)
      return bool(value)

    Sigh.
    """
    return "true" if boolean else "false"

class DescriptiveEnum(enum.Enum):
    """An enum where each member is a `(parameter, description)` tuple."""

    @property
    def parameter(self):
        return self.value[0]

    @property
    def description(self):
        return self.value[1]

    @classmethod
    def lookup(cls, parameter):
        for member in cls:
            if member.parameter == parameter:
                return member
        else:
            raise KeyError(
                "%s value %r not recognised."
                % (cls.__name__, parameter))

class MAASType(ObjectType):
    """Metaclass for `MAAS`."""

    @typed
    async def get_name(cls) -> str:
        """The name of the MAAS instance."""
        return await cls.get_config("maas_name")

    @typed
    async def set_name(cls, name: str):
        """See `get_name`."""
        return await cls.set_config("maas_name", name)

    @typed
    async def get_main_archive(cls) -> str:
        """Main archive URL.

        Archive used by nodes to retrieve packages for Intel architectures,
        e.g. http://archive.ubuntu.com/ubuntu.
        """
        return await cls.get_config("main_archive")

    @typed
    async def set_main_archive(cls, url: str):
        """See `get_main_archive`."""
        await cls.set_config("main_archive", url)

    @typed
    async def get_ports_archive(cls) -> str:
        """Ports archive.

        Archive used by nodes to retrieve packages for non-Intel
        architectures. E.g. http://ports.ubuntu.com/ubuntu-ports.
        """
        return await cls.get_config("ports_archive")

    @typed
    async def set_ports_archive(cls, series: str):
        """See `get_ports_archive`."""
        await cls.set_config("ports_archive", series)

    @typed
    async def get_default_os(cls) -> str:
        """Default OS used for deployment."""
        return await cls.get_config("default_osystem")

    @typed
    async def set_default_os(cls, series: str):
        """See `get_default_os`."""
        await cls.set_config("default_osystem", series)

    @typed
    async def get_default_distro_series(cls) -> str:
        """Default OS release used for deployment."""
        return await cls.get_config("default_distro_series")

    @typed
    async def set_default_distro_series(cls, series: str):
        """See `get_default_distro_series`."""
        await cls.set_config("default_distro_series", series)

    @typed
    async def get_commissioning_distro_series(cls) -> str:
        """Default Ubuntu release used for commissioning."""
        return await cls.get_config("commissioning_distro_series")

    @typed
    async def set_commissioning_distro_series(cls, series: str):
        """See `get_commissioning_distro_series`."""
        await cls.set_config("commissioning_distro_series", series)

    @typed
    async def get_http_proxy(cls) -> Optional[str]:
        """Proxy for APT and HTTP/HTTPS.

        This will be passed onto provisioned nodes to use as a proxy for APT
        traffic. MAAS also uses the proxy for downloading boot images. If no
        URL is provided, the built-in MAAS proxy will be used.
        """
        data = await cls.get_config("http_proxy")
        return None if data is None or data == "" else data

    @typed
    async def set_http_proxy(cls, url: Optional[str]):
        """See `get_http_proxy`."""
        await cls.set_config("http_proxy", "" if url is None else url)

    @typed
    async def get_enable_http_proxy(cls) -> bool:
        """Enable the use of an APT and HTTP/HTTPS proxy.

        Provision nodes to use the built-in HTTP proxy (or user specified
        proxy) for APT. MAAS also uses the proxy for downloading boot images.
        """
        return await cls.get_config("enable_http_proxy")

    @typed
    async def set_enable_http_proxy(cls, enabled: bool):
        """See `get_enable_http_proxy`."""
        await cls.set_config("enable_http_proxy", _django_boolean(enabled))

    @typed
    async def get_curtin_verbose(cls) -> bool:
        """Should `curtin` log with high verbosity?

        Run the fast-path installer with higher verbosity. This provides more
        detail in the installation logs.
        """
        return await cls.get_config("curtin_verbose")

    @typed
    async def set_curtin_verbose(cls, verbose: bool):
        """See `get_curtin_verbose`."""
        await cls.set_config("curtin_verbose", _django_boolean(verbose))

    @typed
    async def get_kernel_options(cls) -> Optional[str]:
        """Kernel options.

        Boot parameters to pass to the kernel by default.
        """
        data = await cls.get_config("kernel_opts")
        return None if data is None or data == "" else data

    @typed
    async def set_kernel_options(cls, options: Optional[str]):
        """See `get_kernel_options`."""
        await cls.set_config("kernel_opts", "" if options is None else options)

    @typed
    async def get_upstream_dns(cls) -> list:
        """Upstream DNS server addresses.

        Upstream DNS servers used to resolve domains not managed by this MAAS
        (space-separated IP addresses). Only used when MAAS is running its own
        DNS server. This value is used as the value of 'forwarders' in the DNS
        server config.
        """
        data = await cls.get_config("upstream_dns")
        return [] if data is None else re.split("[,\s]+", data)

    @typed
    async def set_upstream_dns(cls, addresses: Optional[Sequence[str]]):
        """See `get_upstream_dns`."""
        await cls.set_config("upstream_dns", (
            "" if addresses is None else",".join(addresses)))

    class DNSSEC(DescriptiveEnum):
        """DNSSEC validation settings.

        See `get_dnssec_validation` and `set_dnssec_validation`.
        """

        AUTO = "auto", "Automatic (use default root key)"
        YES = "yes", "Yes (manually configured root key)"
        NO = "no", "No (disable DNSSEC)"

    @typed
    async def get_dnssec_validation(cls) -> DNSSEC:
        """Enable DNSSEC validation of upstream zones.

        Only used when MAAS is running its own DNS server. This value is used
        as the value of 'dnssec_validation' in the DNS server config.
        """
        data = await cls.get_config("dnssec_validation")
        return cls.DNSSEC.lookup(data)

    @typed
    async def set_dnssec_validation(cls, validation: DNSSEC):
        """See `get_dnssec_validation`."""
        await cls.set_config("dnssec_validation", validation.parameter)

    @typed
    async def get_default_dns_ttl(cls) -> int:
        """Default Time-To-Live for DNS records.

        If no TTL value is specified at a more specific point this is how long
        DNS responses are valid, in seconds.
        """
        return int(await cls.get_config("default_dns_ttl"))

    @typed
    async def set_default_dns_ttl(cls, ttl: int):
        """See `get_default_dns_ttl`."""
        await cls.set_config("default_dns_ttl", str(ttl))

    @typed
    async def get_enable_disk_erasing_on_release(cls) -> bool:
        """Should nodes' disks be erased prior to releasing."""
        return await cls.get_config("enable_disk_erasing_on_release")

    @typed
    async def set_enable_disk_erasing_on_release(cls, erase: bool):
        """Should nodes' disks be erased prior to releasing."""
        await cls.set_config(
            "enable_disk_erasing_on_release", _django_boolean(erase))

    @typed
    async def get_windows_kms_host(cls) -> Optional[str]:
        """Windows KMS activation host.

        FQDN or IP address of the host that provides the KMS Windows
        activation service. (Only needed for Windows deployments using KMS
        activation.)
        """
        data = await cls.get_config("windows_kms_host")
        return None if data is None or data == "" else data

    @typed
    async def set_windows_kms_host(cls, host: Optional[str]):
        """See `get_windows_kms_host`."""
        await cls.set_config("windows_kms_host", "" if host is None else host)

    @typed
    async def get_boot_images_auto_import(cls) -> bool:
        """Automatically import/refresh the boot images every 60 minutes."""
        return await cls.get_config("boot_images_auto_import")

    @typed
    async def set_boot_images_auto_import(cls, auto: bool):
        """See `get_boot_images_auto_import`."""
        await cls.set_config("boot_images_auto_import", _django_boolean(auto))

    @typed
    async def get_ntp_server(cls) -> str:
        """Address of NTP server.

        NTP server address passed to nodes via a DHCP response. e.g.
        ntp.ubuntu.com.
        """
        return await cls.get_config("ntp_server")

    @typed
    async def set_ntp_server(cls, server: str):
        """See `get_ntp_server`."""
        await cls.set_config("ntp_server", server)

    class StorageLayout(DescriptiveEnum):

        FLAT = "flat", "Flat layout"
        LVM = "lvm", "LVM layout"
        BCACHE = "bcache", "Bcache layout"

    @typed
    async def get_default_storage_layout(cls) -> StorageLayout:
        """Default storage layout.

        Storage layout that is applied to a node when it is deployed.
        """
        data = await cls.get_config("default_storage_layout")
        return cls.StorageLayout.lookup(data)

    @typed
    async def set_default_storage_layout(cls, series: StorageLayout):
        """See `get_default_storage_layout`."""
        await cls.set_config("default_storage_layout", series.parameter)

    @typed
    async def get_default_min_hwe_kernel(cls) -> Optional[str]:
        """Default minimum kernel version.

        The minimum kernel version used on new and commissioned nodes.
        """
        data = await cls.get_config("default_min_hwe_kernel")
        return None if data is None or data == "" else data

    @typed
    async def set_default_min_hwe_kernel(cls, version: Optional[str]):
        """See `get_default_min_hwe_kernel`."""
        await cls.set_config(
            "default_min_hwe_kernel", "" if version is None else version)

    @typed
    async def get_enable_third_party_drivers(cls) -> bool:
        """Enable the installation of proprietary drivers, e.g. HPVSA."""
        return await cls.get_config("enable_third_party_drivers")

    @typed
    async def set_enable_third_party_drivers(cls, enabled: bool):
        """See `get_enable_third_party_drivers`."""
        await cls.set_config(
            "enable_third_party_drivers", _django_boolean(enabled))

    async def get_config(cls, name: str):
        """Get a configuration value from MAAS.

        Consult your MAAS server for recognised settings. Alternatively, use
        the pre-canned functions also defined on this object.
        """
        return await cls._handler.get_config(name=[name])

    async def set_config(cls, name: str, value):
        """Set a configuration value in MAAS.

        Consult your MAAS server for recognised settings. Alternatively, use
        the pre-canned functions also defined on this object.
        """
        return await cls._handler.set_config(name=[name], value=[value])

    async def _roundtrip(cls):
        """Testing helper: gets each value and sets it again."""
        getters = {
            name[4:]: getattr(cls, name) for name in dir(cls)
            if name.startswith("get_") and name != "get_config"
        }
        setters = {
            name[4:]: getattr(cls, name) for name in dir(cls)
            if name.startswith("set_") and name != "set_config"
        }

        for name, getter in getters.items():
            print(">>>", name)
            value = await getter()
            print(" ->", repr(value))
            print("   ", type(value))
            setter = setters[name]
            try:
                await setter(value)
            except CallError as error:
                print(error)
                print(error.content.decode("utf-8", "replace"))
            else:
                value2 = await getter()
                if value2 != value:
                    print(
                        "!!! Round-trip failed:", repr(value),
                        "-->", repr(value2))

        getters_without_setters = set(getters).difference(setters)
        if getters_without_setters:
            print(
                "!!! Getters without setters:",
                " ".join(getters_without_setters))

        setters_without_getters = set(setters).difference(getters)
        if setters_without_getters:
            print(
                "!!! Setters without getters:",
                " ".join(setters_without_getters))

class MAAS(Object, metaclass=MAASType):
    """The current MAAS."""
<|endoftext|>