#!/usr/bin/env python3

'''
Cytomic Orion MISP Module
An expansion module to enrich attributes in MISP and share indicators of compromise with Cytomic Orion


'''

from pymisp import MISPAttribute, MISPEvent, MISPObject
import json
import requests
import sys

misperrors = {'error': 'Error'}
mispattributes = {'input': ['md5'], 'format': 'misp_standard'}
moduleinfo = {'version': '0.3', 'author': 'Koen Van Impe',
              'description': 'an expansion module to enrich attributes in MISP and share indicators of compromise with Cytomic Orion',
              'module-type': ['expansion']}
moduleconfig = ['api_url', 'token_url', 'clientid', 'clientsecret', 'clientsecret', 'username', 'password', 'upload_timeframe', 'upload_tag', 'delete_tag', 'upload_ttlDays', 'upload_threat_level_id', 'limit_upload_events', 'limit_upload_attributes']
# There are more config settings in this module than used by the enrichment
# There is also a PyMISP module which reuses the module config, and requires additional configuration, for example used for pushing indicators to the API


class CytomicParser():
    def __init__(self, attribute, config_object):
        self.misp_event = MISPEvent()
        self.attribute = MISPAttribute()
        self.attribute.from_dict(**attribute)
        self.misp_event.add_attribute(**self.attribute)

        self.config_object = config_object

        if self.config_object:
            self.token = self.get_token()
        else:
            sys.exit('Missing configuration')

    def get_token(self):
        try:
            scope = self.config_object['scope']
            grant_type = self.config_object['grant_type']
            username = self.config_object['username']
            password = self.config_object['password']
            token_url = self.config_object['token_url']
            clientid = self.config_object['clientid']
            clientsecret = self.config_object['clientsecret']

            if scope and grant_type and username and password:
                data = {'scope': scope, 'grant_type': grant_type, 'username': username, 'password': password}

                if token_url and clientid and clientsecret:
                    access_token_response = requests.post(token_url, data=data, verify=False, allow_redirects=False, auth=(clientid, clientsecret))
                    tokens = json.loads(access_token_response.text)
                    if 'access_token' in tokens:
                        return tokens['access_token']
                    else:
                        self.result = {'error': 'No token received.'}
                        return
                else:
                    self.result = {'error': 'No token_url, clientid or clientsecret supplied.'}
                    return
            else:
                self.result = {'error': 'No scope, grant_type, username or password supplied.'}
                return
        except Exception:
            self.result = {'error': 'Unable to connect to token_url.'}
            return

    def get_results(self):
        if hasattr(self, 'result'):
            return self.result
        event = json.loads(self.misp_event.to_json())
        results = {key: event[key] for key in ('Attribute', 'Object')}
        return {'results': results}

    def parse(self, searchkey):

        if self.token:

            endpoint_fileinformation = self.config_object['endpoint_fileinformation']
            endpoint_machines = self.config_object['endpoint_machines']
            endpoint_machines_client = self.config_object['endpoint_machines_client']
            query_machines = self.config_object['query_machines']
            query_machine_info = self.config_object['query_machine_info']

            # Update endpoint URLs
            query_endpoint_fileinformation = endpoint_fileinformation.format(md5=searchkey)
            query_endpoint_machines = endpoint_machines.format(md5=searchkey)

            # API calls
            api_call_headers = {'Authorization': 'Bearer ' + self.token}
            result_query_endpoint_fileinformation = requests.get(query_endpoint_fileinformation, headers=api_call_headers, verify=False)
            json_result_query_endpoint_fileinformation = json.loads(result_query_endpoint_fileinformation.text)

            if json_result_query_endpoint_fileinformation:

                cytomic_object = MISPObject('cytomic-orion-file')

                cytomic_object.add_attribute('fileName', type='text', value=json_result_query_endpoint_fileinformation['fileName'])
                cytomic_object.add_attribute('fileSize', type='text', value=json_result_query_endpoint_fileinformation['fileSize'])
                cytomic_object.add_attribute('last-seen', type='datetime', value=json_result_query_endpoint_fileinformation['lastSeen'])
                cytomic_object.add_attribute('first-seen', type='datetime', value=json_result_query_endpoint_fileinformation['firstSeen'])
                cytomic_object.add_attribute('classification', type='text', value=json_result_query_endpoint_fileinformation['classification'])
                cytomic_object.add_attribute('classificationName', type='text', value=json_result_query_endpoint_fileinformation['classificationName'])
                self.misp_event.add_object(**cytomic_object)

                result_query_endpoint_machines = requests.get(query_endpoint_machines, headers=api_call_headers, verify=False)
                json_result_query_endpoint_machines = json.loads(result_query_endpoint_machines.text)

                if query_machines and json_result_query_endpoint_machines and len(json_result_query_endpoint_machines) > 0:
                    for machine in json_result_query_endpoint_machines:

                        if query_machine_info and machine['muid']:
                            query_endpoint_machines_client = endpoint_machines_client.format(muid=machine['muid'])
                            result_endpoint_machines_client = requests.get(query_endpoint_machines_client, headers=api_call_headers, verify=False)
                            json_result_endpoint_machines_client = json.loads(result_endpoint_machines_client.text)

                            if json_result_endpoint_machines_client:

                                cytomic_machine_object = MISPObject('cytomic-orion-machine')

                                clienttag = [{'name': json_result_endpoint_machines_client['clientName']}]

                                cytomic_machine_object.add_attribute('machineName', type='target-machine', value=json_result_endpoint_machines_client['machineName'], Tag=clienttag)
                                cytomic_machine_object.add_attribute('machineMuid', type='text', value=machine['muid'])
                                cytomic_machine_object.add_attribute('clientName', type='target-org', value=json_result_endpoint_machines_client['clientName'], Tag=clienttag)
                                cytomic_machine_object.add_attribute('clientId', type='text', value=machine['clientId'])
                                cytomic_machine_object.add_attribute('machinePath', type='text', value=machine['lastPath'])
                                cytomic_machine_object.add_attribute('first-seen', type='datetime', value=machine['firstSeen'])
                                cytomic_machine_object.add_attribute('last-seen', type='datetime', value=machine['lastSeen'])
                                cytomic_machine_object.add_attribute('creationDate', type='datetime', value=json_result_endpoint_machines_client['creationDate'])
                                cytomic_machine_object.add_attribute('clientCreationDateUTC', type='datetime', value=json_result_endpoint_machines_client['clientCreationDateUTC'])
                                cytomic_machine_object.add_attribute('lastSeenUtc', type='datetime', value=json_result_endpoint_machines_client['lastSeenUtc'])
                                self.misp_event.add_object(**cytomic_machine_object)
        else:
            self.result = {'error': 'No (valid) token.'}
            return


