<commit_before>#     Copyright 2014 Netflix, Inc.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.

from security_monkey.views import AuthenticatedService
from security_monkey.views import ACCOUNT_FIELDS
from security_monkey.datastore import Account, AccountType
from security_monkey.datastore import User
from security_monkey.account_manager import get_account_by_id, delete_account_by_id
from security_monkey import db, rbac

from flask import request
from flask_restful import marshal, reqparse
import json

class AccountGetPutDelete(AuthenticatedService):
    decorators = [
        rbac.allow(["View"], ["GET"]),
        rbac.allow(["Admin"], ["PUT", "DELETE"])
    ]
    def __init__(self):
        self.reqparse = reqparse.RequestParser()
        super(AccountGetPutDelete, self).__init__()

    def get(self, account_id):
        """
            .. http:get:: /api/1/account/<int:id>

            Get a list of Accounts matching the given criteria

            **Example Request**:

            .. sourcecode:: http

                GET /api/1/account/1 HTTP/1.1
                Host: example.com
                Accept: application/json, text/javascript

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 200 OK
                Vary: Accept
                Content-Type: application/json

                {
                    third_party: false,
                    name: "example_name",
                    notes: null,
                    identifier: "111111111111",
                    active: true,
                    id: 1,
                    account_type: "AWS",
                    auth: {
                        authenticated: true,
                        user: "user@example.com"
                    }
                }

            :statuscode 200: no error
            :statuscode 401: Authentication failure. Please login.
        """

        result = get_account_by_id(account_id)

        account_marshaled = marshal(result.__dict__, ACCOUNT_FIELDS)
        account_marshaled = dict(
            account_marshaled.items() +
            {'account_type': result.account_type.name}.items()
        )

        custom_fields_marshaled = []
        for field in result.custom_fields:
            field_marshaled = {
                                  'name': field.name,
                                  'value': field.value,
                              }
            custom_fields_marshaled.append(field_marshaled)
        account_marshaled['custom_fields'] = custom_fields_marshaled

        account_marshaled['auth'] = self.auth_dict
        return account_marshaled, 200

    def put(self, account_id):
        """
            .. http:put:: /api/1/account/1

            Edit an existing account.

            **Example Request**:

            .. sourcecode:: http

                PUT /api/1/account/1 HTTP/1.1
                Host: example.com
                Accept: application/json

                {
                    'name': 'edited_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                }

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 200 OK
                Vary: Accept
                Content-Type: application/json

                {
                    'name': 'edited_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                    'account_type': 'AWS'
                }

            :statuscode 200: no error
            :statuscode 401: Authentication Error. Please Login.
        """

        args = json.loads(request.json)
        account_type = args['account_type']
        name = args['name']
        identifier = args['identifier']
        notes = args['notes']
        active = args['active']
        third_party = args['third_party']
        account_id = args['id']
        custom_fields = args['custom_fields']

        from security_monkey.account_manager import account_registry
        account_manager = account_registry.get(account_type)()
        account = account_manager.update(account_id, account_type, name, active,
                    third_party, notes, identifier, custom_fields=custom_fields)

        if not account:
            return {'status': 'error. Account ID not found.'}, 404

        from security_monkey.common.audit_issue_cleanup import clean_account_issues
        clean_account_issues(account)

        marshaled_account = marshal(account.__dict__, ACCOUNT_FIELDS)
        marshaled_account['auth'] = self.auth_dict

        return marshaled_account, 200

    def delete(self, account_id):
        """
            .. http:delete:: /api/1/account/1

            Delete an account.

            **Example Request**:

            .. sourcecode:: http

                DELETE /api/1/account/1 HTTP/1.1
                Host: example.com
                Accept: application/json

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 202 Accepted
                Vary: Accept
                Content-Type: application/json

                {
                    'status': 'deleted'
                }

            :statuscode 202: accepted
            :statuscode 401: Authentication Error. Please Login.
        """
        delete_account_by_id(account_id)
        return {'status': 'deleted'}, 202

class AccountPostList(AuthenticatedService):
    decorators = [
        rbac.allow(["View"], ["GET"]),
        rbac.allow(["Admin"], ["POST"])
    ]

    def __init__(self):
        super(AccountPostList, self).__init__()
        self.reqparse = reqparse.RequestParser()

    def post(self):
        """
            .. http:post:: /api/1/account/

            Create a new account.

            **Example Request**:

            .. sourcecode:: http

                POST /api/1/account/ HTTP/1.1
                Host: example.com
                Accept: application/json

                {
                    'name': 'new_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                    'account_type': 'AWS'
                }

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 201 Created
                Vary: Accept
                Content-Type: application/json

                {
                    'name': 'new_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                    'account_type': 'AWS'
                    ''
                }

            :statuscode 201: created
            :statuscode 401: Authentication Error. Please Login.
        """

        args = json.loads(request.json)
        account_type = args['account_type']
        name = args['name']
        identifier = args['identifier']
        notes = args['notes']
        active = args['active']
        third_party = args['third_party']
        custom_fields = args['custom_fields']

        from security_monkey.account_manager import account_registry
        account_manager = account_registry.get(account_type)()
        account = account_manager.create(account_type, name, active, third_party,
                    notes, identifier, custom_fields=custom_fields)

        marshaled_account = marshal(account.__dict__, ACCOUNT_FIELDS)
        marshaled_account['auth'] = self.auth_dict
        return marshaled_account, 200

    def get(self):
        """
            .. http:get:: /api/1/accounts

            Get a list of Accounts matching the given criteria

            **Example Request**:

            .. sourcecode:: http

                GET /api/1/accounts HTTP/1.1
                Host: example.com
                Accept: application/json, text/javascript

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 200 OK
                Vary: Accept
                Content-Type: application/json

                {
                    count: 1,
                    items: [
                        {
                            third_party: false,
                            name: "example_name",
                            notes: null,
                            role_name: null,
                            identifier: "111111111111",
                            active: true,
                            id: 1,
                            s3_name: "example_name"
                        },
                    ],
                    total: 1,
                    page: 1,
                    auth: {
                        authenticated: true,
                        user: "user@example.com"
                    }
                }

            :statuscode 200: no error
            :statuscode 401: Authentication failure. Please login.
        """

        self.reqparse.add_argument('count', type=int, default=30, location='args')
        self.reqparse.add_argument('page', type=int, default=1, location='args')
        self.reqparse.add_argument('order_by', type=str, default=None, location='args')
        self.reqparse.add_argument('order_dir', type=str, default='desc', location='args')

        args = self.reqparse.parse_args()
        page = args.pop('page', None)
        count = args.pop('count', None)
        order_by = args.pop('order_by', None)
        order_dir = args.pop('order_dir', None)
        for k, v in args.items():
            if not v:
                del args[k]

        query = Account.query

        if order_by and hasattr(Account, order_by):
            if order_dir.lower() == 'asc':
                if order_by == 'account_type':
                    query = query.join(Account.account_type).order_by(getattr(AccountType, 'name').asc())
                else:
                    query = query.order_by(getattr(Account, order_by).asc())
            else:
                if order_by == 'account_type':
                    query = query.join(Account.account_type).order_by(getattr(AccountType, 'name').desc())
                else:
                    query = query.order_by(getattr(Account, order_by).desc())
        else:
            query = query.order_by(Account.id)

        result = query.paginate(page, count, error_out=False)

        items = []
        for account in result.items:
            account_marshaled = marshal(account.__dict__, ACCOUNT_FIELDS)
            account_marshaled = dict(
                account_marshaled.items() +
                {'account_type': account.account_type.name}.items()
            )

            items.append(account_marshaled)

        marshaled_dict = {
            'total': result.total,
            'count': len(items),
            'page': result.page,
            'items': items,
            'auth': self.auth_dict
        }

        return marshaled_dict, 200
<commit_msg>Adding active and third_party flags to account view<commit_after>#     Copyright 2014 Netflix, Inc.
#
#     Licensed under the Apache License, Version 2.0 (the "License");
#     you may not use this file except in compliance with the License.
#     You may obtain a copy of the License at
#
#         http://www.apache.org/licenses/LICENSE-2.0
#
#     Unless required by applicable law or agreed to in writing, software
#     distributed under the License is distributed on an "AS IS" BASIS,
#     WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
#     See the License for the specific language governing permissions and
#     limitations under the License.

