OwlCyberSecurity - MANAGER

Edit File: api_client.py

# -*- coding: utf-8 -*-

# Copyright © Cloud Linux GmbH & Cloud Linux Software, Inc 2010-2021 All Rights Reserved
#
# Licensed under CLOUD LINUX LICENSE AGREEMENT
# http://cloudlinux.com/docs/LICENSE.TXT

"""
This module contains class implementing MongoDB API interaction
"""
import dbm
import hashlib
import pwd
import json
import logging
import urllib.parse
import uuid
from functools import partial
from typing import List, Any, Iterable

from requests import Session, Response
from requests.adapters import HTTPAdapter
from requests.exceptions import RequestException
from requests.packages.urllib3.util.retry import Retry
from schema import Schema, SchemaError
from clwpos.papi import (
    is_feature_visible,
    is_feature_hidden_server_wide,
)

from ..adviser.advice_types import supported as supported_advice_types, get_advice_instance
from xray.apiclient.schemas import (
    detailed_advice_schema,
    user_sites_info_schema,
    advice_list_schema
)
from ..internal.constants import api_server, proto, adviser_api_server
from ..internal.exceptions import XRayAPIError, XRayAPIEmptyResponse, TaskNotFoundError
from ..internal.local_counters import open_local_storage
from ..internal.types import Task
from ..internal.user_plugin_utils import (
    get_xray_exec_user,
    user_mode_verification
)
from ..internal.utils import read_jwt_token

def _does_user_exist_on_server(username):
    try:
        pwd.getpwnam(username)
        return True
    except KeyError:
        return False

class Client:
    """
    Base client class
    """

def __init__(self, *, system_id: str,
                 tracing_task_id: str = 'unavailable'):
        self.system_id = system_id
        self.task_id = tracing_task_id
        self.logger = logging.getLogger('api_client')

retry_conf = Retry(total=3,
                           allowed_methods=frozenset(['GET', 'POST']),
                           status_forcelist=frozenset([502, 503, 504]),
                           backoff_factor=3)  # sleeps 0s, 6s, 18s
        adapter = HTTPAdapter(max_retries=retry_conf)
        self.session = Session()
        self.session.mount(f'{proto}://', adapter)
        self.session.request = partial(self.session.request, timeout=10)

def __repr__(self):
        return f'{self.__class__.__name__}::{self.main_endpoint}::tracing_task_id={self.task_id}'

def __str__(self):
        return f'{self.__class__.__name__}::{self.task_id}'

@property
    def main_endpoint(self) -> str:
        """
        Base endpoint
        """
        raise NotImplementedError('instances are to set their endpoint!')

@property
    def _headers(self) -> dict:
        """
        Updated request headers
        :return: dict with headers
        """
        return {
            'X-Auth': read_jwt_token()
        }

def _query_string(self, **kwargs) -> str:
        """
        Construct query string
        :return: string including system_id and given kwargs
        """
        initial = {'system_id': self.system_id}
        if kwargs:
            initial.update(kwargs)
        return urllib.parse.urlencode(initial)

def _preprocess_response(self, response: Response,
                             with_status_check: bool = True) -> Any:
        """
        Perform preprocessing checks of received Response:
        - is it OK, e.g. 200 OK
        - was it successful, e.g. status == ok
        and extract JSON from the Response object
        :return: full JSON representation of response
        """
        if not response.ok:
            # to not duplicate in logs
            self.logger.debug('Server responded: %s', response.text)
            raise XRayAPIError(
                f'Unable to connect to server with {response.status_code}:{response.reason}',
                extra={'resp_data': response.text})
        else:
            self.logger.info('[%s:%s] Response received %s',
                             response.status_code, response.reason,
                             response.url)

# sometimes our services could return 204 code, which returns NO_CONTENT => json() fails
        if response.status_code != 204:
            result = response.json()
        else:
            result = {'status': 'ok'}
        if with_status_check:
            if result['status'] != 'ok':
                raise XRayAPIError(f"Received unsuccessful response: {result}")
        return result

