from __future__ import print_function

from future import standard_library

standard_library.install_aliases()
from builtins import object
import configparser
import json
import sys
from datetime import datetime

import dateutil.parser
import dateutil.tz
import pytz
import requests
import requests.auth
import textract


class TruStar(object):
    """
    Main class you to instantiate the TruStar API
    """

    def __init__(self, config_file="trustar.conf", config_role="trustar"):

        self.enclaveIds = []
        self.attributedToMe = False

        config_parser = configparser.RawConfigParser()
        config_parser.read(config_file)

        try:
            # parse required properties
            self.auth = config_parser.get(config_role, 'auth_endpoint')
            self.base = config_parser.get(config_role, 'api_endpoint')
            self.apikey = config_parser.get(config_role, 'user_api_key')
            self.apisecret = config_parser.get(config_role, 'user_api_secret')

            # parse enclave an attribution properties
            if config_parser.has_option(config_role, 'enclave_ids'):
                self.enclaveIds = filter(None, config_parser.get(config_role, 'enclave_ids').split(','))

            if config_parser.has_option(config_role, 'attribute_reports'):
                self.attributedToMe = config_parser.getboolean(config_role, 'attribute_reports')
        except Exception as e:
            print("Problem reading config file: %s", e)
            sys.exit(1)

    @staticmethod
    def normalize_timestamp(datetime_str):
        """
        Attempt to convert a string timestamp in to a TruSTAR compatible format for submission.
        Will return current time with UTC time zone if None
        :param datetime_str: raw timestamp containing date, time and ideally timezone
        """
        try:
            datetime_dt = dateutil.parser.parse(datetime_str)
        except Exception as e:
            datetime_dt = datetime.now()

        if not datetime_dt.tzinfo:
            datetime_dt = datetime_dt.replace(tzinfo=pytz.utc)

        return datetime_dt.isoformat()

    def get_token(self):
        """
        Retrieves the OAUTH token generated by your API key and API secret.
        this function has to be called before any API calls can be made
        """
        client_auth = requests.auth.HTTPBasicAuth(self.apikey, self.apisecret)
        post_data = {"grant_type": "client_credentials"}
        resp = requests.post(self.auth, auth=client_auth, data=post_data)
        token_json = resp.json()
        return token_json["access_token"]

    def get_latest_reports(self, access_token):
        """
        Retrieves the latest 10 reports submitted to the TruSTAR community
        :param access_token: OAuth API token
        """

        headers = {"Authorization": "Bearer " + access_token}
        resp = requests.get(self.base + "/reports/latest", headers=headers)
        return json.loads(resp.content.decode('utf8'))

    def get_correlated_reports(self, access_token, indicator):
        """
        Retrieves all TruSTAR reports that contain the searched indicator. You can specify multiple indicators
        separated by commas
        :param indicator:
        :param access_token:
        """

        headers = {"Authorization": "Bearer " + access_token}
        payload = {'q': indicator}
        resp = requests.get(self.base + "/reports/correlate", payload, headers=headers)
        return json.loads(resp.content)

    def query_indicator(self, access_token, indicator, limit):
        """
        Finds all reports that contain the indicators and returns correlated indicators from those reports.
        you can specify the limit of indicators returned.
        :param limit:
        :param indicator:
        :param access_token:
        """

        headers = {"Authorization": "Bearer " + access_token}
        payload = {'q': indicator, 'limit': limit}

        resp = requests.get(self.base + "/indicators", payload, headers=headers)
        return json.loads(resp.content)

    def query_latest_indicators(self,
                                access_token,
                                source,
                                indicator_types,
                                limit,
                                interval_size):
        """
        Finds all latest indicators
        :param self:
        :param access_token:
        :param source: source of the indicators which can either be INCIDENT_REPORT or OSINT
        :param interval_size: time interval on returned indicators. Max is set to 24 hours
        :param limit: limit on the number of indicators. Max is set to 5000
        :param indicator_types: a list of indicators or a string equal to "ALL" to query all indicator types extracted
        by TruSTAR
        :return json response of the result
        """

        headers = {"Authorization": "Bearer " + access_token}
        payload = {'source': source, 'types': indicator_types, 'limit': limit, 'intervalSize': interval_size}
        resp = requests.get(self.base + "/indicators/latest", payload, headers=headers)
        return json.loads(resp.content)

    def submit_report(self, access_token, report_body_txt, report_name, discovered_time_str=None,
                      enclave=False):
        """
        Wraps supplied text as a JSON-formatted TruSTAR Incident Report and submits it to TruSTAR Station
        By default, this submits to the TruSTAR community. To submit to your enclave, pass in your enclave_id
        :param discovered_time_str:
        :param enclave: boolean - whether or not to submit report to user's enclaves (see 'enclave_ids' config property)
        :param report_name:
        :param report_body_txt:
        :param access_token:
        """

        # Convert timestamps
        distribution_type = 'ENCLAVE' if enclave else 'COMMUNITY'
        if distribution_type == 'ENCLAVE' and len(self.enclaveIds) < 1:
            raise Exception("Must specify one or more enclave IDs to submit enclave reports into")

        headers = {'Authorization': 'Bearer ' + access_token, 'content-Type': 'application/json'}

        payload = {'incidentReport': {
            'title': report_name,
            'timeDiscovered': self.normalize_timestamp(discovered_time_str),
            'timeBegan': self.normalize_timestamp(discovered_time_str),
            'reportBody': report_body_txt,
            'distributionType': distribution_type},
            'enclaveIds': self.enclaveIds,
            'attributedToMe': self.attributedToMe}

        print("Submitting report %s to TruSTAR Station..." % report_name)
        resp = requests.post(self.base + "/reports/submit", json.dumps(payload, encoding="ISO-8859-1"), headers=headers,
                             timeout=60)
        return resp.json()

    @staticmethod
    def process_file(source_file):
        if source_file.endswith(('.pdf', '.PDF')):
            txt = textract.process(source_file, method='pdfminer')
        elif source_file.endswith(('.doc', '.docx', '.txt', '.eml', '.xls', '.xlsx', '.csv', '.rtf', '.json')):
            txt = textract.process(source_file)
        else:
            raise ValueError('UNSUPPORTED FILE EXTENSION')
        return txt