from security_monkey.views import AuthenticatedService
from security_monkey.views import ACCOUNT_FIELDS
from security_monkey.datastore import Account, AccountType
from security_monkey.datastore import User
from security_monkey.account_manager import get_account_by_id, delete_account_by_id
from security_monkey import db, rbac

from flask import request
from flask_restful import marshal, reqparse
import json

class AccountGetPutDelete(AuthenticatedService):
    decorators = [
        rbac.allow(["View"], ["GET"]),
        rbac.allow(["Admin"], ["PUT", "DELETE"])
    ]
    def __init__(self):
        self.reqparse = reqparse.RequestParser()
        super(AccountGetPutDelete, self).__init__()

    def get(self, account_id):
        """
            .. http:get:: /api/1/account/<int:id>

            Get a list of Accounts matching the given criteria

            **Example Request**:

            .. sourcecode:: http

                GET /api/1/account/1 HTTP/1.1
                Host: example.com
                Accept: application/json, text/javascript

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 200 OK
                Vary: Accept
                Content-Type: application/json

                {
                    third_party: false,
                    name: "example_name",
                    notes: null,
                    identifier: "111111111111",
                    active: true,
                    id: 1,
                    account_type: "AWS",
                    auth: {
                        authenticated: true,
                        user: "user@example.com"
                    }
                }

            :statuscode 200: no error
            :statuscode 401: Authentication failure. Please login.
        """

        result = get_account_by_id(account_id)

        account_marshaled = marshal(result.__dict__, ACCOUNT_FIELDS)
        account_marshaled = dict(
            account_marshaled.items() +
            {'account_type': result.account_type.name}.items()
        )

        custom_fields_marshaled = []
        for field in result.custom_fields:
            field_marshaled = {
                                  'name': field.name,
                                  'value': field.value,
                              }
            custom_fields_marshaled.append(field_marshaled)
        account_marshaled['custom_fields'] = custom_fields_marshaled

        account_marshaled['auth'] = self.auth_dict
        return account_marshaled, 200

    def put(self, account_id):
        """
            .. http:put:: /api/1/account/1

            Edit an existing account.

            **Example Request**:

            .. sourcecode:: http

                PUT /api/1/account/1 HTTP/1.1
                Host: example.com
                Accept: application/json

                {
                    'name': 'edited_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                }

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 200 OK
                Vary: Accept
                Content-Type: application/json

                {
                    'name': 'edited_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                    'account_type': 'AWS'
                }

            :statuscode 200: no error
            :statuscode 401: Authentication Error. Please Login.
        """

        args = json.loads(request.json)
        account_type = args['account_type']
        name = args['name']
        identifier = args['identifier']
        notes = args['notes']
        active = args['active']
        third_party = args['third_party']
        account_id = args['id']
        custom_fields = args['custom_fields']

        from security_monkey.account_manager import account_registry
        account_manager = account_registry.get(account_type)()
        account = account_manager.update(account_id, account_type, name, active,
                    third_party, notes, identifier, custom_fields=custom_fields)

        if not account:
            return {'status': 'error. Account ID not found.'}, 404

        from security_monkey.common.audit_issue_cleanup import clean_account_issues
        clean_account_issues(account)

        marshaled_account = marshal(account.__dict__, ACCOUNT_FIELDS)
        marshaled_account['auth'] = self.auth_dict

        return marshaled_account, 200

    def delete(self, account_id):
        """
            .. http:delete:: /api/1/account/1

            Delete an account.

            **Example Request**:

            .. sourcecode:: http

                DELETE /api/1/account/1 HTTP/1.1
                Host: example.com
                Accept: application/json

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 202 Accepted
                Vary: Accept
                Content-Type: application/json

                {
                    'status': 'deleted'
                }

            :statuscode 202: accepted
            :statuscode 401: Authentication Error. Please Login.
        """
        delete_account_by_id(account_id)
        return {'status': 'deleted'}, 202

class AccountPostList(AuthenticatedService):
    decorators = [
        rbac.allow(["View"], ["GET"]),
        rbac.allow(["Admin"], ["POST"])
    ]

    def __init__(self):
        super(AccountPostList, self).__init__()
        self.reqparse = reqparse.RequestParser()

    def post(self):
        """
            .. http:post:: /api/1/account/

            Create a new account.

            **Example Request**:

            .. sourcecode:: http

                POST /api/1/account/ HTTP/1.1
                Host: example.com
                Accept: application/json

                {
                    'name': 'new_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                    'account_type': 'AWS'
                }

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 201 Created
                Vary: Accept
                Content-Type: application/json

                {
                    'name': 'new_account'
                    'identifier': '0123456789',
                    'notes': 'this account is for ...',
                    'active': true,
                    'third_party': false
                    'account_type': 'AWS'
                    ''
                }

            :statuscode 201: created
            :statuscode 401: Authentication Error. Please Login.
        """

        args = json.loads(request.json)
        account_type = args['account_type']
        name = args['name']
        identifier = args['identifier']
        notes = args['notes']
        active = args['active']
        third_party = args['third_party']
        custom_fields = args['custom_fields']

        from security_monkey.account_manager import account_registry
        account_manager = account_registry.get(account_type)()
        account = account_manager.create(account_type, name, active, third_party,
                    notes, identifier, custom_fields=custom_fields)

        marshaled_account = marshal(account.__dict__, ACCOUNT_FIELDS)
        marshaled_account['auth'] = self.auth_dict
        return marshaled_account, 200

    def get(self):
        """
            .. http:get:: /api/1/accounts

            Get a list of Accounts matching the given criteria

            **Example Request**:

            .. sourcecode:: http

                GET /api/1/accounts HTTP/1.1
                Host: example.com
                Accept: application/json, text/javascript

            **Example Response**:

            .. sourcecode:: http

                HTTP/1.1 200 OK
                Vary: Accept
                Content-Type: application/json

                {
                    count: 1,
                    items: [
                        {
                            third_party: false,
                            name: "example_name",
                            notes: null,
                            role_name: null,
                            identifier: "111111111111",
                            active: true,
                            id: 1,
                            s3_name: "example_name"
                        },
                    ],
                    total: 1,
                    page: 1,
                    auth: {
                        authenticated: true,
                        user: "user@example.com"
                    }
                }

            :statuscode 200: no error
            :statuscode 401: Authentication failure. Please login.
        """

        self.reqparse.add_argument('count', type=int, default=30, location='args')
        self.reqparse.add_argument('page', type=int, default=1, location='args')
        self.reqparse.add_argument('order_by', type=str, default=None, location='args')
        self.reqparse.add_argument('order_dir', type=str, default='desc', location='args')
        self.reqparse.add_argument('active', type=str, default=None, location='args')
        self.reqparse.add_argument('third_party', type=str, default=None, location='args')

        args = self.reqparse.parse_args()
        page = args.pop('page', None)
        count = args.pop('count', None)
        order_by = args.pop('order_by', None)
        order_dir = args.pop('order_dir', None)
        for k, v in args.items():
            if not v:
                del args[k]

        query = Account.query
        if 'active' in args:
            active = args['active'].lower() == "true"
            query = query.filter(Account.active == active)
        if 'third_party' in args:
            third_party = args['third_party'].lower() == "true"
            query = query.filter(Account.third_party == third_party)

        if order_by and hasattr(Account, order_by):
            if order_dir.lower() == 'asc':
                if order_by == 'account_type':
                    query = query.join(Account.account_type).order_by(getattr(AccountType, 'name').asc())
                else:
                    query = query.order_by(getattr(Account, order_by).asc())
            else:
                if order_by == 'account_type':
                    query = query.join(Account.account_type).order_by(getattr(AccountType, 'name').desc())
                else:
                    query = query.order_by(getattr(Account, order_by).desc())
        else:
            query = query.order_by(Account.id)

        result = query.paginate(page, count, error_out=False)

        items = []
        for account in result.items:
            account_marshaled = marshal(account.__dict__, ACCOUNT_FIELDS)
            account_marshaled = dict(
                account_marshaled.items() +
                {'account_type': account.account_type.name}.items()
            )

            items.append(account_marshaled)

        marshaled_dict = {
            'total': result.total,
            'count': len(items),
            'page': result.page,
            'items': items,
            'auth': self.auth_dict
        }

        return marshaled_dict, 200
<|endoftext|>


<commit_before>import numpy as np
import pandas as pd
import lmfit

from sklearn.base import BaseEstimator, RegressorMixin

from epistasis.stats import pearson
from ..mapping import EpistasisMap, mutations_to_sites
from .base import BaseModel
from epistasis.model_matrix_ext import get_model_matrix
from .utils import X_fitter, X_predictor

