# *-* coding: utf-8 *-*
from __future__ import unicode_literals

import sys
import hashlib
import time
from datetime import datetime

import requests
import pytz
from asn1crypto import cms, algos, core, keys, pem, tsp, x509, util
from oscrypto import asymmetric
from cryptography.hazmat import backends
from cryptography.hazmat.primitives import hashes, serialization
from cryptography.hazmat.primitives.asymmetric import padding, utils


def cert2asn(cert, cert_bytes=True):
    if isinstance(cert, x509.Certificate):
        return cert
    if cert_bytes:
        cert_bytes = cert.public_bytes(serialization.Encoding.PEM)
    else:
        cert_bytes = cert
    if pem.detect(cert_bytes):
        _, _, cert_bytes = pem.unarmor(cert_bytes)
    return x509.Certificate.load(cert_bytes)

def sign(datau, key, cert, othercerts, hashalgo, attrs=True, signed_value=None, hsm=None, pss=False, timestampurl=None):
    if signed_value is None:
        signed_value = getattr(hashlib, hashalgo)(datau).digest()
    signed_time = datetime.now(tz=util.timezone.utc)

    if hsm is not None:
        keyid, cert = hsm.certificate()
        cert = cert2asn(cert, False)
    else:
        cert = cert2asn(cert)

    certificates = []
    certificates.append(cert)
    for i in range(len(othercerts)):
        certificates.append(cert2asn(othercerts[i]))

    hashalgo = unicode(hashalgo) if sys.version[0] < '3' else hashalgo

    signer = {
        'version': 'v1',
        'sid': cms.SignerIdentifier({
            'issuer_and_serial_number': cms.IssuerAndSerialNumber({
                'issuer': cert.issuer,
                'serial_number': cert.serial_number,
            }),
        }),
        'digest_algorithm': algos.DigestAlgorithm({'algorithm': hashalgo}),
        'signature': signed_value,
    }
    if not pss:
        signer['signature_algorithm'] = algos.SignedDigestAlgorithm({'algorithm': 'rsassa_pkcs1v15'})
    else:
        if isinstance(key, keys.PrivateKeyInfo):
            salt_length = key.byte_size - hashes.SHA512.digest_size - 2
            salt_length = hashes.SHA512.digest_size
        else:
            salt_length = padding.calculate_max_pss_salt_length(key, hashes.SHA512)
        signer['signature_algorithm'] = algos.SignedDigestAlgorithm({
            'algorithm': 'rsassa_pss',
            'parameters': algos.RSASSAPSSParams({
                'hash_algorithm': algos.DigestAlgorithm({'algorithm': 'sha512'}),
                'mask_gen_algorithm': algos.MaskGenAlgorithm({
                    'algorithm': algos.MaskGenAlgorithmId('mgf1'),
                    'parameters': {
                        'algorithm': algos.DigestAlgorithmId('sha512'),
                    }
                }),
                'salt_length': algos.Integer(salt_length),
                'trailer_field': algos.TrailerField(1)
            })
        })

    if attrs:
        if attrs is True:
            signer['signed_attrs'] = [
                cms.CMSAttribute({
                    'type': cms.CMSAttributeType('content_type'),
                    'values': ('data',),
                }),
                cms.CMSAttribute({
                    'type': cms.CMSAttributeType('message_digest'),
                    'values': (signed_value,),
                }),
                cms.CMSAttribute({
                    'type': cms.CMSAttributeType('signing_time'),
                    'values': (cms.Time({'utc_time': core.UTCTime(signed_time)}),)
                }),
            ]
        else:
            signer['signed_attrs'] = attrs


    config = {
        'version': 'v1',
        'digest_algorithms': cms.DigestAlgorithms((
            algos.DigestAlgorithm({'algorithm': hashalgo}),
        )),
        'encap_content_info': {
            'content_type': 'data',
        },
        'certificates': certificates,
        # 'crls': [],
        'signer_infos': [
            signer,
        ],
    }
    datas = cms.ContentInfo({
        'content_type': cms.ContentType('signed_data'),
        'content': cms.SignedData(config),
    })
    if attrs:
        tosign = datas['content']['signer_infos'][0]['signed_attrs'].dump()
        tosign = b'\x31' + tosign[1:]
    else:
        tosign = datau
    if hsm is not None:
        signed_value_signature = hsm.sign(keyid, tosign, hashalgo)
    elif isinstance(key, keys.PrivateKeyInfo):
        key = asymmetric.load_private_key(key)
        if pss:
            signed_value_signature = asymmetric.rsa_pss_sign(
                key,
                tosign,
                'sha512'
            )
        else:
            signed_value_signature = asymmetric.rsa_pkcs1v15_sign(
                key,
                tosign,
                hashalgo.lower()
            )
    else:
        if pss:
            hasher = hashes.Hash(hashes.SHA512(), backend=backends.default_backend())
            hasher.update(tosign)
            digest = hasher.finalize()
            signed_value_signature = key.sign(
                digest,
                padding.PSS(
                    mgf=padding.MGF1(hashes.SHA512()),
                    salt_length=salt_length
                ),
                utils.Prehashed(hashes.SHA512())
            )
        else:
            signed_value_signature = key.sign(
                tosign,
                padding.PKCS1v15(),
                getattr(hashes, hashalgo.upper())()
            )

    if timestampurl is not None:
        signed_value = getattr(hashlib, hashalgo)(signed_value_signature).digest()
        tspreq = tsp.TimeStampReq({
            "version": 1,
            "message_imprint": tsp.MessageImprint({
                "hash_algorithm": algos.DigestAlgorithm({'algorithm': hashalgo}),
                "hashed_message": signed_value,
            }),
            #'req_policy', ObjectIdentifier, {'optional': True}),
            "nonce": int(time.time()*1000),
            "cert_req": True,
            #'extensions': tsp.Extensions()
        })
        tspreq = tspreq.dump()

        tspheaders = {"Content-Type": "application/timestamp-query"}
        tspresp = requests.post(timestampurl, data=tspreq, headers=tspheaders)
        if tspresp.headers.get('Content-Type', None) == 'application/timestamp-reply':
            tspresp = tsp.TimeStampResp.load(tspresp.content)

            if tspresp['status']['status'].native == 'granted':
                attrs = [
                    cms.CMSAttribute({
                        'type': cms.CMSAttributeType('signature_time_stamp_token'),
                        'values': cms.SetOfContentInfo([
                            cms.ContentInfo({
                                'content_type': cms.ContentType('signed_data'),
                                'content': tspresp["time_stamp_token"]["content"],
                            })
                        ])
                    })
                ]
                datas['content']['signer_infos'][0]['unsigned_attrs'] = attrs

    # signed_value_signature = core.OctetString(signed_value_signature)
    datas['content']['signer_infos'][0]['signature'] = signed_value_signature

    #open('signed-content-info', 'wb').write(datas.dump())
    return datas.dump()