def _process_response(self, response: Response,
                          with_status_check: bool = True) -> dict:
        """
        Check received response
        :param response: a requests.Response object
        :return: 'result' dict from JSON representation of response
        """
        result = self._preprocess_response(response,
                                           with_status_check=with_status_check)
        try:
            data = result['result']
            self.logger.info('[%s] Received response data %s',
                             response.status_code, data)
            if data is None:
                raise XRayAPIEmptyResponse(result=result)
            return data
        except KeyError:
            return dict()

def _post(self, endpoint: str, payload: dict = None, log_data: bool = True,
              with_status_check: bool = True, include_jwt: bool = True) -> dict:
        """
        Perform POST request to given endpoint.
        Add payload as JSON if given
        :param endpoint: target URL
        :param payload: dict with date to POST
        :param log_data: whether to log POST data or not
        :return: 'result' dict from JSON representation of response
        """
        self.logger.info('Sending POST request to %s',
                         endpoint)
        if log_data and payload:
            self.logger.info('Data attached to POST: %s', payload)
        try:
            if include_jwt:
                headers = self._headers
            else:
                headers = {}
            if payload is None:
                resp = self.session.post(endpoint, headers=headers)
            else:
                resp = self.session.post(endpoint,
                                         json=payload, headers=headers, timeout=60)
        except RequestException as e:
            raise self._give_xray_exception(e, endpoint, 'POST failed',
                                            'Failed to POST data to X-Ray API server') from e
        return self._process_response(resp,
                                      with_status_check=with_status_check)

def _delete(self, endpoint: str, log_data: bool = True,
                with_status_check: bool = True):
        self.logger.info('Sending DELETE request to %s',
                         endpoint)
        try:
            resp = self.session.delete(endpoint, headers=self._headers)
        except RequestException as e:
            raise self._give_xray_exception(e, endpoint, 'DELETE failed',
                                            'Failed to DE:ETE data to X-Ray API server') from e
        return self._process_response(resp,
                                      with_status_check=with_status_check)

def _raw_get(self, endpoint: str = None) -> Response:
        """
        GET request to endpoint or to main endpoint if no endpoint given
        :param endpoint: target URL
        :return: a requests Response object
        """
        if endpoint is None:
            endpoint = f'{self.main_endpoint}?{self._query_string()}'
        self.logger.info('Sending GET request to %s', endpoint)
        try:
            resp = self.session.get(endpoint, headers=self._headers)
        except RequestException as e:
            raise self._give_xray_exception(e, endpoint, 'GET failed',
                                            'Failed to GET data from X-Ray API server') from e
        return resp

def _get_full(self, endpoint: str = None,
                  with_status_check: bool = True) -> Any:
        """
        GET request to endpoint or to main endpoint if no endpoint given
        :param endpoint: target URL
        :return: full dict from JSON representation of response without any
                 processing
        """
        resp = self._raw_get(endpoint)
        return self._preprocess_response(resp,
                                         with_status_check=with_status_check)

def _get(self, endpoint: str = None) -> dict:
        """
        GET request to endpoint or to main endpoint if no endpoint given
        :param endpoint: target URL
        :return: 'result' dict from JSON representation of response
        """
        resp = self._raw_get(endpoint)
        try:
            return self._process_response(resp)
        except XRayAPIEmptyResponse as e:
            raise TaskNotFoundError(
                task_id=self.task_id
            ) from e

def _give_xray_exception(self, exc, api_endpoint, log_message,
                             exc_message):
        """
        Process received exception
        :param exc: original exception
        :param api_endpoint: requested endpoint
        :param log_message: text for logging the error
        :param exc_message: text for internal exception
        """
        self.logger.error('%s with %s', log_message, exc,
                          extra={'endpoint': api_endpoint})
        try:
            exc_info = exc.args[0].reason
        except (IndexError, AttributeError):
            exc_info = exc

return XRayAPIError(
            f'{exc_message}: {exc_info}. Please, try again later')

class TaskMixin:
    """
    A mixin class with Task related methods
    """