def handler(q=False):
    if q is False:
        return False
    request = json.loads(q)

    if not request.get('attribute'):
        return {'error': 'Unsupported input.'}

    attribute = request['attribute']
    if not any(input_type == attribute['type'] for input_type in mispattributes['input']):
        return {'error': 'Unsupported attributes type'}

    if not request.get('config'):
        return {'error': 'Missing configuration'}

    config_object = {
        'clientid': request["config"].get("clientid"),
        'clientsecret': request["config"].get("clientsecret"),
        'scope': 'orion.api',
        'password': request["config"].get("password"),
        'username': request["config"].get("username"),
        'grant_type': 'password',
        'token_url': request["config"].get("token_url"),
        'endpoint_fileinformation': '{api_url}{endpoint}'.format(api_url=request["config"].get("api_url"), endpoint='/forensics/md5/{md5}/info'),
        'endpoint_machines': '{api_url}{endpoint}'.format(api_url=request["config"].get("api_url"), endpoint='/forensics/md5/{md5}/muids'),
        'endpoint_machines_client': '{api_url}{endpoint}'.format(api_url=request["config"].get("api_url"), endpoint='/forensics/muid/{muid}/info'),
        'query_machines': True,
        'query_machine_info': True
    }

    cytomic_parser = CytomicParser(attribute, config_object)
    cytomic_parser.parse(attribute['value'])

    return cytomic_parser.get_results()


def introspection():
    return mispattributes


def version():
    moduleinfo['config'] = moduleconfig
    return moduleinfo