"""
    Client for the ARC data-delivery service.
"""

import argparse
import httplib
import logging
import os
import re
import six
import socket
import ssl
import sys
import time
import uuid
import xml.etree.ElementTree as ET

from cryptography import x509
from cryptography.hazmat.backends import default_backend
from cryptography.hazmat.primitives import serialization
from datetime import datetime, timedelta

logging.basicConfig(
    format='%(asctime)s - %(name)s - %(levelname)s - %(message)s')
LOGGER = logging.getLogger(__name__)

ARC_DELEGATION_NAMESPACE = "http://www.nordugrid.org/schemas/delegation"


class HTTPSClientAuthConnection(httplib.HTTPSConnection):
    """
    Class to make a HTTPS connection,
    with support for full client-based SSL Authentication
    """

    def __init__(self, host='127.0.0.1', port=3000, key_file='', cert_file='', ca_file='', timeout=None):
        httplib.HTTPSConnection.__init__(self, host, port, key_file=key_file, cert_file=cert_file)
        self.key_file = key_file
        self.cert_file = cert_file
        self.ca_file = ca_file
        self.port = port
        self.timeout = timeout

    def connect(self):
        """ Connect to a host on a given (SSL) port.
            If ca_file is pointing somewhere, use it to check Server Certificate.

            Redefined/copied and extended from httplib.py:1105 (Python 2.6.x).
            This is needed to pass cert_reqs=ssl.CERT_REQUIRED as parameter to ssl.wrap_socket(),
            which forces SSL to check server certificate against our client certificate.
        """
        sock = socket.create_connection((self.host, self.port), self.timeout)
        # If there's no CA File, don't force Server Certificate Check
        if self.ca_file:
            self.sock = ssl.wrap_socket(
                sock, self.key_file, self.cert_file,
                ca_certs=self.ca_file, cert_reqs=ssl.CERT_REQUIRED
                )
        else:
            self.sock = ssl.wrap_socket(
                sock, self.key_file, self.cert_file, cert_reqs=ssl.CERT_NONE
                )


class ARCException(Exception):
    """
    Exceptions generated by this code.
    """
    pass