@property
    def task_fields(self) -> tuple:
        """
        Limit processed fields
        """
        return ("url", "status", "client_ip", "tracing_by", "tracing_count",
                "starttime", "ini_location", "initial_count", "request_count",
                "auto_task", "user")

def _task(self, dict_view: dict) -> Task:
        """
        Turn dictionary structure into valid Task type
        """
        task_view = {k: v for k, v in dict_view.items() if
                     k in self.task_fields}
        task_view['task_id'] = dict_view['tracing_task_id']
        return Task(**task_view)

class TasksClient(Client, TaskMixin):
    """
    'tasks' endpoint client
    """

@property
    def main_endpoint(self) -> str:
        """
        Base endpoint: tasks
        """
        return f'{proto}://{api_server}/api/xray/tasks'

def _query_string(self, **kwargs) -> str:
        """
        Construct query string.
        Aimed to get auto tasks only
        :return: string including system_id and type=auto
        """
        return super()._query_string(type='auto')

def get_tasks(self) -> List[Task]:
        """
        Get list of Tasks
        """
        data = super()._get()
        return [self._task(item) for item in data]

class DBMClient(Client):
    """
    Client class using local dbm storage instead of remote API
    """

def __init__(self, *, system_id: str,
                 tracing_task_id: str = 'unavailable'):
        super().__init__(system_id=system_id, tracing_task_id=tracing_task_id)
        self.task_object = None
        self._db = self._db_open()

def __del__(self):
        self._db.close()

@staticmethod
    def _db_open() -> 'gdbm object':
        """
        Open dbm DB
        :return: corresponding object
        """
        return dbm.open('/root/local_mongo', 'c')

def _post(self, post_data: dict) -> None:
        """
        Update a DBM task with given data
        """
        self._db[self.task_id] = json.dumps(post_data)

def _get(self) -> dict:
        """
        Get saved DBM data
        :return: dict
        """
        try:
            return json.loads(self._db[self.task_id].decode())
        except (KeyError, json.JSONDecodeError) as e:
            raise XRayAPIError('Failed to load task') from e

@staticmethod
    def _id() -> str:
        return uuid.uuid1().hex

def get_task(self) -> Task:
        """
        Get saved task
        :return:
        """
        saved_task = self._get()
        saved_task['task_id'] = self.task_id
        self.task_object = Task(**saved_task)
        return self.task_object

def create(self, task: Task) -> str:
        """
        Create new task and get unique ID
            url --> URL
            client_ip --> IP
            tracing_by --> time|request_qty
            tracing_count --> COUNT
            ini_location --> PATH
            status --> processing
        :param task: a Task instance
        :return: task ID
        """
        self.task_id = self._id()
        task.task_id = self.task_id
        task.status = 'hold'
        self._post(task.as_dict())
        self.task_object = task
        return self.task_id

def update(self, starttime: int) -> None:
        """
        Update started|continued task
            status --> running
            starttime --> new timestamp
        :return:
        """
        if self.task_object is None:
            self.get_task()
        self.task_object.status = 'running'
        self.task_object.starttime = starttime
        self._post(self.task_object.as_dict())

def stop(self, count: int) -> None:
        """
        Update stopped task
            status --> stopped
            tracing_count --> new value
        :return:
        """
        if self.task_object is None:
            self.get_task()
        self.task_object.status = 'stopped'
        self.task_object.tracing_count = count
        self._post(self.task_object.as_dict())

def complete(self) -> None:
        """
        Complete tracing task
            status --> completed
        :return:
        """
        if self.task_object is None:
            self.get_task()
        self.task_object.status = 'completed'
        self._post(self.task_object.as_dict())

def delete(self) -> None:
        """
        Delete tracing task
        :return:
        """
        del self._db[self.task_id]

class APIClient(Client, TaskMixin):
    """
    X-Ray task API client class
    """

@property
    def main_endpoint(self) -> str:
        """
        Base endpoint: task
        """
        return f'{proto}://{api_server}/api/xray/task'

def _query_string(self, **kwargs) -> str:
        """
        Construct query string
        :return: string either including system_id and task_id
                or system_id only
        """
        if self.task_id != 'unavailable':
            return super()._query_string(tracing_task_id=self.task_id)
        return super()._query_string()

