# Copyright (c) 2017, The MITRE Corporation
# For license information, see the LICENSE.txt file

import argparse
import io
import os
import sys
import traceback
import datetime
import libtaxii.clients as tc
import six
from six.moves.configparser import SafeConfigParser
from six.moves.urllib.parse import urlparse

import libtaxii as t
from libtaxii.common import gen_filename
from libtaxii.constants import *


EXIT_SUCCESS = 0
EXIT_FAILURE = 1

#: When writing content blocks, clobber any existing files
W_CLOBBER = 0
#: When writing content blocks, skip any existing files
W_SKIP = 1
#: When writing content blocks, prompt before clobbering any existing files
W_PROMPT = 2
#: Do not write any files
W_NONE = 3
W_CHOICES = (W_CLOBBER, W_PROMPT, W_SKIP, W_NONE)

MSG_FILE_OVERWRITTEN = "File overwritten: "
MSG_FILE_SKIPPED = "File write skipped: "
MSG_FILE_CREATED = "File created: "


def write_type(s):
    s = s.lower()
    if s == 'clobber':
        return W_CLOBBER
    elif s == 'skip':
        return W_SKIP
    elif s == 'prompt':
        return W_PROMPT
    elif s == 'none':
        return W_NONE
    else:
        raise ValueError('Argument must be one of \'clobber\', \'skip\', \'prompt\', or \'none\'')


def add_poll_response_args(parser):
    """
    Adds arguments to the parser that are commonly used for handling Poll Response messages:
    - --dest-dir
    - --write

    :param parser:
    :return:
    """
    parser.add_argument("--dest-dir",
                        dest="dest_dir",
                        default="",
                        help="The directory to save Content Blocks to. Defaults to the current directory.")

    parser.add_argument("--write",
                        dest="write_type",
                        default='clobber',
                        type=write_type,
                        help="The write type. Must be one of \'clobber\', \'skip\', \'prompt\', or \'none\'. Defaults "
                             "to \'clobber\'")


class ArgParserConfig(SafeConfigParser):

    def as_args(self, section, raw=False, vars=None):
        """
        Returns format friendly list for parse_args e.g., ["--url", "http://hailataxii.com:80", "--arg1", "something"]
        """
        arg_list = []
        for arg, val in self.items(section, raw, vars):
            arg_list.append("--" + arg)
            arg_list.append(val)
        return arg_list


class ProxyAction(argparse.Action):

    def __call__(self, parser, namespace, values, option_string=None):
        """
        Turns the 'None' string argument into the None object
        that the API expects.
        """
        if option_string == "--proxy" and values == "None":
            values = None

        setattr(namespace, self.dest, values)


class LoadFromFile(argparse.Action):

    def __call__(self, parser, namespace, values, option_string=None):
        config_file = ArgParserConfig()
        with values as f:
            config_file.readfp(f)

        # Overrides initial values if present on config file
        parser.parse_known_args(config_file.as_args('libtaxii'), namespace=namespace)