class State(EpistasisMap):
    """A state in an EpistasisEnsembleModel.

    This represents/models an unknown excited state that is perturbed by
    mutations in the genotype-phenotype map and contributes to the overall
    phenotype.

    Parameters
    ----------
    name : string
        name of the state. Should have the form 'state_X' where X is a letter.

    sites : list of lists
        each element in list is a list of genotype site indexes. These sites
        are additive/epistatic coefficients contributing to this state.

    Attributes
    ----------
    See EpistasisMap for more details.

    keys : list
        list of lmfit.Parameter keys.
    """
    def __init__(self, name, sites, *args, **kwargs):
        # Call super init.
        super(State, self).__init__(sites=sites, *args, **kwargs)

        # Set name.
        self.name = name

        # Construct parameters object
        self.parameters = lmfit.Parameters()

        # fill parameters .
        for key in self.keys:
            self.parameters.add(key, value=0)

    @property
    def keys(self):
        """Coefficient Parameter keys."""
        keys = []
        for sites in self.sites:
            key = "".join([str(ch) for ch in sites])
            name = "{}_{}".format(self.name, key)
            keys.append(name)
        return keys

class EpistasisEnsembleRegression(BaseEstimator, BaseModel):
    """Ensemble epistasis model. It models variation in a genotype-phenotype map
    as a statistical ensemble of nstates contributing to each genotype's
    phenotype.

    Attributes
    ----------
    See BaseModel for more details.

    parameters : lmfit.Parameters
        Parameters resulting from fit.
    """

    _ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']

    def __init__(self, order=1, nstates=2):
        self.nstates = nstates
        self.model_type = 'local'
        self.order = order
        self.states = {}
        self.Xbuilt = {}
        self.parameters = lmfit.Parameters()

    def add_gpm(self, gpm):
        """Add genotype-phenotype map to model object."""
        super(EpistasisEnsembleRegression, self).add_gpm(gpm)

        # Add states to model.
        for i in range(self.nstates):
            # State name
            name = "state_{}".format(self._ALPHABET[i])

            # Add state.
            self.add_state(name)

        return self

    def add_state(self, name):
        """ Add a state to the model."""
        sites = self.Xcolumns

        # Create state.
        state = State(name, sites, model_type=self.model_type)

        # Store state.
        self.states[name] = state

        # Set as attribute.
        setattr(self, name, state)

        return self

    @property
    def num_of_params(self):
        """Return number of parameters in model."""
        n = 0
        n += len(self.parameters)
        return n

    @property
    def parameters(self):
        """All parameters in the value."""
        parameters = lmfit.Parameters()

        # Get parameter data.
        for state in self.states.values():
            parameters.add_many(*state.parameters.values())

        return parameters

    @parameters.setter
    def parameters(self, parameters):
        """Set parameters for all states."""

        # Add state parameters.
        for state in self.states.values():
            keys = state.keys
            parameters_ = lmfit.Parameters()
            values = []
            for key in state.keys:
                p = parameters[key]
                parameters_.add_many(p)
                values.append(p.value)

            parameters_.add_many(*[parameters[key] for key in state.keys])
            state.parameters = parameters_
            state.values = values

    def functional_form(self, thetas, X=None):
        """Ensemble function calculating phenotypes using a Boltzmann weighted
        ensemble of states, where each state is constructed from linear
        epistasis models.

        Parameters
        ----------
        thetas : array
            array of parameters values ordered alphabetically.

        X : 2d-array (default None)
            X matrix to use for linear portion of models. If None, uses the
            matrix stored under the 'fit' key.
        """
        length = self.states['state_A'].n
        nstates = len(self.states)

        if X is None:
            X = self.Xbuilt['fit']

        # Calculate a partition function
        Z = 0
        for state_i in range(nstates):
            # Get parameter indexes
            idx_start = state_i * length
            idx_stop = state_i * length + length
            dDG = thetas[idx_start:idx_stop]

            # Get state.
            state = self.states['state_{}'.format(self._ALPHABET[state_i])]

            # Additive model.
            additive = X @ dDG

            # add to ensemble
            Z += np.exp(-additive)

        # Ensemble model.
        y = np.log(Z)
        return y

    @X_fitter
    def fit(self, X='obs', y='obs', **kwargs):
        """Fit ensemble model to data.

        Parameters
        ----------
        X : 2-d array
            independent data; samples.
        y : array
            dependent data; observations.
        """
        # X matrix.
        self.add_X(X=X, key='fit')

        # Storing failed residuals
        last_residual_set = None

        # Residual function to minimize.
        def residual(params, func, y=None):
            # Fit model
            parvals = list(params.values())
            ymodel = func(parvals)

            # Store items in case of error.
            nonlocal last_residual_set
            last_residual_set = (params, ymodel)

            return y - ymodel

        y = self.gpm.phenotypes

        # Minimize the above residual function.
        self.results = lmfit.minimize(
            residual, self.parameters,
            args=[self.functional_form],
            kws={'y': y})

        # Set parameters fitted by model.
        self.parameters = self.results.params

        return self

    def fit_transform(self, X='obs', y='obs', **kwargs):
        """Same as calling fit in ensemble model.
        """
        self.fit(X=X, y=y, **kwargs)

    @X_predictor
    def predict(self, X='obs'):
        """Predict phenotypes using fitted model.
        """
        return self.functional_form(list(self.parameters.values()), X=X)

    def predict_transform(self, X='obs', y='obs'):
        """Same as calling predict."""
        return self.predict(X=X)

    @X_fitter
    def score(self, X='obs', y='obs', **kwargs):
        """Calculate the pearson coefficient between the models predictions and
        a given y array.
        """
        return pearson(self.gpm.phenotypes, self.predict(X=X))**2
<commit_msg>update ensemble module<commit_after>import numpy as np
import pandas as pd
import lmfit

from sklearn.base import BaseEstimator, RegressorMixin

from epistasis.stats import pearson
from ..mapping import EpistasisMap, mutations_to_sites
from .base import BaseModel
from epistasis.model_matrix_ext import get_model_matrix
from .utils import arghandler

class State(EpistasisMap):
    """A state in an EpistasisEnsembleModel.

    This represents/models an unknown excited state that is perturbed by
    mutations in the genotype-phenotype map and contributes to the overall
    phenotype.

    Parameters
    ----------
    name : string
        name of the state. Should have the form 'state_X' where X is a letter.

    sites : list of lists
        each element in list is a list of genotype site indexes. These sites
        are additive/epistatic coefficients contributing to this state.

    Attributes
    ----------
    See EpistasisMap for more details.

    keys : list
        list of lmfit.Parameter keys.
    """
    def __init__(self, name, sites, *args, **kwargs):
        # Call super init.
        super(State, self).__init__(sites=sites, *args, **kwargs)

        # Set name.
        self.name = name

        # Construct parameters object
        self.parameters = lmfit.Parameters()

        # fill parameters .
        for key in self.keys:
            self.parameters.add(key, value=0)

    @property
    def keys(self):
        """Coefficient Parameter keys."""
        keys = []
        for sites in self.sites:
            key = "".join([str(ch) for ch in sites])
            name = "{}_{}".format(self.name, key)
            keys.append(name)
        return keys

