api_wrapper.py

"""
Copyright (c) 2012-2013 RockStor, Inc. <http://rockstor.com>
This file is part of RockStor.

RockStor is free software; you can redistribute it and/or modify
it under the terms of the GNU General Public License as published
by the Free Software Foundation; either version 2 of the License,
or (at your option) any later version.

RockStor is distributed in the hope that it will be useful, but
WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
General Public License for more details.

You should have received a copy of the GNU General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
"""

import requests
import time
import json
import base64
from storageadmin.exceptions import RockStorAPIException
from storageadmin.models import OauthApp
from django.conf import settings


class APIWrapper(object):

    def __init__(self, client_id=None, client_secret=None, url=None):
        self.access_token = None
        self.expiration = time.time()
        self.client_id = client_id
        self.client_secret = client_secret
        # directly connect to gunicorn, bypassing nginx as we are on the same
        # host.
        self.url = 'http://127.0.0.1:443'
        if (url is not None):
            # for remote urls.
            self.url = url

    def set_token(self):
        if (self.client_id is None or self.client_secret is None):
            app = OauthApp.objects.get(name=settings.OAUTH_INTERNAL_APP)
            self.client_id = app.application.client_id
            self.client_secret = app.application.client_secret

        token_request_data = {
            'grant_type': 'client_credentials',
            'client_id': self.client_id,
            'client_secret': self.client_secret,
        }

        user_pass = '{0}:{1}'.format(self.client_id, self.client_secret)
        auth_string = base64.b64encode(user_pass.encode('utf-8'))
        auth_headers = {'HTTP_AUTHORIZATION':
                        'Basic ' + auth_string.decode("utf-8"), }
        content = None
        try:
            response = requests.post('%s/o/token/' % self.url,
                                     data=token_request_data,
                                     headers=auth_headers, verify=False)
            content = json.loads(response.content.decode("utf-8"))
            self.access_token = content['access_token']
            self.expiration = int(time.time()) + content['expires_in'] - 600
        except Exception as e:
            msg = ('Exception while setting access_token for url(%s): %s. '
                   'content: %s' % (self.url, e.__str__(), content))
            raise Exception(msg)

    def api_call(self, url, data=None, calltype='get', headers=None,
                 save_error=True):
        if (self.access_token is None or
                time.time() > self.expiration):
            self.set_token()

        api_auth_header = {'Authorization': 'Bearer ' + self.access_token, }
        call = getattr(requests, calltype)
        url = ('%s/api/%s' % (self.url, url))
        try:
            if (headers is not None):
                headers.update(api_auth_header)
                if (headers['content-type'] == 'application/json'):
                    r = call(url, verify=False, data=json.dumps(data),
                             headers=headers)
                else:
                    r = call(url, verify=False, data=data,
                             headers=headers)
            else:
                r = call(url, verify=False, headers=api_auth_header, data=data)
        except requests.exceptions.ConnectionError:
            print('Error connecting to Rockstor. Is it running?')
            raise

        if (r.status_code == 404):
            msg = ('Invalid api end point: %s' % url)
            raise RockStorAPIException(detail=msg)

        if (r.status_code != 200):
            try:
                error_d = json.loads(r.text)
                if (settings.DEBUG is True and save_error is True):
                    cur_time = str(int(time.time()))
                    err_file = '/tmp/err-%s.html' % cur_time
                    with open(err_file, 'w') as efo:
                        for line in r.text.split('\n'):
                            efo.write('%s\n' % line)
                        print('Error detail is saved at %s' % err_file)
                if ('detail' in error_d):
                    if (error_d['detail'] == 'Authentication credentials were not provided.'):  # noqa E501
                        self.set_token()
                        return self.api_call(url, data=data, calltype=calltype,
                                             headers=headers,
                                             save_error=save_error)
                    raise RockStorAPIException(detail=error_d['detail'])
            except ValueError as e:
                raise RockStorAPIException(detail='Internal Server Error: %s'
                                           % e.__str__())
            r.raise_for_status()

        try:
            ret_val = r.json()
        except ValueError:
            ret_val = {}
        return ret_val