class ARC:  # (TransferTool)

    def __init__(self, url, proxy=None, ca_file=None):
        # TransferTool.__init__()
        self.url = url
        self.protocol, self.host, self.port, self.path = self._split_url(url)
        self.ca_file = ca_file
        self.proxy = proxy
        with open(self.proxy, 'rb') as f:
            self.proxy_pem = f.read()

    def _split_url(self, url):
        """
        Split url into (protocol, host, port, path) and return this tuple.
        """
        match = re.match('(\w*)://([^/?#:]*):?(\d*)/?(.*)', url)
        if match is None:
            raise ARCException('URL '+url+' is malformed')
        port_s = match.group(3)
        if (len(port_s) > 0):
            port = int(port_s)
        else:
            port = None
        urltuple = (match.group(1), match.group(2), port, '/'+match.group(4))
        return urltuple

    def _add_ETElement(self, root, child, text):
        """
        Utility function to add a child element with text to the root element.
        Returns the child element.
        """
        sub = ET.SubElement(root, child)
        sub.text = text
        return sub

    def _check_SOAP_fault(self, element):
        """
        Checks the response from a service given in element for SOAP faults, and
        raises an exception if there is one.
        """
        response_body = element[0]
        if response_body.find('{http://schemas.xmlsoap.org/soap/envelope/}Fault') is not None:
            fault = response_body.find('{http://schemas.xmlsoap.org/soap/envelope/}Fault')
            faultcode = fault.find('{http://schemas.xmlsoap.org/soap/envelope/}faultcode').text
            faultstring = fault.find('{http://schemas.xmlsoap.org/soap/envelope/}faultstring').text
            raise ARCException('SOAP error: '+faultcode+' - '+faultstring)

    def _send_request(self, request):
        '''
        Send request.
        '''
        conn = HTTPSClientAuthConnection(
            self.host, self.port, key_file=self.proxy,
            cert_file=self.proxy, ca_file=self.ca_file
            )
        try:
            LOGGER.debug(request)
            conn.request('POST', self.path, request)
            resp = conn.getresponse()

            # On SOAP fault 500 is returned - this is caught in checkSOAPFault
            if resp.status != 200 and resp.status != 500:
                raise ARCException('Error code '+str(resp.status)+' returned: '+resp.reason)
            xmldata = resp.read()
            LOGGER.debug(xmldata)
            response = ET.XML(xmldata)
            self._check_SOAP_fault(response)
            return response
        except Exception, e:
            raise ARCException('Error connecting to service at ' + self.host + ': ' + str(e))
        finally:
            conn.close()

    def ping(self):
        """
        Ping the data-delivery service at the given endpoint.
        """
        # create request with etree
        soap = ET.Element('soap-env:Envelope', attrib={
            'xmlns:echo': 'urn:echo',
            'xmlns:soap-enc': 'http://schemas.xmlsoap.org/soap/encoding/',
            'xmlns:soap-env': 'http://schemas.xmlsoap.org/soap/envelope/',
            'xmlns:xsd': 'http://www.w3.org/2001/XMLSchema',
            'xmlns:xsi': 'http://www.w3.org/2001/XMLSchema-instance'})
        body = ET.SubElement(soap, 'soap-env:Body')
        ET.SubElement(body, 'DataDeliveryPing')
        request = ET.tostring(soap)

        response = self._send_request(request)
        try:
            pingresult = response.find('{http://schemas.xmlsoap.org/soap/envelope/}Body').\
                find('DataDeliveryPingResponse').\
                find('DataDeliveryPingResult').\
                find('Result')
        except:
            raise ARCException('Error with XML structure received from datadelivery service')

        result = {}
        if pingresult.find('ResultCode') is not None:
            result['ResultCode'] = pingresult.find('ResultCode').text

        if pingresult.find('LoadAvg') is not None:
            result['LoadAvg'] = pingresult.find('LoadAvg').text

        result['AllowedDirs'] = []
        for allowed_dir in pingresult.findall('AllowedDir'):
            result['AllowedDirs'].append(allowed_dir.text)

        return result

    def query(self, transfer_ids):
        """
        Query transfer statuses.
        """
        soap = ET.Element(
            'soap-env:Envelope',
            attrib={'xmlns:echo': 'urn:echo', 'xmlns:soap-enc': 'http://schemas.xmlsoap.org/soap/encoding/',
                    'xmlns:soap-env': 'http://schemas.xmlsoap.org/soap/envelope/',
                    'xmlns:xsd': 'http://www.w3.org/2001/XMLSchema',
                    'xmlns:xsi': 'http://www.w3.org/2001/XMLSchema-instance'})

        body = ET.SubElement(soap, 'soap-env:Body')
        datadeliveryquery = ET.SubElement(body, 'DataDeliveryQuery')
        for transfer_id in transfer_ids:
            dtr = ET.SubElement(datadeliveryquery, 'DTR')
            self._add_ETElement(dtr, 'ID', transfer_id)

        request = ET.tostring(soap)
        response = self._send_request(request)
        try:
            transfer_result = response.find('{http://schemas.xmlsoap.org/soap/envelope/}Body').\
                find('DataDeliveryQueryResponse').\
                find('DataDeliveryQueryResult')
            results = transfer_result.findall('Result')
        except:
            raise ARCException('Error with XML structure received from datadelivery service')

        transfer_statuses = {}
        for result in results:
            result_id = result.find('ID').text
            result_code = result.find('ResultCode').text

            transfer_statuses[result_id] = {'ResultCode': result_code}

            if (result.find('ErrorDescription') is not None
                    and result.find('ErrorDescription').text):
                transfer_statuses[result_id]['ErrorDescription'] = result.find('ErrorDescription').text

            if result.find('Log') is not None:
                transfer_statuses[result_id]['Log'] = result.find('Log').text

            if result.find('BytesTransferred') is not None:
                transfer_statuses[result_id]['BytesTransferred'] = result.find('BytesTransferred').text

        return transfer_statuses

    def cancel(self, transfer_ids):
        """
        Cancel transfer statuses.
        """
        soap = ET.Element(
            'soap-env:Envelope',
            attrib={'xmlns:echo': 'urn:echo', 'xmlns:soap-enc': 'http://schemas.xmlsoap.org/soap/encoding/',
                    'xmlns:soap-env': 'http://schemas.xmlsoap.org/soap/envelope/',
                    'xmlns:xsd': 'http://www.w3.org/2001/XMLSchema',
                    'xmlns:xsi': 'http://www.w3.org/2001/XMLSchema-instance'})

        body = ET.SubElement(soap, 'soap-env:Body')
        datadeliverycancel = ET.SubElement(body, 'DataDeliveryCancel')
        for transfer_id in transfer_ids:
            dtr = ET.SubElement(datadeliverycancel, 'DTR')
            self._add_ETElement(dtr, 'ID', transfer_id)

        request = ET.tostring(soap)
        response = self._send_request(request)
        try:
            cancel_result = response.find('{http://schemas.xmlsoap.org/soap/envelope/}Body').\
                find('DataDeliveryCancelResponse').\
                find('DataDeliveryCancelResult')
            results = cancel_result.findall('Result')
        except:
            raise ARCException('Error with XML structure received from datadelivery service')

        cancel_statuses = {}
        for result in results:
            cancel_statuses[result.find('ID').text] = result.find('ResultCode').text

        return cancel_statuses

    def _get_delegation(self):
        """
        Get deleguation id and proxy request.
        """
        soap = ET.Element(
            'soap-env:Envelope',
            attrib={'xmlns:deleg':  ARC_DELEGATION_NAMESPACE,
                    'xmlns:echo': 'urn:echo',
                    'xmlns:soap-enc': 'http://schemas.xmlsoap.org/soap/encoding/',
                    'xmlns:soap-env': 'http://schemas.xmlsoap.org/soap/envelope/',
                    'xmlns:xsd': 'http://www.w3.org/2001/XMLSchema',
                    'xmlns:xsi': 'http://www.w3.org/2001/XMLSchema-instance'})

        body = ET.SubElement(soap, 'soap-env:Body')
        ET.SubElement(body, 'deleg:DelegateCredentialsInit')
        request = ET.tostring(soap)

        response = self._send_request(request)
        result = response.find('{http://schemas.xmlsoap.org/soap/envelope/}Body').\
            find('{http://www.nordugrid.org/schemas/delegation}DelegateCredentialsInitResponse').\
            find('{http://www.nordugrid.org/schemas/delegation}TokenRequest')

        delegation_id = result.find('{http://www.nordugrid.org/schemas/delegation}Id').text
        x509_request = result.find('{http://www.nordugrid.org/schemas/delegation}Value').text

        return delegation_id, x509_request

    def _check_rfc_proxy(self, proxy):
        """
        Check for X509 RFC 3820 proxy.
        """
        for ext in proxy.extensions:
            if ext.oid.dotted_string == "1.3.6.1.5.5.7.1.14":
                return True
        raise ARCException('Invalid X509 RFC 3820 proxy.')

    def _sign_request(self, x509_request, lifetime=24):
        """
        Sign proxy.
        """
        now = datetime.utcnow()
        csr = x509.load_pem_x509_csr(x509_request, default_backend())
        if not csr.is_signature_valid:
            raise ARCException('Invalid request signature.')

        proxy = x509.load_pem_x509_certificate(self.proxy_pem, default_backend())
        self._check_rfc_proxy(proxy)
        key = serialization.load_pem_private_key(self.proxy_pem, password=None, backend=default_backend())
        key_id = x509.SubjectKeyIdentifier.from_public_key(key.public_key())
        subject_attributes = list(proxy.subject)
        subject_attributes.append(
            x509.NameAttribute(x509.oid.NameOID.COMMON_NAME, six.u(str(int(time.time())))))

        new_cert = x509.CertificateBuilder() \
            .issuer_name(proxy.subject) \
            .not_valid_before(now) \
            .not_valid_after(now + timedelta(hours=lifetime)) \
            .serial_number(proxy.serial_number) \
            .public_key(csr.public_key()) \
            .subject_name(x509.Name(subject_attributes)) \
            .add_extension(x509.BasicConstraints(ca=False, path_length=None),
                           critical=True) \
            .add_extension(x509.KeyUsage(digital_signature=True,
                                         content_commitment=False,
                                         key_encipherment=False,
                                         data_encipherment=False,
                                         key_agreement=True,
                                         key_cert_sign=False,
                                         crl_sign=False,
                                         encipher_only=False,
                                         decipher_only=False),
                           critical=True) \
            .add_extension(x509.AuthorityKeyIdentifier(
                            key_identifier=key_id.digest,
                            authority_cert_issuer=[x509.DirectoryName(proxy.issuer)],
                            authority_cert_serial_number=proxy.serial_number
                            ),
                           critical=False) \
            .add_extension(x509.extensions.UnrecognizedExtension(
                x509.ObjectIdentifier("1.3.6.1.5.5.7.1.14"),
                b"0\x0c0\n\x06\x08+\x06\x01\x05\x05\x07\x15\x01"),
            critical=True) \
            .sign(private_key=key,
                  algorithm=proxy.signature_hash_algorithm,
                  backend=default_backend())
        return new_cert.public_bytes(serialization.Encoding.PEM) + self.proxy_pem
        # proxy.public_bytes(serialization.Encoding.PEM)
