ody_migr_mmth_endpoint.py

#!/usr/bin/env python
"""Main script to import json data for any data category from Mammoth."""


import os
import requests
import logging

from ody_migr_utils import generate_hash, get_valid_json
from ody_migr_config import (ENV_MIGR_SECRET_KEY,
                             DATA_OPT_MEMDATA,
                             DATA_OPT_SWITCH_ACCOUNTS,
                             DATA_OPT_CREATE_HIST,
                             DATA_OPT_GET_HIST,
                             DATA_OPT_GETCULIST,
                             DEFAULT_ERROR_STR)

LOGGER = logging.getLogger(__name__)


class MammothMigration:
    """Retrieve response from Mammoth."""

    def __init__(self,
                 _cu,
                 _data_category,
                 _server,
                 _username,
                 _password,
                 **kwargs):
        self.cu = _cu
        self.data_category = _data_category
        self.server = _server

        # being extra cautious here (execution should not reach inside this
        # condition through the options provided in the Monitor screen)
        if self.server == "www" and self.data_category != DATA_OPT_GETCULIST:
            err_msg = "{} server must only be used for {} operation.".format(
                self.server, DATA_OPT_GETCULIST)
            LOGGER.error(err_msg)
            raise SystemExit(err_msg)

        if self.data_category == DATA_OPT_GETCULIST:
            # required for generating valid passphrase (consistency
            # with Mammoth's process
            assert self.cu == ""

        self.username = _username
        self.password = _password
        self.payload = None
        self.getstream = False
        self.histstatus = ""
        self.endpoint_script = kwargs.get(
            "endpoint_script", "hcuadm/mOdysseyMigrExp.prg")

        if kwargs != {}:
            if (_data_category == DATA_OPT_MEMDATA or
                    _data_category == DATA_OPT_CREATE_HIST or
                    _data_category == DATA_OPT_SWITCH_ACCOUNTS):

                assert "params" in kwargs
                self.payload = kwargs.get("params")
                assert self.payload is not None

            elif _data_category == DATA_OPT_GET_HIST:
                assert "stream" in kwargs
                self.getstream = kwargs.get("stream")

        self.response = {}

    def _dest_url(self):
        """Generate a Mammoth endpoint url"""
        hash = generate_hash(self.data_category, self.cu,
                             os.environ.get(ENV_MIGR_SECRET_KEY))
        server_dest = "https://{}.homecu.net".format(self.server)
        if self.data_category == DATA_OPT_GETCULIST:
            resource_dest = ("{}?"
                             "passphrase={}&"
                             "action={}").format(self.endpoint_script,
                                                 hash,
                                                 self.data_category)

        else:
            resource_dest = ("{}?"
                             "cu={}&"
                             "passphrase={}&"
                             "action={}").format(self.endpoint_script,
                                                 self.cu,
                                                 hash,
                                                 self.data_category)
        return os.path.join(server_dest, resource_dest)

    def _validate_response(self):
        """Validate response data returned from Mammoth.

        Raises:
             SystemExit on error
        """

        # at this point, response dict must have already been populated
        error_flag = False
        resp_error = DEFAULT_ERROR_STR

        if self.data_category == DATA_OPT_GET_HIST:
            # check for error in application/json response
            if dict(self.response.headers)[
                    "Content-Type"].strip() == "application/json":
                resp_error = self.response.json()["error"]
                if resp_error != DEFAULT_ERROR_STR:
                    if "Process is already started" in resp_error:
                        # do not consider this case as error
                        if "ERROR:" in resp_error:
                            resp_error = resp_error.split("ERROR:")[1].strip()
                        self.histstatus = resp_error
                    else:
                        error_flag = True

            # check for error in text/plain response
            if dict(self.response.headers)[
                    "Content-Type"].split(";")[0].strip() == "text/plain":
                resp_error = self.response.content.decode()
                if "Process is already started" in resp_error:
                    # do not consider this case as error
                    if "ERROR:" in resp_error:
                        resp_error = resp_error.split("ERROR:")[1].strip()
                    self.histstatus = resp_error
                else:
                    error_flag = True

        else:
            assert self.response != {}
            resp_error = self.response["error"]
            if resp_error != DEFAULT_ERROR_STR:
                if "Process is already started" in resp_error:
                    # do not consider this case as error
                    if "ERROR:" in resp_error:
                        resp_error = resp_error.split("ERROR:")[1].strip()
                    self.histstatus = resp_error
                else:
                    error_flag = True

        if error_flag:
            LOGGER.error("Mammoth endpoint status: {}".format(resp_error))
            raise SystemExit(
                "Migration Request with `cu = {} and "
                "data_category = {}` returned with error. Error: {}"
                .format(
                    self.cu,
                    self.data_category,
                    resp_error))
        else:
            resp_error = "NORMAL!" if resp_error == DEFAULT_ERROR_STR \
                else resp_error
            LOGGER.info("Mammoth endpoint status: {}".format(resp_error))

    def _initiate_http_request(self):
        """Establish connection with Mammoth and obtain response object"""

        try:
            # we run http connection within the requests context manager
            # in order to avoid possible resource leaks
            with requests.Session() as s:
                url = self._dest_url()
                if self.payload is None:
                    http_resp = s.get(
                        url,
                        auth=(self.username, self.password),
                        stream=self.getstream
                    )
                else:
                    http_resp = s.post(
                        url,
                        auth=(self.username, self.password),
                        data=self.payload
                    )

                status_code = http_resp.status_code
                if int(status_code) >= 300:
                    http_resp.raise_for_status()

                LOGGER.info("HTTP response url: {}".format(http_resp.url))
                if int(status_code) == 200:
                    if self.data_category == DATA_OPT_GET_HIST:
                        self.response = http_resp

                    else:
                        self.response = http_resp.json()

        except (requests.exceptions.HTTPError,
                requests.exceptions.ConnectionError) as err:
            LOGGER.error("Mammoth endpoint status: {}".format(err))
            raise SystemExit(err)

    def run(self):
        """Connect to mammoth endpoint and fetch response json data"""
        self._initiate_http_request()
        self._validate_response()
        return self.response

    def __str__(self):
        return get_valid_json(self.response)