class TaxiiScript(object):

    #: taxii version to put in the request headers. Defaults to TAXII 1.1
    taxii_version = VID_TAXII_XML_11
    #: parser description
    parser_description = ''
    #: default path
    path = '/'

    def get_arg_parser(self,
                       parser_description,
                       path="/taxii-discovery-service",
                       host="hailataxii.com",
                       port="80",
                       https=False,
                       cert=None,
                       key=None,
                       username=None,
                       password=None,
                       proxy='noproxy',
                       xml_output=False,
                       from_file=None):
        """
        Parser things common to all scripts. Parsers for specific TAXII Services should
        add their own arguments.
        """
        parser = argparse.ArgumentParser(description=parser_description)
        url = "https" if https else "http" + "://" + host + ":" + str(port) + path
        parser.add_argument("-u",
                            "--url",
                            dest="url",
                            default=url,
                            help="The URL to connect to. Defaults to %s." % url)
        parser.add_argument("--host",
                            dest="host",
                            default=host,
                            help=argparse.SUPPRESS
                            )
        parser.add_argument("--port",
                            dest="port",
                            default=port,
                            type=int,
                            help=argparse.SUPPRESS)
        parser.add_argument("--path",
                            dest="path",
                            default=path,
                            help=argparse.SUPPRESS)
        parser.add_argument("--https",
                            dest="https",
                            default=https,
                            type=bool,
                            help=argparse.SUPPRESS)
        parser.add_argument("--cert",
                            dest="cert",
                            default=cert,
                            help="The file location of the certificate to use. Defaults to %s." % cert)
        parser.add_argument("--key",
                            dest="key",
                            default=key,
                            help="The file location of the private key to use. Defaults to %s." % key)
        parser.add_argument("--username",
                            dest="username",
                            default=username,
                            help="The username to authenticate with. Defaults to %s." % username)
        parser.add_argument("--pass",
                            dest="password",
                            default=password,
                            help="The password to authenticate with. Defaults to %s." % password)
        parser.add_argument("--proxy",
                            dest="proxy",
                            action=ProxyAction, default=proxy,
                            help="The proxy to use (e.g., http://myproxy.example.com:80/), or 'noproxy' to not use "
                                 "any proxy. If omitted, 'noproxy' will be used.")
        parser.add_argument("--xml-output",
                            dest="xml_output",
                            action='store_true',
                            default=xml_output,
                            help="If present, the raw XML of the response will be printed to standard out. "
                                 "Otherwise, a \"Rich\" output will be presented.")
        parser.add_argument("--from-file",
                            dest="from_file",
                            action=LoadFromFile,
                            default=from_file,
                            type=open,
                            help="Use a configuration file to load arguments into the script. The contents of the "
                                 "configuration file will take precedence over passed flags.")

        return parser

    def handle_response(self, response, args):
        """
        Default response handler. Just prints the response
        """
        print("Response:\n")
        if args.xml_output is False:
            print(response.to_text())
        else:
            print(response.to_xml(pretty_print=True))

    def create_client(self, use_https, proxy, cert=None, key=None, username=None, password=None):
        client = tc.HttpClient()
        client.set_use_https(use_https)
        client.set_proxy(proxy)
        tls = (cert is not None and key is not None)
        basic = (username is not None and password is not None)
        if tls and basic:
            client.set_auth_type(tc.HttpClient.AUTH_CERT_BASIC)
            client.set_auth_credentials({'key_file': key,
                                         'cert_file': cert,
                                         'username': username,
                                         'password': password})
        elif tls:
            client.set_auth_type(tc.HttpClient.AUTH_CERT)
            client.set_auth_credentials({'key_file': key, 'cert_file': cert})
        elif basic:
            client.set_auth_type(tc.HttpClient.AUTH_BASIC)
            client.set_auth_credentials({'username': username, 'password': password})

        return client

    @staticmethod
    def get_write_and_message(filename, write_type_):
        """
        Depending on whether the file exists and the write_type, decide whether the file should be created or not
        and provide a message to be printed out to the user.
        """
        file_exists = os.path.exists(filename)

        if file_exists and write_type_ == W_CLOBBER:
            write = True
            message = MSG_FILE_OVERWRITTEN
        elif file_exists and write_type_ == W_PROMPT:
            var = None
            while var not in ('y', 'n'):
                var = six.moves.input("Overwrite file (%s)? (y/n): " % filename)
            if var == 'y':
                write = True
                message = MSG_FILE_OVERWRITTEN
            else:
                write = False
                message = MSG_FILE_SKIPPED
        elif not file_exists and write_type_ in (W_CLOBBER, W_PROMPT):
            write = True
            message = MSG_FILE_CREATED
        elif write_type_ in (W_SKIP, W_NONE):  # This happens no matter if the file exists or not
            write = False
            message = MSG_FILE_SKIPPED
        else:
            raise ValueError("Unknown combination of file_exists (%s) and write_type (%s)" % (file_exists, write_type_))

        return write, message

    def write_cbs_from_poll_response_10(self, poll_response, dest_dir, write_type_=W_CLOBBER):
        """
        This function writes content blocks to file from a TAXII 1.0 Poll Response
        """

        for cb in poll_response.content_blocks:
            if cb.content_binding == CB_STIX_XML_10:
                format_ = '_STIX10_'
                ext = '.xml'
            elif cb.content_binding == CB_STIX_XML_101:
                format_ = '_STIX101_'
                ext = '.xml'
            elif cb.content_binding == CB_STIX_XML_11:
                format_ = '_STIX11_'
                ext = '.xml'
            elif cb.content_binding == CB_STIX_XML_111:
                format_ = '_STIX111_'
                ext = '.xml'
            elif cb.content_binding == CB_STIX_XML_12:
                format_ = '_STIX12_'
            else:  # Format and extension are unknown
                format_ = ''
                ext = ''

        if cb.timestamp_label:
            date_string = 't' + cb.timestamp_label.isoformat()
        else:
            date_string = 's' + datetime.datetime.now().isoformat()

        filename = gen_filename(poll_response.collection_name,
                                format_,
                                date_string,
                                ext)
        filename = os.path.join(dest_dir, filename)
        write, message = TaxiiScript.get_write_and_message(filename, write_type_)

        if write:
            with io.open(filename, 'wb') as f:
                f.write(cb.content)

        print("%s%s" % (message, filename))

    def write_cbs_from_poll_response_11(self, poll_response, dest_dir, write_type_=W_CLOBBER):
        """
        This function writes content blocks to file from a TAXII 1.1 Poll Response
        """

        for cb in poll_response.content_blocks:
            if cb.content_binding.binding_id == CB_STIX_XML_10:
                format_ = '_STIX10_'
                ext = '.xml'
            elif cb.content_binding.binding_id == CB_STIX_XML_101:
                format_ = '_STIX101_'
                ext = '.xml'
            elif cb.content_binding.binding_id == CB_STIX_XML_11:
                format_ = '_STIX11_'
                ext = '.xml'
            elif cb.content_binding.binding_id == CB_STIX_XML_111:
                format_ = '_STIX111_'
                ext = '.xml'
            elif cb.content_binding.binding_id == CB_STIX_XML_12:
                format_ = '_STIX12_'
                ext = '.xml'
            else:  # Format and extension are unknown
                format_ = ''
                ext = ''

            if cb.timestamp_label:
                date_string = 't' + cb.timestamp_label.isoformat()
            else:
                date_string = 's' + datetime.datetime.now().isoformat()

            filename = gen_filename(poll_response.collection_name,
                                    format_,
                                    date_string,
                                    ext)
            filename = os.path.join(dest_dir, filename)
            write, message = TaxiiScript.get_write_and_message(filename, write_type_)

            if write:
                with io.open(filename, 'wb') as f:
                    f.write(cb.content)

            print("%s%s" % (message, filename))

    def create_request_message(self, args):
        """
        This function should create a request message.
        Should be implemented by child classes.
        """
        raise NotImplementedError

    @staticmethod
    def _parse_url_info(args, parser):
        """
        In libtaxii 1.1.106, the command line args --port, --path, --host, and --https were deprecated in favor of
        --url. This function performs the following actions:

        1. Print a warning if any deprecated arguments are used (i.e., not one of the defaults).
        2. Raise a ValueError if --url is used in combination with any deprecated argument
        3. If deprecated arguments were used, translate them into the --url argument
        """

        deprecated_args = ['port', 'path', 'host', 'https']
        deprecated_arg_used = False
        for arg in deprecated_args:
            if str(getattr(args, arg)) != str(parser._optionals.get_default(arg)):
                print("WARNING: use of deprecated argument: %s. Use --url instead!" % arg)
                deprecated_arg_used = True

        url_used = getattr(args, 'url') != parser._optionals.get_default('url')
        if url_used and deprecated_arg_used:
            raise ValueError("Error: Both --url and a deprecated argument was used. Please use one or the other!")

        if deprecated_arg_used:
            url = "https" if args.https else "http" + "://" + args.host + ":" + str(args.port) + args.path
            args.url = url
            print("Deprecated arguments transformed into: '--url %s'" % url)

    def __call__(self):
        """
        Invoke a TAXII Service based on the arguments
        """
        try:
            parser = self.get_arg_parser(parser_description=self.parser_description, path=self.path)
            args = parser.parse_args()
            TaxiiScript._parse_url_info(args, parser)
            request_message = self.create_request_message(args)
            url = urlparse(args.url)
            client = self.create_client(url.scheme == 'https',
                                        args.proxy,
                                        args.cert,
                                        args.key,
                                        args.username,
                                        args.password)

            print("Request:\n")
            if args.xml_output is False:
                print(request_message.to_text())
            else:
                print(request_message.to_xml(pretty_print=True))

            resp = client.call_taxii_service2(url.hostname,
                                              url.path,
                                              self.taxii_version,
                                              request_message.to_xml(pretty_print=True),
                                              url.port)
            r = t.get_message_from_http_response(resp, '0')

            self.handle_response(r, args)
        except Exception as ex:
            traceback.print_exc()
            sys.exit(EXIT_FAILURE)

        sys.exit(EXIT_SUCCESS)

# TODO: These are stubs that will eventually need to be moved out into their own files / scripts


class SubscriptionClient11Script(TaxiiScript):
    parser_description = 'TAXII 1.1 Subscription Management Client'
    path = '/services/collection-management/'


class SubscriptionClient10Script(TaxiiScript):
    taxii_version = VID_TAXII_XML_10
    parser_description = 'TAXII 1.0 Subscription Management Client'
    path = '/services/feed-management/'


# No poll fulfillment in TAXII 1.0