#        self.proxy_pem

    def _update_delegation(self, delegation_id, proxy):
        """
        Update the proxy.
        """
        soap = ET.Element('soap-env:Envelope',
                          attrib={'xmlns:deleg': ARC_DELEGATION_NAMESPACE,
                                  'xmlns:echo': 'urn:echo',
                                  'xmlns:soap-enc': 'http://schemas.xmlsoap.org/soap/encoding/',
                                  'xmlns:soap-env': 'http://schemas.xmlsoap.org/soap/envelope/',
                                  'xmlns:xsd': 'http://www.w3.org/2001/XMLSchema',
                                  'xmlns:xsi': 'http://www.w3.org/2001/XMLSchema-instance'})
        body = ET.SubElement(soap, 'soap-env:Body')
        delegation = ET.SubElement(body, 'deleg:UpdateCredentials')
        token = ET.SubElement(delegation, 'deleg:DelegatedToken', attrib={'deleg:Format': 'x509'})
        self._add_ETElement(token, 'deleg:Id', delegation_id)
        self._add_ETElement(token, 'deleg:Value', proxy)
        response = self._send_request(ET.tostring(soap))
        return ET.tostring(response)

    def _submit_transfers(self, files, delegation_id, proxy):
        '''
        Submit transfers.
        '''
        soap = ET.Element(
            'soap-env:Envelope',
            attrib={'xmlns:echo': 'urn:echo',
                    'xmlns:soap-enc': 'http://schemas.xmlsoap.org/soap/encoding/',
                    'xmlns:soap-env': 'http://schemas.xmlsoap.org/soap/envelope/',
                    'xmlns:xsd': 'http://www.w3.org/2001/XMLSchema',
                    'xmlns:xsi': 'http://www.w3.org/2001/XMLSchema-instance',
                    'xmlns:deleg': ARC_DELEGATION_NAMESPACE})
        body = ET.SubElement(soap, 'soap-env:Body')
        datadeliverystart = ET.SubElement(body, 'DataDeliveryStart')

        # add delegation
        token = ET.SubElement(
            datadeliverystart,
            'deleg:DelegatedToken',
            attrib={'deleg:Format': 'x509'}
            )
        self._add_ETElement(token, 'deleg:Id', delegation_id)
        self._add_ETElement(token, 'deleg:Value', proxy)

        for fname, attrs in files.items():
            file = ET.SubElement(datadeliverystart, 'DTR')
            self._add_ETElement(file, 'ID', attrs['id'])
            self._add_ETElement(file, 'Source', attrs['src'])
            self._add_ETElement(file, 'Destination', attrs['dest'])
            # uid and gid are not used if caching is enabled
            self._add_ETElement(file, 'Uid', 0)
            self._add_ETElement(file, 'Gid', 0)
            self._add_ETElement(file, 'Caching', 'true')
            self._add_ETElement(file, 'Size', attrs['size'])
            self._add_ETElement(file, 'CheckSum', attrs['checksum'])
            # hard-coded timeouts, probably can be taken from job_params
            self._add_ETElement(file, 'MinAverageSpeed', '100')
            self._add_ETElement(file, 'AverageTime', '60')
            self._add_ETElement(file, 'MinCurrentSpeed', '100')
            self._add_ETElement(file, 'MaxInactivityTime', '120')

        response = self._send_request(ET.tostring(soap))
        try:
            transfer_result = response.\
                find('{http://schemas.xmlsoap.org/soap/envelope/}Body').\
                find('DataDeliveryStartResponse').\
                find('DataDeliveryStartResult')
            results = transfer_result.findall('Result')
        except:
            raise ARCException('Error with XML structure received from data delivery service')

        if len(results) == 0:
            raise ARCException('No results returned')

        transferids = []
        for result in results:
            result_code = result.find('ResultCode').text
            result_text = result.find('ErrorDescription') is not None and result.find('ErrorDescription').text
            if result_code == 'OK':
                transferid = result.find('ID').text
                transferids.append(transferid)
            else:
                # Return error
                print(result_text)
        return transferids

    def submit(self, files):
        """
        Submit transfers.

        Returns a list of transfer identifiers.
        """

        delegation_id, x509_request = self._get_delegation()

        proxy = self._sign_request(x509_request)

        self._update_delegation(delegation_id, proxy)

        transferids = self._submit_transfers(files, delegation_id, proxy)

        # with open('/tmp/proxy.pem', 'w') as fd:
        #    fd.write(proxy)


        return transferids