def _post_create(self, post_data: dict) -> None:
        """
        POST request to "create a task" API endpoint with given data
        Obtains a task ID
        :param post_data: dict with POST data
        """
        endpoint = f'{self.main_endpoint}/create?{self._query_string()}'
        response_data = self._post(endpoint,
                                   {k: v for k, v in post_data.items() if
                                    k != 'starttime'})
        self.task_id = response_data['tracing_task_id']

def _post_update(self, post_data: dict) -> None:
        """
        POST request to "update a task" API endpoint with given data
        :param post_data: dict with POST data
        """
        endpoint = f'{self.main_endpoint}/update?{self._query_string()}'
        self._post(endpoint, post_data)

def _share(self) -> None:
        """
        GET request to "share a task" API endpoint
        """
        share_endpoint = self.main_endpoint[:-5]
        endpoint = f'{share_endpoint}/share-request?{self._query_string()}'
        self._get(endpoint)

def _delete(self) -> None:
        """
        POST request to "delete a task" API endpoint
        """
        endpoint = f'{self.main_endpoint}/delete?{self._query_string()}'
        self._post(endpoint)

@user_mode_verification
    def get_task(self) -> Task:
        """
        Get saved task
        :return:
        """
        return self._task(self._get())

def create(self, task: Task) -> Task:
        """
        Create new task and get unique ID
            url --> URL
            client_ip --> IP
            tracing_by --> time|request_qty
            tracing_count --> COUNT
            ini_location --> PATH
            status --> processing
        :param task: a Task instance
        :return: updated Task instance
        """
        task.status = 'hold'
        self._post_create({k: v for k, v in task.as_dict().items() if
                           k in self.task_fields})
        return self.task_id

def update(self, starttime: int) -> None:
        """
        Update started|continued task
            status --> running
            starttime --> new timestamp
        :param starttime: time of starting the Task
        """
        self._post_update({'status': 'running',
                           'starttime': starttime})

def update_count_only(self, count: int) -> None:
        """
        Update tracing_count only. No status updated
            tracing_count --> new value
        :return:
        """
        self._post_update({'tracing_count': count})

def update_counts_only(self, *, request_count: int,
                           tracing_count: int = None) -> None:
        """
        Update tracing_count only. No status updated
            request_count --> new value
            tracing_count --> new value if given
        :param request_count: number of requests already traced
        :param tracing_count: number of requests left to trace
        """
        if tracing_count is None:
            data = {'request_count': request_count}
        else:
            data = {'request_count': request_count,
                    'tracing_count': tracing_count}
        self._post_update(data)

def stop(self, count: int) -> None:
        """
        Update stopped task
            status --> stopped
            tracing_count --> new value
        :return:
        """
        self._post_update({'status': 'stopped',
                           'tracing_count': count})

def complete(self) -> None:
        """
        Complete tracing task
            status --> completed
        :return:
        """
        self._post_update({'status': 'completed'})

def share(self) -> None:
        """
        Share tracing task
        :return:
        """
        self._share()

def delete(self) -> None:
        """
        Delete tracing task
        :return:
        """
        self._delete()

class SendClient(Client):
    """
    X-Ray requests API client class
    """

@property
    def main_endpoint(self) -> str:
        """
        Base endpoint: requests
        """
        return f'{proto}://{api_server}/api/xray/requests'

def __call__(self, data: dict) -> None:
        """
        Send given data to ClickHouse
        :param data: dict with data
        """
        endpoint = f'{self.main_endpoint}?{self._query_string()}'
        self._post(endpoint, data, log_data=False)

class UIAPIClient(Client):
    """
    X-Ray User plugin API client class
    """

@property
    def main_endpoint(self) -> str:
        """
        Base endpoint: requests
        """
        return f'{proto}://{api_server}/api/xray'

def _query_string(self, **kwargs) -> str:
        """
        Construct query string
        :return: string including system_id
                 and given kwargs, filtered by non-empty values
        """
        filter_empty = {k: v for k, v in kwargs.items() if v is not None}
        return super()._query_string(**filter_empty)