class EpistasisEnsembleRegression(BaseModel):
    """Ensemble epistasis model. It models variation in a genotype-phenotype map
    as a statistical ensemble of nstates contributing to each genotype's
    phenotype.

    Attributes
    ----------
    See BaseModel for more details.

    parameters : lmfit.Parameters
        Parameters resulting from fit.
    """

    _ALPHABET = ['A', 'B', 'C', 'D', 'E', 'F', 'G', 'H', 'I', 'J']

    def __init__(self, order=1, nstates=2):
        self.nstates = nstates
        self.model_type = 'local'
        self.order = order
        self.states = {}
        self.Xbuilt = {}
        self.parameters = lmfit.Parameters()

    def add_gpm(self, gpm):
        """Add genotype-phenotype map to model object."""
        super(EpistasisEnsembleRegression, self).add_gpm(gpm)

        # Add states to model.
        for i in range(self.nstates):
            # State name
            name = "state_{}".format(self._ALPHABET[i])

            # Add state.
            self.add_state(name)

        return self

    def add_state(self, name):
        """ Add a state to the model."""
        sites = self.Xcolumns

        # Create state.
        state = State(name, sites, model_type=self.model_type)

        # Store state.
        self.states[name] = state

        # Set as attribute.
        setattr(self, name, state)

        return self

    @property
    def num_of_params(self):
        """Return number of parameters in model."""
        n = 0
        n += len(self.parameters)
        return n

    @property
    def parameters(self):
        """All parameters in the value."""
        parameters = lmfit.Parameters()

        # Get parameter data.
        for state in self.states.values():
            parameters.add_many(*state.parameters.values())

        return parameters

    @parameters.setter
    def parameters(self, parameters):
        """Set parameters for all states."""

        # Add state parameters.
        for state in self.states.values():
            keys = state.keys
            parameters_ = lmfit.Parameters()
            values = []
            for key in state.keys:
                p = parameters[key]
                parameters_.add_many(p)
                values.append(p.value)

            parameters_.add_many(*[parameters[key] for key in state.keys])
            state.parameters = parameters_
            state.values = values

    def functional_form(self, thetas, X=None):
        """Ensemble function calculating phenotypes using a Boltzmann weighted
        ensemble of states, where each state is constructed from linear
        epistasis models.

        Parameters
        ----------
        thetas : array
            array of parameters values ordered alphabetically.

        X : 2d-array (default None)
            X matrix to use for linear portion of models. If None, uses the
            matrix stored under the 'fit' key.
        """
        length = self.states['state_A'].n
        nstates = len(self.states)

        if X is None:
            X = self.Xbuilt['fit']

        # Calculate a partition function
        Z = 0
        for state_i in range(nstates):
            # Get parameter indexes
            idx_start = state_i * length
            idx_stop = state_i * length + length
            dDG = thetas[idx_start:idx_stop]

            # Get state.
            state = self.states['state_{}'.format(self._ALPHABET[state_i])]

            # Additive model.
            additive = X @ dDG

            # add to ensemble
            Z += np.exp(-additive)

        # Ensemble model.
        y = np.log(Z)
        return y

    @arghandler
    def fit(self, X=None, y=None, **kwargs):
        """Fit ensemble model to data.

        Parameters
        ----------
        X : 2-d array
            independent data; samples.
        y : array
            dependent data; observations.
        """
        # X matrix.
        self.add_X(X=X, key='fit')

        # Storing failed residuals
        last_residual_set = None

        # Residual function to minimize.
        def residual(params, func, y=None):
            # Fit model
            parvals = list(params.values())
            ymodel = func(parvals)

            # Store items in case of error.
            nonlocal last_residual_set
            last_residual_set = (params, ymodel)

            return y - ymodel

        y = self.gpm.phenotypes

        # Minimize the above residual function.
        self.results = lmfit.minimize(
            residual, self.parameters,
            args=[self.functional_form],
            kws={'y': y})

        # Set parameters fitted by model.
        self.parameters = self.results.params

        return self

    def fit_transform(self, X=None, y=None, **kwargs):
        """Same as calling fit in ensemble model.
        """
        self.fit(X=X, y=y, **kwargs)

    @arghandler
    def predict(self, X=None):
        """Predict phenotypes using fitted model.
        """
        return self.functional_form(list(self.parameters.values()), X=X)

    def predict_transform(self, X=None, y=None):
        """Same as calling predict."""
        return self.predict(X=X)

    @arghandler
    def score(self, X=None, y=None, **kwargs):
        """Calculate the pearson coefficient between the models predictions and
        a given y array.
        """
        return pearson(self.gpm.phenotypes, self.predict(X=X))**2
<|endoftext|>


<commit_before>style = '''<meta id="OGTitle" property="og:title" content="SeleniumBase">
        <title>Test Report</title>
        <link rel="SHORTCUT ICON"
        href="http://cdn2.hubspot.net/hubfs/100006/images/results_logo4.ico" />
        <style type="text/css">
        html {
            background-color: #9988ad;
        }
        html, body {
            font-size: 100%;
            box-sizing: border-box;
        }
        body {
            background-image: none;
            background-origin: padding-box;
            background-color: #ecf2ff;
            padding: 30;
            margin: 10;
            font-family: "Proxima Nova","proxima-nova",
            "Helvetica Neue",Helvetica,Arial,sans-serif !important;
            text-rendering: optimizelegibility;
            -moz-osx-font-smoothing: grayscale;
            box-shadow: 0px 2px 5px 1px rgba(0, 0, 0, 0.24),
            1px 2px 12px 0px rgba(0, 0, 0, 0.18) !important;
        }
        table {
            border-collapse: collapse;
            border-spacing: 0;
            box-shadow: 0px 2px 5px 1px rgba(0, 0, 0, 0.25),
            1px 2px 12px 0px rgba(0, 0, 0, 0.19) !important;
            transition: all 0.15s ease-out 0s;
            transition-property: all;
            transition-duration: 0.1s;
            transition-timing-function: ease-out;
            transition-delay: 0s;
        }
        table:hover {
            box-shadow: 0px 2px 5px 1px rgba(0, 0, 0, 0.35),
            1px 2px 12px 0px rgba(0, 0, 0, 0.28) !important;
        }
        thead th, thead td {
            padding: 0.5rem 0.625rem 0.625rem;
            font-weight: bold;
            text-align: left;
        }
        thead {
            text-align: center;
            border: 1px solid #e1e1e1;
            background-color: #fefefe;
            width: 150%;
            color: #0C8CDB;
            background-color: #cbf2ff;
        }
        tbody tr:nth-child(even) {
            background-color: #f1f1f1;
        }
        tbody tr:nth-child(odd) {
            background-color: #ffffff;
        }
        tbody tr:nth-child(even):hover {
            background-color: #f8f8d2;
        }
        tbody tr:nth-child(odd):hover {
            background-color: #ffffe0;
        }
        tbody th, tbody td {
            padding: 0.5rem 0.625rem 0.625rem;
        }
        tbody {
            border: 1px solid #e1e1e1;
            background-color: #fefefe;
        }
        td {
            padding: 5px 5px 5px 0;
            vertical-align: top;
        }
        h1 table {
            font-size: 27px;
            text-align: left;
            padding: 0.5rem 0.625rem 0.625rem;
            font-weight: bold;
            padding-right: 10px;
            padding-left: 20px;
            padding: 15px 15px 15px 0;
        }
        h2 table {
            color: #0C8CDB;
            font-size: 16px;
            text-align: left;
            font-weight: bold;
            padding: 5px 5px 5px 0;
            padding-right: 10px;
            padding-left: 20px;
        }
        </style>'''
<commit_msg>Set table width to 100%<commit_after>style = '''<meta id="OGTitle" property="og:title" content="SeleniumBase">
        <title>Test Report</title>
        <link rel="SHORTCUT ICON"
        href="http://cdn2.hubspot.net/hubfs/100006/images/results_logo4.ico" />
        <style type="text/css">
        html {
            background-color: #9988ad;
        }
        html, body {
            font-size: 100%;
            box-sizing: border-box;
        }
        body {
            background-image: none;
            background-origin: padding-box;
            background-color: #ecf2ff;
            padding: 30;
            margin: 10;
            font-family: "Proxima Nova","proxima-nova",
            "Helvetica Neue",Helvetica,Arial,sans-serif !important;
            text-rendering: optimizelegibility;
            -moz-osx-font-smoothing: grayscale;
            box-shadow: 0px 2px 5px 1px rgba(0, 0, 0, 0.24),
            1px 2px 12px 0px rgba(0, 0, 0, 0.18) !important;
        }
        table {
            width: 100%;
            border-collapse: collapse;
            border-spacing: 0;
            box-shadow: 0px 2px 5px 1px rgba(0, 0, 0, 0.25),
            1px 2px 12px 0px rgba(0, 0, 0, 0.19) !important;
            transition: all 0.15s ease-out 0s;
            transition-property: all;
            transition-duration: 0.1s;
            transition-timing-function: ease-out;
            transition-delay: 0s;
        }
        table:hover {
            box-shadow: 0px 2px 5px 1px rgba(0, 0, 0, 0.35),
            1px 2px 12px 0px rgba(0, 0, 0, 0.28) !important;
        }
        thead th, thead td {
            padding: 0.5rem 0.625rem 0.625rem;
            font-weight: bold;
            text-align: left;
        }
        thead {
            text-align: center;
            border: 1px solid #e1e1e1;
            background-color: #fefefe;
            width: 150%;
            color: #0C8CDB;
            background-color: #cbf2ff;
        }
        tbody tr:nth-child(even) {
            background-color: #f1f1f1;
        }
        tbody tr:nth-child(odd) {
            background-color: #ffffff;
        }
        tbody tr:nth-child(even):hover {
            background-color: #f8f8d2;
        }
        tbody tr:nth-child(odd):hover {
            background-color: #ffffe0;
        }
        tbody th, tbody td {
            padding: 0.5rem 0.625rem 0.625rem;
        }
        tbody {
            border: 1px solid #e1e1e1;
            background-color: #fefefe;
        }
        td {
            padding: 5px 5px 5px 0;
            vertical-align: top;
        }
        h1 table {
            font-size: 27px;
            text-align: left;
            padding: 0.5rem 0.625rem 0.625rem;
            font-weight: bold;
            padding-right: 10px;
            padding-left: 20px;
            padding: 15px 15px 15px 0;
        }
        h2 table {
            color: #0C8CDB;
            font-size: 16px;
            text-align: left;
            font-weight: bold;
            padding: 5px 5px 5px 0;
            padding-right: 10px;
            padding-left: 20px;
        }
        </style>'''