def get_parser():
    '''
    Get parser
    '''
    oparser = argparse.ArgumentParser(prog=os.path.basename(sys.argv[0]), add_help=True)
    oparser.add_argument('--debug', action='store_true', help='print debug messages to stderr.')
    subparsers = oparser.add_subparsers()

    # Ping command
    ping_parser = subparsers.add_parser('ping', help='Ping data delivery service.')
    ping_parser.set_defaults(which='ping')
    ping_parser.add_argument('--url', dest="url", help="The service url.",
                             default='https://localhost:8443/datadeliveryservice')
    ping_parser.add_argument(
        '--proxy', dest='proxy',
        default='/tmp/x509up_u' + str(os.getuid()),
        help='Client X509 proxy file.')
    ping_parser.add_argument(
        '--ca-certificate', dest='ca_certificate',
        default=None,
        help='CA certificate to verify peer against (SSL).'
        )

    # Submit command
    submit_parser = subparsers.add_parser('submit', help='Submit transfer.')
    submit_parser.set_defaults(which='submit')
    submit_parser.add_argument('--url', dest="url", help="The service url.",
                               default='https://localhost:8443/datadeliveryservice')
    submit_parser.add_argument('--proxy', dest='proxy',
                               default='/tmp/x509up_u' + str(os.getuid()),
                               help='Client X509 proxy file.')
    submit_parser.add_argument('--ca-certificate', dest='ca_certificate', default=None,
                               help='CA certificate to verify peer against (SSL).')
    submit_parser.add_argument('--source', dest='source', required=True)
    submit_parser.add_argument('--destination', dest='destination', required=True)
    submit_parser.add_argument('--size', dest='size', required=True)
    submit_parser.add_argument('--checksum', dest='checksum', required=True)

    # Query command
    query_parser = subparsers.add_parser('query', help='Query transfers.')
    query_parser.set_defaults(which='query')
    query_parser.add_argument('--url', dest="url", help="The service url.",
                              default='https://localhost:8443/datadeliveryservice')
    query_parser.add_argument(
        '--proxy', dest='proxy',
        default='/tmp/x509up_u' + str(os.getuid()),
        help='Client X509 proxy file.'
        )
    query_parser.add_argument(
        '--ca-certificate', dest='ca_certificate',
        default=None,
        help='CA certificate to verify peer against (SSL).'
        )
    query_parser.add_argument(dest='transfer_ids', nargs='+',
                              help="Transfer ids", action='store')

    # Cancel command
    cancel_parser = subparsers.add_parser('cancel', help='Cancel transfers.')
    cancel_parser.set_defaults(which='cancel')
    cancel_parser.add_argument('--url', dest="url", help="The service url.", required=True)
    cancel_parser.add_argument(
        '--proxy', dest='proxy',
        default='/tmp/x509up_u' + str(os.getuid()),
        help='Client X509 proxy file.'
        )
    cancel_parser.add_argument(
        '--ca-certificate', dest='ca_certificate',
        default=None,
        help='CA certificate to verify peer against (SSL).')
    cancel_parser.add_argument(dest='transfer_ids', nargs='+',
                               help="Transfer ids", action='store')
    return oparser


def main():
    ''' Main method '''
    oparser = get_parser()
    args = oparser.parse_args(sys.argv[1:])

    if args.debug:
        LOGGER.setLevel(logging.DEBUG)

    transfertool = ARC(args.url, proxy=args.proxy, ca_file=args.ca_certificate)

    if args.which == 'ping':
        print("Pinging {}".format(args.url))
        print(transfertool.ping())

    elif args.which == 'submit':
        files = {'test': {'id': str(uuid.uuid4()),
                          'src': args.source,
                          'dest': args.destination,
                          'size': args.size,
                          'checksum': args.checksum}}
        transfer_ids = transfertool.submit(files)
        for transfer_id in transfer_ids:
            print(transfer_id)

    elif args.which == 'query':
        transfer_statuses = transfertool.query(args.transfer_ids)
        for transfer_id in transfer_statuses:
            print('%s:' % transfer_id)
            for item in transfer_statuses[transfer_id]:
                print('\t %s: %s' % (item, str(transfer_statuses[transfer_id][item])))

    elif args.which == 'cancel':
        cancel_statuses = transfertool.cancel(args.transfer_ids)
        for transfer_id in cancel_statuses:
            print(transfer_id, cancel_statuses[transfer_id])

    sys.exit(0)