def get_task_list(self) -> dict:
        """
        Get list of tasks and return not processed (full) response
        from API server
        """
        qs = self._query_string(user=get_xray_exec_user())
        endpoint = f'{self.main_endpoint}/tasks?{qs}'
        response = self._get_full(endpoint)
        for task in response['result']:
            # mix up local data for all tasks except completed
            # completed tasks have all the data actual in mongo
            if task['status'] == 'completed':
                continue

fake_id = hashlib.blake2b(task['tracing_task_id'].encode(),
                                      digest_size=10).hexdigest()
            with open_local_storage(fake_id) as storage:
                if task['tracing_by'] != 'time':
                    task['tracing_count'] = task['initial_count'] - storage.processed_requests
                task['request_count'] = storage.processed_requests
        return response

def get_request_list(self, task_id: str) -> dict:
        """
        Get list of requests collected for given tracing task
        """
        qs = self._query_string(tracing_task_id=task_id)
        endpoint = f'{self.main_endpoint}/requests?{qs}'
        return self._get_full(endpoint)

def get_request_data(self, task_id: str, request_id: int) -> dict:
        """
        Get collected statistics for given request ID of given tracing task
        """
        qs = self._query_string(tracing_task_id=task_id,
                                request_id=request_id)
        endpoint = f'{self.main_endpoint}/request?{qs}'
        return self._get_full(endpoint)

class SmartAdviceAPIClient(Client):
    """
    X-Ray Adviser API client class
    """

def __init__(self):
        super().__init__(system_id='not_needed')

def _validate(self, data: Any, schema: Schema) -> Any:
        """Validate given data using given schema"""
        try:
            return schema.validate(data)
        except SchemaError as e:
            self.logger.error('Failed to validate API response: %s', data)
            msg = e.errors[-1] or e.autos[-1]
            raise XRayAPIError(f'Malformed API response: {msg}')

@property
    def main_endpoint(self) -> str:
        """
        Base endpoint: requests
        """
        return f'https://{adviser_api_server}/api'

@property
    def fields_allowed(self) -> tuple:
        """
        Limit fields available for update
        """
        return ("status", "source", "reason")

def _query_string(self, **kwargs) -> str:
        """
        Construct query string
        :return: string including types and given kwargs
        """
        initial = [('type', _t) for _t in supported_advice_types]
        user_context = get_xray_exec_user()
        if user_context:
            initial.append(('username', user_context))
        initial.extend([(k, v) for k, v in kwargs.items() if v])
        return urllib.parse.urlencode(initial, safe=',')

def __call__(self, data: dict) -> None:
        """
        Send given data to Adviser microservice
        :param data: dict with data
        """
        endpoint = f'{self.main_endpoint}/requests/add'
        self._post(endpoint, data, log_data=False,
                   with_status_check=False)

def _patch(self, endpoint: str, payload: dict = None) -> Any:
        """
        Perform PATCH request to given endpoint.
        Add payload as JSON.
        :param endpoint: target URL
        :param payload: dict with data to PATCH
        :return: full response
        """
        self.logger.info('Sending PATCH request to %s',
                         endpoint)

try:
            resp = self.session.patch(endpoint,
                                      json=payload,
                                      headers=self._headers)
        except RequestException as e:
            raise self._give_xray_exception(e, endpoint, 'PATCH failed',
                                            'Failed to PATCH data to Smart Advice API server') from e

return self._preprocess_response(resp, with_status_check=False)

def send_stat(self, data: dict) -> None:
        """
        Send statistics to Adviser microservice
        """
        endpoint = f'{self.main_endpoint}/requests/metadata'
        self._post(endpoint, data, with_status_check=False)

def _filter_advice_list(self, advice_list: List[dict]):
        """
        Loop over advices and remove those which have
        non-existing users and those which are invisible.
        :param advice_list: list of advices received from API
        """

visible_advices = []
        for item in advice_list:
            # skip advices which appear to be linked with non existing users
            if not _does_user_exist_on_server(item['metadata']['username']):
                logging.info('User %s does not exist anymore on this server,'
                             'skipping advices for him', item['metadata']['username'])
                continue