<|endoftext|>


<commit_before>from __future__ import unicode_literals

from django.db import models
from django.conf import settings
from django.utils.translation import ugettext_lazy as _
from django.urls import reverse
from django.utils import timezone
from django.utils.encoding import python_2_unicode_compatible

@python_2_unicode_compatible
class TeamMembership(models.Model):
    ROLE_OWNER = 'owner'
    ROLE_EDITOR = 'editor'
    ROLE_VIEWER = 'viewer'
    ROLES = (
        (ROLE_OWNER, _('owner')),
        (ROLE_EDITOR, _('editor')),
        (ROLE_VIEWER, _('viewer')),
    )
    MEMBERSHIP_STATUS_INACTIVE = 'inactive'
    MEMBERSHIP_STATUS_INVITED = 'invited'
    MEMBERSHIP_STATUS_ACTIVE = 'active'
    MEMBERSHIP_STATUS = (
        (MEMBERSHIP_STATUS_INACTIVE, _('inactive')),
        (MEMBERSHIP_STATUS_INVITED, _('invited')),
        (MEMBERSHIP_STATUS_ACTIVE, _('active')),
    )
    user = models.ForeignKey(settings.AUTH_USER_MODEL, null=True, blank=True,
                             on_delete=models.CASCADE)
    email = models.CharField(max_length=255, blank=True)
    team = models.ForeignKey('Team', on_delete=models.CASCADE)
    role = models.CharField(max_length=30, choices=ROLES)
    status = models.CharField(max_length=30, choices=MEMBERSHIP_STATUS)
    created = models.DateTimeField(default=timezone.now)
    updated = models.DateTimeField(default=timezone.now)

    def __str__(self):
        return '%s in %s' % (self.user, self.team)

    def is_active(self):
        return self.status == self.MEMBERSHIP_STATUS_ACTIVE

    def is_invited(self):
        return self.status == self.MEMBERSHIP_STATUS_INVITED

    def is_owner(self):
        return self.role == self.ROLE_OWNER

    def send_invite_mail(self):
        pass

class TeamManager(models.Manager):
    def get_for_user(self, user, *args, **kwargs):
        return self.get_queryset().filter(
            *args,
            teammembership__user=user,
            teammembership__status=TeamMembership.MEMBERSHIP_STATUS_ACTIVE,
            **kwargs
        )

@python_2_unicode_compatible
class Team(models.Model):
    name = models.CharField(max_length=255)
    created = models.DateTimeField(default=timezone.now)

    members = models.ManyToManyField(settings.AUTH_USER_MODEL,
                                     through=TeamMembership)

    objects = TeamManager()

    class Meta:
        verbose_name = _('team')
        verbose_name_plural = _('teams')
        permissions = (
            ("can_use_teams", _("Can use teams")),
        )

    def __str__(self):
        return self.name

    def get_absolute_url(self):
        return reverse('team-detail', kwargs={'pk': self.pk})

    @property
    def member_count(self):
        return self.members.count()

    def can_do(self, user, *args):
        return self.teammembership_set.filter(
            *args,
            status=TeamMembership.MEMBERSHIP_STATUS_ACTIVE,
            user=user,
        ).exists()

    def can_read(self, user):
        return self.can_do(user)

    def can_write(self, user):
        return self.can_do(
            user,
            models.Q(role=TeamMembership.ROLE_EDITOR) |
            models.Q(role=TeamMembership.ROLE_OWNER)
        )

    def can_manage(self, user):
        return self.can_do(
            user,
            models.Q(role=TeamMembership.ROLE_OWNER)
        )
<commit_msg>Fix syntax for Python 2<commit_after>from __future__ import unicode_literals

from django.db import models
from django.conf import settings
from django.utils.translation import ugettext_lazy as _
from django.urls import reverse
from django.utils import timezone
from django.utils.encoding import python_2_unicode_compatible

@python_2_unicode_compatible
class TeamMembership(models.Model):
    ROLE_OWNER = 'owner'
    ROLE_EDITOR = 'editor'
    ROLE_VIEWER = 'viewer'
    ROLES = (
        (ROLE_OWNER, _('owner')),
        (ROLE_EDITOR, _('editor')),
        (ROLE_VIEWER, _('viewer')),
    )
    MEMBERSHIP_STATUS_INACTIVE = 'inactive'
    MEMBERSHIP_STATUS_INVITED = 'invited'
    MEMBERSHIP_STATUS_ACTIVE = 'active'
    MEMBERSHIP_STATUS = (
        (MEMBERSHIP_STATUS_INACTIVE, _('inactive')),
        (MEMBERSHIP_STATUS_INVITED, _('invited')),
        (MEMBERSHIP_STATUS_ACTIVE, _('active')),
    )
    user = models.ForeignKey(settings.AUTH_USER_MODEL, null=True, blank=True,
                             on_delete=models.CASCADE)
    email = models.CharField(max_length=255, blank=True)
    team = models.ForeignKey('Team', on_delete=models.CASCADE)
    role = models.CharField(max_length=30, choices=ROLES)
    status = models.CharField(max_length=30, choices=MEMBERSHIP_STATUS)
    created = models.DateTimeField(default=timezone.now)
    updated = models.DateTimeField(default=timezone.now)

    def __str__(self):
        return '%s in %s' % (self.user, self.team)

    def is_active(self):
        return self.status == self.MEMBERSHIP_STATUS_ACTIVE

    def is_invited(self):
        return self.status == self.MEMBERSHIP_STATUS_INVITED

    def is_owner(self):
        return self.role == self.ROLE_OWNER

    def send_invite_mail(self):
        pass

class TeamManager(models.Manager):
    def get_for_user(self, user, *args, **kwargs):
        return self.get_queryset().filter(
            *args,
            teammembership__user=user,
            teammembership__status=TeamMembership.MEMBERSHIP_STATUS_ACTIVE,
            **kwargs
        )

@python_2_unicode_compatible
class Team(models.Model):
    name = models.CharField(max_length=255)
    created = models.DateTimeField(default=timezone.now)

    members = models.ManyToManyField(settings.AUTH_USER_MODEL,
                                     through=TeamMembership)

    objects = TeamManager()

    class Meta:
        verbose_name = _('team')
        verbose_name_plural = _('teams')
        permissions = (
            ("can_use_teams", _("Can use teams")),
        )

    def __str__(self):
        return self.name

    def get_absolute_url(self):
        return reverse('team-detail', kwargs={'pk': self.pk})

    @property
    def member_count(self):
        return self.members.count()

    def can_do(self, user, *args):
        kwargs = dict(
            status=TeamMembership.MEMBERSHIP_STATUS_ACTIVE,
            user=user
        )
        return self.teammembership_set.filter(
            *args,
            **kwargs
        ).exists()

    def can_read(self, user):
        return self.can_do(user)

    def can_write(self, user):
        return self.can_do(
            user,
            models.Q(role=TeamMembership.ROLE_EDITOR) |
            models.Q(role=TeamMembership.ROLE_OWNER)
        )

    def can_manage(self, user):
        return self.can_do(
            user,
            models.Q(role=TeamMembership.ROLE_OWNER)
        )
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
An incredibly basic, not very flexible database for keeping track of all
the initial conditions that I create.  This is not user friendly, but mostly
for my personal use

EXAMPLE

import icdb

# Create a new database looking at 2 base directories (and all subdirectories)

# Finds all files matching '*_settings.p'

db = icdb.database('*_settings.p', '~/simulations', \
'/home/ibackus/hyak/gscratch/vsm/ibackus/simulations')

# Query the full paths to IC files:

db.data('filename')

# Query Qmin

db.data('sigma.Qmin')

# Access the settings/information for the first simulation

sim_settings = db.data[0]

# Refresh database

db.refresh()

Created on Mon Sep 15 16:06:30 2014

@author: ibackus
"""
import numpy as np
import ICgen_settings
import os
import fnmatch

class database():

    """
    IC database class.  Initialize with:

    db = icdb.database(filefilter, dir1, dir2, ...)

    filefilter is a string used to filter IC settings files, ie '*_settings.p'
    dir1, dir2 are base directories which are recursively searched to find
    settings files
    """

    def __init__(self, filefilter='IC_settings.p', *directories, **kwargs):

        print 'Building database...'
        if len(directories) == 0:

            directories = ['.']
        print 'Number of base directories: ', len(directories)
        self.dirs = directories
        self.filefilter = filefilter
        self._kwargs = kwargs

        matches = []

        for directory in directories:

            print 'Building for ', directory
            new_matches = self._scan_dir(directory, filefilter)
            print '    Found {} files'.format(len(new_matches))
            matches += new_matches

        filelist = list(set(matches))
        datalist = []

        for i, fname in enumerate(filelist):

            a = ICgen_settings.settings()
            a.load(fname)
            a_dir = os.path.split(fname)[0]
            ICname = a.filenames.IC_file_name
            ICname = os.path.join(a_dir, ICname)
            ICname = os.path.realpath(ICname)
            a.filename = ICname

            if os.path.exists(ICname):

                datalist.append(a)

        nfiles = len(datalist)
        print 'Found {} files in total'.format(nfiles)
        data = fancy_array(nfiles)
        data[:] = datalist
        self.data = data

    def nfiles(self):
        """
        Returns the current number of files in the database
        """

        return len(self.data)

    def _scan_dir(self, basedir='.', filefilter='*', kind='IC'):

        if kind is 'IC':

            matches = []

            for root, dirnames, filenames in os.walk(basedir):
                for filename in fnmatch.filter(filenames, filefilter):
                    fname = os.path.join(root, filename)
                    fname = os.path.realpath(fname)

                    matches.append(fname)

            return matches

    def refresh(self):
        """
        Refreshes the database by re-initializing it
        """

        self.__init__(self.filefilter, *self.dirs, **self._kwargs)

class fancy_array(np.ndarray):

    def __new__(subtype, shape, buffer=None, offset=0, \
    strides=None, order=None):

        if not isinstance(shape, int):

            raise ValueError, 'Shape must be an integer.  1D arrays only'

        dtype = object

        obj = np.ndarray.__new__(subtype, shape, dtype, buffer, offset, \
        strides, order)

        return obj

    def __array_finalize__(self, obj):

        if obj is None:

            return

        #self.info = getattr(obj, 'info', None)
        #self.info = obj.info

    def __call__(self, attr):

        outarray = np.ndarray(self.shape, dtype=object)

        for i, entry in enumerate(self):

            if hasattr_nested(entry, attr):

                outarray[i] = getattr_nested(entry, attr)

            else:

                outarray[i] = None

        return outarray

def hasattr_nested(obj, attr):
    """
    Check whether an object contains a specified (possibly nested) attribute

    ie:

    hasattr_nested(obj, 'x.y.z') will check if obj.x.y.z exists
    """

    parts = attr.split('.')

    for part in parts:

        if hasattr(obj, part):
            obj = getattr(obj, part)
        else:
            return False
    else:
        return True

def getattr_nested(obj, attr):
    """
    Get attribute (possibly nested) from object

    ie:

    a = getattr_nested(obj, 'x.y.z')

    is equivalent to a = obj.x.y.z
    """

    parts = attr.split('.')

    for part in parts:

        if hasattr(obj, part):
            obj = getattr(obj, part)
        else:
            return False
    else:
        return obj
            <commit_msg>Simple database to keep track of ICs that have been generated<commit_after># -*- coding: utf-8 -*-
"""
An incredibly basic, not very flexible database for keeping track of all
the initial conditions that I create.  This is not user friendly, but mostly
for my personal use

EXAMPLE

import icdb

# Create a new database looking at 2 base directories (and all subdirectories)

# Finds all files matching '*_settings.p'

db = icdb.database('*_settings.p', '~/simulations', \
'/home/ibackus/hyak/gscratch/vsm/ibackus/simulations')

# Query the full paths to IC files:

db.data('filename')

# Query Qmin

db.data('sigma.Qmin')

# Access the settings/information for the first simulation

sim_settings = db.data[0]

# Refresh database

db.refresh()

Created on Mon Sep 15 16:06:30 2014