advice_instance = get_advice_instance(item['advice']['type'])
            if is_feature_visible(advice_instance.module_name,
                                  item['metadata']['username']) and \
                    not is_feature_hidden_server_wide(advice_instance.module_name):
                visible_advices.append(item)
        return visible_advices

def advice_list(self, filtered: bool = True, show_all: bool = False) -> List:
        """
        Get list of advice
        :param filtered: Automatically removes invisible advices
                         and those which are inked to non-existing users.
        """
        endpoint = f'{self.main_endpoint}/advice/list?{self._query_string(show_all=show_all)}'
        response = self._get_full(endpoint, with_status_check=False)

response = self._validate(
            data=response,
            schema=advice_list_schema)

if filtered:
            response = self._filter_advice_list(response)
        return response

def site_info(self, username) -> List:
        """
        Get urls/advices information per user`s site
        """
        endpoint = f'{self.main_endpoint}/advice/site_info/{username}'

response = self._get_full(endpoint, with_status_check=False)
        return self._validate(
            data=response,
            schema=user_sites_info_schema)

def advice_details(self, advice_id: int) -> dict:
        """
        Get details of an advice by given advice_id
        """
        endpoint = f'{self.main_endpoint}/v2/advice/{advice_id}/details'

response = self._get_full(endpoint, with_status_check=False)
        return self._validate(
            data=response,
            schema=detailed_advice_schema)

def update_advice(self, advice_id: int, **kwargs) -> Any:
        """
        Partial update of an advice by given advice_id.
        Fields allowed for update are limited by fields_allowed property
        """
        data = {k: v for k, v in kwargs.items() if
                k in self.fields_allowed}
        endpoint = f'{self.main_endpoint}/advice/{advice_id}'
        return self._patch(endpoint, data)

def report(self, data: dict) -> Any:
        """
        Sends analytics data to the microservice
        """
        endpoint = f'{self.main_endpoint}/analytics/events'
        return self._post(endpoint, data, with_status_check=False, include_jwt=False)

class AWPProvisionAPIClient(Client):
    """
    X-Ray Adviser API client class
    """

def __init__(self):
        super().__init__(system_id='not_needed')

def _query_string(self, **kwargs) -> str:
        return urllib.parse.urlencode(kwargs)

@property
    def main_endpoint(self) -> str:
        """
        Base endpoint: requests
        """
        return f'https://{adviser_api_server}/awp'

def _process_response(self, response: Response,
                          with_status_check: bool = True) -> dict:
        return self._preprocess_response(response,
                                         with_status_check=with_status_check)

def get_create_pullzone(self, account_id: str, domain: str, website: str):
        """
        Gets pullzone if already exists, otherwise creates it
        """
        endpoint = f'{self.main_endpoint}/cdn/pullzone'
        return self._post(endpoint, {'account_id': account_id,
                                     'original_url': domain,
                                     'website': website},
                          log_data=False,
                   with_status_check=False)

def remove_pullzone(self, account_id: str, domain: str, website: str):
        """
        Gets pullzone if already exists, otherwise creates it
        """
        endpoint = f'{self.main_endpoint}/cdn/pullzone'
        return self._delete(f'{endpoint}?{self._query_string(account_id=account_id, original_url=domain, website=website)}',
                            log_data=False,
                            with_status_check=False)

def purge_cdn_cache(self, account_id: str, domain: str, website: str):
        endpoint = f'{self.main_endpoint}/cdn/purge'
        return self._post(endpoint, {'account_id': account_id,
                                     'original_url': domain,
                                     'website': website},
                          log_data=False,
                          with_status_check=False)

def sync_account(self, account_id: Iterable[str]):
        endpoint = f'{self.main_endpoint}/public/account/sync'
        return self._post(endpoint, {'account_id': account_id},
                          log_data=False,
                          with_status_check=False)

def get_usage(self, account_id: str):
        endpoint = f'{self.main_endpoint}/cdn/usage?{self._query_string(account_id=account_id)}'
        return self._get(endpoint)