@author: ibackus
"""
import numpy as np
import ICgen_settings
import os
import fnmatch

class database():

    """
    IC database class.  Initialize with:

    db = icdb.database(filefilter, dir1, dir2, ...)

    filefilter is a string used to filter IC settings files, ie '*_settings.p'
    dir1, dir2 are base directories which are recursively searched to find
    settings files
    """

    def __init__(self, filefilter='IC_settings.p', *directories, **kwargs):

        print 'Building database...'
        if len(directories) == 0:

            directories = ['.']
        print 'Number of base directories: ', len(directories)
        self.dirs = directories
        self.filefilter = filefilter
        self._kwargs = kwargs

        matches = []

        for directory in directories:

            print 'Building for ', directory
            new_matches = self._scan_dir(directory, filefilter)
            print '    Found {} files'.format(len(new_matches))
            matches += new_matches

        filelist = list(set(matches))
        datalist = []

        for i, fname in enumerate(filelist):

            a = ICgen_settings.settings()
            a.load(fname)
            a_dir = os.path.split(fname)[0]
            ICname = a.filenames.IC_file_name
            ICname = os.path.join(a_dir, ICname)
            ICname = os.path.realpath(ICname)

            if os.path.exists(ICname):

                a.dir = os.path.realpath(a_dir)
                datalist.append(a)

        nfiles = len(datalist)
        print 'Found {} files in total'.format(nfiles)
        data = fancy_array(nfiles)
        data[:] = datalist
        self.data = data

    def nfiles(self):
        """
        Returns the current number of files in the database
        """

        return len(self.data)

    def _scan_dir(self, basedir='.', filefilter='*', kind='IC'):

        if kind is 'IC':

            matches = []

            for root, dirnames, filenames in os.walk(basedir):
                for filename in fnmatch.filter(filenames, filefilter):
                    fname = os.path.join(root, filename)
                    fname = os.path.realpath(fname)

                    matches.append(fname)

            return matches

    def refresh(self):
        """
        Refreshes the database by re-initializing it
        """

        self.__init__(self.filefilter, *self.dirs, **self._kwargs)

class fancy_array(np.ndarray):

    def __new__(subtype, shape, buffer=None, offset=0, \
    strides=None, order=None):

        if not isinstance(shape, int):

            raise ValueError, 'Shape must be an integer.  1D arrays only'

        dtype = object

        obj = np.ndarray.__new__(subtype, shape, dtype, buffer, offset, \
        strides, order)

        return obj

    def __array_finalize__(self, obj):

        if obj is None:

            return

        #self.info = getattr(obj, 'info', None)
        #self.info = obj.info

    def __call__(self, attr):

        outarray = np.ndarray(self.shape, dtype=object)

        for i, entry in enumerate(self):

            if hasattr_nested(entry, attr):

                outarray[i] = getattr_nested(entry, attr)

            else:

                outarray[i] = None

        return outarray

def hasattr_nested(obj, attr):
    """
    Check whether an object contains a specified (possibly nested) attribute

    ie:

    hasattr_nested(obj, 'x.y.z') will check if obj.x.y.z exists
    """

    parts = attr.split('.')

    for part in parts:

        if hasattr(obj, part):
            obj = getattr(obj, part)
        else:
            return False
    else:
        return True

def getattr_nested(obj, attr):
    """
    Get attribute (possibly nested) from object

    ie:

    a = getattr_nested(obj, 'x.y.z')

    is equivalent to a = obj.x.y.z
    """

    parts = attr.split('.')

    for part in parts:

        if hasattr(obj, part):
            obj = getattr(obj, part)
        else:
            return False
    else:
        return obj
            <|endoftext|>


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

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

class DAO(object):

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

all_query = {
    "query" : {
        "match_all" : { }
    }
}<commit_msg>add dao operation which wraps query results in the object type<commit_after>import uuid, json
from esprit import raw, util
from copy import deepcopy

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

class DAO(object):

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

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

    @classmethod
    def object_query(cls, q='', terms=None, should_terms=None, facets=None, conn=None, **kwargs):
        j = cls.query(q=q, terms=terms, should_terms=should_terms, facets=facets, conn=conn, **kwargs)
        res = raw.unpack_json_result(j)
        return [cls(r) for r in res]

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

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

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

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

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

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

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

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


<commit_before>#!/usr/bin/env python3
import json
import lzma
import hashlib
import requests
import struct
import subprocess
import os
from typing import Generator

SPARSE_CHUNK_FMT = struct.Struct('H2xI4x')

class StreamingDecompressor:
  def __init__(self, url: str) -> None:
    self.buf = b""

    self.req = requests.get(url, stream=True, headers={'Accept-Encoding': None})
    self.it = self.req.iter_content(chunk_size=1024 * 1024)
    self.decompressor = lzma.LZMADecompressor(format=lzma.FORMAT_AUTO)
    self.eof = False
    self.sha256 = hashlib.sha256()

  def read(self, length: int) -> bytes:
    while len(self.buf) < length:
      self.req.raise_for_status()

      try:
        compressed = next(self.it)
      except StopIteration:
        self.eof = True
        break
      out = self.decompressor.decompress(compressed)
      self.buf += out

    result = self.buf[:length]
    self.buf = self.buf[length:]

    self.sha256.update(result)
    return result

def unsparsify(f: StreamingDecompressor) -> Generator[bytes, None, None]:
  # https://source.android.com/devices/bootloader/images#sparse-format
  magic = struct.unpack("I", f.read(4))[0]
  assert(magic == 0xed26ff3a)

  # Version
  major = struct.unpack("H", f.read(2))[0]
  minor = struct.unpack("H", f.read(2))[0]
  assert(major == 1 and minor == 0)

  f.read(2)  # file header size
  f.read(2)  # chunk header size

  block_sz = struct.unpack("I", f.read(4))[0]
  f.read(4)  # total blocks
  num_chunks = struct.unpack("I", f.read(4))[0]
  f.read(4)  # crc checksum

  for _ in range(num_chunks):
    chunk_type, out_blocks = SPARSE_CHUNK_FMT.unpack(f.read(12))

    if chunk_type == 0xcac1:  # Raw
      # TODO: yield in smaller chunks. Yielding only block_sz is too slow. Largest observed data chunk is 252 MB.
      yield f.read(out_blocks * block_sz)
    elif chunk_type == 0xcac2:  # Fill
      filler = f.read(4) * (block_sz // 4)
      for _ in range(out_blocks):
        yield filler
    elif chunk_type == 0xcac3:  # Don't care
      yield b""
    else:
      raise Exception("Unhandled sparse chunk type")

def get_target_slot_number() -> int:
  current_slot = subprocess.check_output(["abctl", "--boot_slot"], encoding='utf-8').strip()
  return 1 if current_slot == "_a" else 0

def slot_number_to_suffix(slot_number: int) -> str:
  assert slot_number in (0, 1)
  return '_a' if slot_number == 0 else '_b'

def get_partition_path(target_slot_number: int, partition: dict) -> str:
  path = f"/dev/disk/by-partlabel/{partition['name']}"

  if partition.get('has_ab', True):
    path += slot_number_to_suffix(target_slot_number)

  return path

def verify_partition(target_slot_number: int, partition: dict) -> bool:
  full_check = partition['full_check']
  path = get_partition_path(target_slot_number, partition)
  partition_size = partition['size']

  with open(path, 'rb+') as out:
    if full_check:
      raw_hash = hashlib.sha256()

      pos = 0
      chunk_size = 1024 * 1024
      while pos < partition_size:
        n = min(chunk_size, partition_size - pos)
        raw_hash.update(out.read(n))
        pos += n

      return raw_hash.hexdigest().lower() == partition['hash_raw'].lower()
    else:
      out.seek(partition_size)
      return out.read(64) == partition['hash_raw'].lower().encode()

def clear_partition_hash(target_slot_number: int, partition: dict) -> None:
  path = get_partition_path(target_slot_number, partition)
  with open(path, 'wb+') as out:
    partition_size = partition['size']

    out.seek(partition_size)
    out.write(b"\x00" * 64)
    os.sync()

def flash_partition(target_slot_number: int, partition: dict, cloudlog):
  cloudlog.info(f"Downloading and writing {partition['name']}")

  if verify_partition(target_slot_number, partition):
    cloudlog.info(f"Already flashed {partition['name']}")
    return

  downloader = StreamingDecompressor(partition['url'])

  # Clear hash before flashing in case we get interrupted
  full_check = partition['full_check']
  if not full_check:
    clear_partition_hash(target_slot_number, partition)

  path = get_partition_path(target_slot_number, partition)
  with open(path, 'wb+') as out:
    partition_size = partition['size']

    # Flash partition
    if partition['sparse']:
      raw_hash = hashlib.sha256()
      for chunk in unsparsify(downloader):
        raw_hash.update(chunk)
        out.write(chunk)
        p = int(out.tell() / partition_size * 100)
        print(f"Installing {partition['name']}: {p}")

      if raw_hash.hexdigest().lower() != partition['hash_raw'].lower():
        raise Exception(f"Unsparse hash mismatch '{raw_hash.hexdigest().lower()}'")
    else:
      while not downloader.eof:
        out.write(downloader.read(1024 * 1024))

    if downloader.sha256.hexdigest().lower() != partition['hash'].lower():
      raise Exception("Uncompressed hash mismatch")

    if out.tell() != partition['size']:
      raise Exception("Uncompressed size mismatch")

    # Write hash after successfull flash
    os.sync()
    if not full_check:
      out.write(partition['hash_raw'].lower().encode())

def swap(manifest_path: str, target_slot_number: int, cloudlog) -> None:
  update = json.load(open(manifest_path))
  for partition in update:
    if not partition.get('full_check', False):
      clear_partition_hash(target_slot_number, partition)

  while True:
    out = subprocess.check_output(f"abctl --set_active {target_slot_number}", shell=True, stderr=subprocess.STDOUT, encoding='utf8')
    if ("No such file or directory" not in out) and ("lun as boot lun" in out):
      cloudlog.info(f"Swap successfull {out}")
      break
    else:
      cloudlog.error(f"Swap failed {out}")

def flash_agnos_update(manifest_path: str, target_slot_number: int, cloudlog) -> None:
  update = json.load(open(manifest_path))

  cloudlog.info(f"Target slot {target_slot_number}")

  # set target slot as unbootable
  os.system(f"abctl --set_unbootable {target_slot_number}")

  for partition in update:
    success = False

    for retries in range(10):
      try:
        flash_partition(target_slot_number, partition, cloudlog)
        success = True
        break

      except requests.exceptions.RequestException:
        cloudlog.exception("Failed")
        cloudlog.info(f"Failed to download {partition['name']}, retrying ({retries})")
        time.sleep(10)

    if not success:
      cloudlog.info(f"Failed to flash {partition['name']}, aborting")
      raise Exception("Maximum retries exceeded")

  cloudlog.info(f"AGNOS ready on slot {target_slot_number}")

def verify_agnos_update(manifest_path: str, target_slot_number: int) -> bool:
  update = json.load(open(manifest_path))
  return all(verify_partition(target_slot_number, partition) for partition in update)

if __name__ == "__main__":
  import logging
  import time
  import argparse

  parser = argparse.ArgumentParser(description="Flash and verify AGNOS update",
                                   formatter_class=argparse.ArgumentDefaultsHelpFormatter)

  parser.add_argument("--verify", action="store_true", help="Verify and perform swap if update ready")
  parser.add_argument("--swap", action="store_true", help="Verify and perform swap, downloads if necessary")
  parser.add_argument("manifest", help="Manifest json")
  args = parser.parse_args()

  logging.basicConfig(level=logging.INFO)

  target_slot_number = get_target_slot_number()
  if args.verify:
    if verify_agnos_update(args.manifest, target_slot_number):
      swap(args.manifest, target_slot_number, logging)
      exit(0)
    exit(1)
  elif args.swap:
    while not verify_agnos_update(args.manifest, target_slot_number):
      logging.error("Verification failed. Flashing AGNOS")
      flash_agnos_update(args.manifest, target_slot_number, logging)

    logging.warning(f"Verification succeeded. Swapping to slot {target_slot_number}")
    swap(args.manifest, target_slot_number, logging)
  else:
    flash_agnos_update(args.manifest, target_slot_number, logging)
<commit_msg>agnos updater: support non-sparse images (#22371)<commit_after>#!/usr/bin/env python3
import json
import lzma
import hashlib
import requests
import struct
import subprocess
import os
from typing import Generator

SPARSE_CHUNK_FMT = struct.Struct('H2xI4x')

class StreamingDecompressor:
  def __init__(self, url: str) -> None:
    self.buf = b""

    self.req = requests.get(url, stream=True, headers={'Accept-Encoding': None})
    self.it = self.req.iter_content(chunk_size=1024 * 1024)
    self.decompressor = lzma.LZMADecompressor(format=lzma.FORMAT_AUTO)
    self.eof = False
    self.sha256 = hashlib.sha256()

  def read(self, length: int) -> bytes:
    while len(self.buf) < length:
      self.req.raise_for_status()

      try:
        compressed = next(self.it)
      except StopIteration:
        self.eof = True
        break
      out = self.decompressor.decompress(compressed)
      self.buf += out

    result = self.buf[:length]
    self.buf = self.buf[length:]

    self.sha256.update(result)
    return result

def unsparsify(f: StreamingDecompressor) -> Generator[bytes, None, None]:
  # https://source.android.com/devices/bootloader/images#sparse-format
  magic = struct.unpack("I", f.read(4))[0]
  assert(magic == 0xed26ff3a)

  # Version
  major = struct.unpack("H", f.read(2))[0]
  minor = struct.unpack("H", f.read(2))[0]
  assert(major == 1 and minor == 0)

  f.read(2)  # file header size
  f.read(2)  # chunk header size

  block_sz = struct.unpack("I", f.read(4))[0]
  f.read(4)  # total blocks
  num_chunks = struct.unpack("I", f.read(4))[0]
  f.read(4)  # crc checksum

  for _ in range(num_chunks):
    chunk_type, out_blocks = SPARSE_CHUNK_FMT.unpack(f.read(12))

    if chunk_type == 0xcac1:  # Raw
      # TODO: yield in smaller chunks. Yielding only block_sz is too slow. Largest observed data chunk is 252 MB.
      yield f.read(out_blocks * block_sz)
    elif chunk_type == 0xcac2:  # Fill
      filler = f.read(4) * (block_sz // 4)
      for _ in range(out_blocks):
        yield filler
    elif chunk_type == 0xcac3:  # Don't care
      yield b""
    else:
      raise Exception("Unhandled sparse chunk type")

# noop wrapper with same API as unsparsify() for non sparse images
def noop(f: StreamingDecompressor) -> Generator[bytes, None, None]:
  while not f.eof:
    yield f.read(1024 * 1024)

def get_target_slot_number() -> int:
  current_slot = subprocess.check_output(["abctl", "--boot_slot"], encoding='utf-8').strip()
  return 1 if current_slot == "_a" else 0

def slot_number_to_suffix(slot_number: int) -> str:
  assert slot_number in (0, 1)
  return '_a' if slot_number == 0 else '_b'

def get_partition_path(target_slot_number: int, partition: dict) -> str:
  path = f"/dev/disk/by-partlabel/{partition['name']}"

  if partition.get('has_ab', True):
    path += slot_number_to_suffix(target_slot_number)

  return path

def verify_partition(target_slot_number: int, partition: dict) -> bool:
  full_check = partition['full_check']
  path = get_partition_path(target_slot_number, partition)
  partition_size = partition['size']

  with open(path, 'rb+') as out:
    if full_check:
      raw_hash = hashlib.sha256()

      pos = 0
      chunk_size = 1024 * 1024
      while pos < partition_size:
        n = min(chunk_size, partition_size - pos)
        raw_hash.update(out.read(n))
        pos += n

      return raw_hash.hexdigest().lower() == partition['hash_raw'].lower()
    else:
      out.seek(partition_size)
      return out.read(64) == partition['hash_raw'].lower().encode()

def clear_partition_hash(target_slot_number: int, partition: dict) -> None:
  path = get_partition_path(target_slot_number, partition)
  with open(path, 'wb+') as out:
    partition_size = partition['size']

    out.seek(partition_size)
    out.write(b"\x00" * 64)
    os.sync()

def flash_partition(target_slot_number: int, partition: dict, cloudlog):
  cloudlog.info(f"Downloading and writing {partition['name']}")

  if verify_partition(target_slot_number, partition):
    cloudlog.info(f"Already flashed {partition['name']}")
    return

  downloader = StreamingDecompressor(partition['url'])

  # Clear hash before flashing in case we get interrupted
  full_check = partition['full_check']
  if not full_check:
    clear_partition_hash(target_slot_number, partition)

  path = get_partition_path(target_slot_number, partition)
  with open(path, 'wb+') as out:
    # Flash partition
    last_p = 0
    raw_hash = hashlib.sha256()
    f = unsparsify if partition['sparse'] else noop
    for chunk in f(downloader):
      raw_hash.update(chunk)
      out.write(chunk)
      p = int(out.tell() / partition['size'] * 100)
      if p != last_p:
        last_p = p
        print(f"Installing {partition['name']}: {p}")

    if raw_hash.hexdigest().lower() != partition['hash_raw'].lower():
      raise Exception(f"Raw hash mismatch '{raw_hash.hexdigest().lower()}'")

    if downloader.sha256.hexdigest().lower() != partition['hash'].lower():
      raise Exception("Uncompressed hash mismatch")

    if out.tell() != partition['size']:
      raise Exception("Uncompressed size mismatch")

    # Write hash after successfull flash
    os.sync()
    if not full_check:
      out.write(partition['hash_raw'].lower().encode())

def swap(manifest_path: str, target_slot_number: int, cloudlog) -> None:
  update = json.load(open(manifest_path))
  for partition in update:
    if not partition.get('full_check', False):
      clear_partition_hash(target_slot_number, partition)

  while True:
    out = subprocess.check_output(f"abctl --set_active {target_slot_number}", shell=True, stderr=subprocess.STDOUT, encoding='utf8')
    if ("No such file or directory" not in out) and ("lun as boot lun" in out):
      cloudlog.info(f"Swap successfull {out}")
      break
    else:
      cloudlog.error(f"Swap failed {out}")

def flash_agnos_update(manifest_path: str, target_slot_number: int, cloudlog) -> None:
  update = json.load(open(manifest_path))

  cloudlog.info(f"Target slot {target_slot_number}")

  # set target slot as unbootable
  os.system(f"abctl --set_unbootable {target_slot_number}")

  for partition in update:
    success = False

    for retries in range(10):
      try:
        flash_partition(target_slot_number, partition, cloudlog)
        success = True
        break

      except requests.exceptions.RequestException:
        cloudlog.exception("Failed")
        cloudlog.info(f"Failed to download {partition['name']}, retrying ({retries})")
        time.sleep(10)

    if not success:
      cloudlog.info(f"Failed to flash {partition['name']}, aborting")
      raise Exception("Maximum retries exceeded")

  cloudlog.info(f"AGNOS ready on slot {target_slot_number}")

def verify_agnos_update(manifest_path: str, target_slot_number: int) -> bool:
  update = json.load(open(manifest_path))
  return all(verify_partition(target_slot_number, partition) for partition in update)

if __name__ == "__main__":
  import logging
  import time
  import argparse

  parser = argparse.ArgumentParser(description="Flash and verify AGNOS update",
                                   formatter_class=argparse.ArgumentDefaultsHelpFormatter)

  parser.add_argument("--verify", action="store_true", help="Verify and perform swap if update ready")
  parser.add_argument("--swap", action="store_true", help="Verify and perform swap, downloads if necessary")
  parser.add_argument("manifest", help="Manifest json")
  args = parser.parse_args()

  logging.basicConfig(level=logging.INFO)

  target_slot_number = get_target_slot_number()
  if args.verify:
    if verify_agnos_update(args.manifest, target_slot_number):
      swap(args.manifest, target_slot_number, logging)
      exit(0)
    exit(1)
  elif args.swap:
    while not verify_agnos_update(args.manifest, target_slot_number):
      logging.error("Verification failed. Flashing AGNOS")
      flash_agnos_update(args.manifest, target_slot_number, logging)

    logging.warning(f"Verification succeeded. Swapping to slot {target_slot_number}")
    swap(args.manifest, target_slot_number, logging)
  else:
    flash_agnos_update(args.manifest, target_slot_number, logging)
<|endoftext|>