#!/usr/bin/env python
#
# This Source Code Form is subject to the terms of the Mozilla Public
# License, v. 2.0. If a copy of the MPL was not distributed with this
# file, You can obtain one at http://mozilla.org/MPL/2.0/.

"""
Reads a specification from stdin and outputs a PKCS7 (CMS) message with
the desired properties.

The specification format is as follows:

sha1:<hex string>
sha256:<hex string>
md5:<hex string>
tamperDigest:sha1 - Only sha1 is supported
erase:{certificate, signerInfo}
signer:
<pycert specification>

Eith or both of sha1 and sha256 may be specified. The value of
each hash directive is what will be put in the messageDigest
attribute of the SignerInfo that corresponds to the signature
algorithm defined by the hash algorithm and key type of the
default key. Together, these comprise the signerInfos field of
the SignedData. If neither hash is specified, the signerInfos
will be an empty SET (i.e. there will be no actual signature
information).
The certificate specification must come last.
The script provides a possibility to tamper the hash while generating
SignedAttributes, such that the SignedAttributes signature will be incorrect
Erase allows specifying which PKCS7 field to strip (supports certificate or signerInfo)
"""

import base64
import sys
from enum import Enum
from io import StringIO

import pycert
import pykey
from pyasn1.codec.der import decoder, encoder
from pyasn1.type import tag, univ
from pyasn1_modules import rfc2315, rfc2459


class Error(Exception):
    """Base class for exceptions in this module."""

    pass


class UnknownDirectiveError(Error):
    """Helper exception type to handle unknown specification
    directives."""

    def __init__(self, directive):
        super(UnknownDirectiveError, self).__init__()
        self.directive = directive

    def __str__(self):
        return "Unknown directive %s" % repr(self.directive)


class FieldStrip(Enum):
    CERTIFICATE = "certificate"
    SIGNER_INFO = "signerInfo"


class HashToTamper(Enum):
    SHA1 = "sha1"


class CMS:
    """Utility class for reading a CMS specification and
    generating a CMS message"""

    def __init__(self, paramStream):
        self.sha1 = ""
        self.sha256 = ""
        self.md5 = ""

        self.fieldStrip = ""
        self.tamperDigest = ""

        signerSpecification = StringIO()
        readingSignerSpecification = False
        for line in paramStream.readlines():
            if readingSignerSpecification:
                print(line.strip(), file=signerSpecification)
            elif line.strip() == "signer:":
                readingSignerSpecification = True
            elif line.startswith("sha1:"):
                self.sha1 = line.strip()[len("sha1:") :]
            elif line.startswith("sha256:"):
                self.sha256 = line.strip()[len("sha256:") :]
            elif line.startswith("md5:"):
                self.md5 = line.strip()[len("md5:") :]
            elif line.startswith("erase:"):
                if line.strip()[len("erase:") :] == FieldStrip.CERTIFICATE.value:
                    self.fieldStrip = FieldStrip.CERTIFICATE
                elif line.strip()[len("erase:") :] == FieldStrip.SIGNER_INFO.value:
                    self.fieldStrip = FieldStrip.SIGNER_INFO
                else:
                    raise UnknownDirectiveError(line.strip())
            elif line.startswith("tamperDigest"):
                if line.strip()[len("tamperDigest:") :] == HashToTamper.SHA1.value:
                    self.tamperDigest = HashToTamper.SHA1
                else:
                    raise UnknownDirectiveError(line.strip())
            else:
                raise UnknownDirectiveError(line.strip())
        signerSpecification.seek(0)
        self.signer = pycert.Certificate(signerSpecification)
        self.signingKey = pykey.keyFromSpecification("default")

    def buildAuthenticatedAttributes(self, value, implicitTag=None):
        """Utility function to build a pyasn1 AuthenticatedAttributes
        object. Useful because when building a SignerInfo, the
        authenticatedAttributes needs to be tagged implicitly, but when
        signing an AuthenticatedAttributes, it needs the explicit SET
        tag."""
        if implicitTag:
            authenticatedAttributes = rfc2315.Attributes().subtype(
                implicitTag=implicitTag
            )
        else:
            authenticatedAttributes = rfc2315.Attributes()
        contentTypeAttribute = rfc2315.Attribute()
        # PKCS#9 contentType
        contentTypeAttribute["type"] = univ.ObjectIdentifier("1.2.840.113549.1.9.3")
        contentTypeAttribute["values"] = univ.SetOf(rfc2459.AttributeValue())
        # PKCS#7 data
        contentTypeAttribute["values"][0] = univ.ObjectIdentifier(
            "1.2.840.113549.1.7.1"
        )
        authenticatedAttributes[0] = contentTypeAttribute
        hashAttribute = rfc2315.Attribute()
        # PKCS#9 messageDigest
        hashAttribute["type"] = univ.ObjectIdentifier("1.2.840.113549.1.9.4")
        hashAttribute["values"] = univ.SetOf(rfc2459.AttributeValue())
        hashAttribute["values"][0] = univ.OctetString(hexValue=value)
        authenticatedAttributes[1] = hashAttribute
        return authenticatedAttributes

    def pykeyHashToDigestAlgorithm(self, pykeyHash):
        """Given a pykey hash algorithm identifier, builds an
        AlgorithmIdentifier for use with pyasn1."""
        if pykeyHash == pykey.HASH_SHA1:
            oidString = "1.3.14.3.2.26"
        elif pykeyHash == pykey.HASH_SHA256:
            oidString = "2.16.840.1.101.3.4.2.1"
        elif pykeyHash == pykey.HASH_MD5:
            oidString = "1.2.840.113549.2.5"
        else:
            raise pykey.UnknownHashAlgorithmError(pykeyHash)
        algorithmIdentifier = rfc2459.AlgorithmIdentifier()
        algorithmIdentifier["algorithm"] = univ.ObjectIdentifier(oidString)
        # Directly setting parameters to univ.Null doesn't currently work.
        nullEncapsulated = encoder.encode(univ.Null())
        algorithmIdentifier["parameters"] = univ.Any(nullEncapsulated)
        return algorithmIdentifier

    def buildSignerInfo(self, certificate, pykeyHash, digestValue):
        """Given a pyasn1 certificate, a pykey hash identifier
        and a hash value, creates a SignerInfo with the
        appropriate values."""
        signerInfo = rfc2315.SignerInfo()
        signerInfo["version"] = 1
        issuerAndSerialNumber = rfc2315.IssuerAndSerialNumber()
        issuerAndSerialNumber["issuer"] = self.signer.getIssuer()
        issuerAndSerialNumber["serialNumber"] = certificate["tbsCertificate"][
            "serialNumber"
        ]
        signerInfo["issuerAndSerialNumber"] = issuerAndSerialNumber
        signerInfo["digestAlgorithm"] = self.pykeyHashToDigestAlgorithm(pykeyHash)
        rsa = rfc2459.AlgorithmIdentifier()
        rsa["algorithm"] = rfc2459.rsaEncryption
        rsa["parameters"] = univ.Null()
        authenticatedAttributes = self.buildAuthenticatedAttributes(
            digestValue,
            implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0),
        )
        authenticatedAttributesTBS = self.buildAuthenticatedAttributes(digestValue)
        signerInfo["authenticatedAttributes"] = authenticatedAttributes
        signerInfo["digestEncryptionAlgorithm"] = rsa
        authenticatedAttributesEncoded = encoder.encode(authenticatedAttributesTBS)
        signature = self.signingKey.sign(authenticatedAttributesEncoded, pykeyHash)
        if self.tamperDigest == HashToTamper.SHA1 and pykeyHash == pykey.HASH_SHA1:
            digestValue = hex((int(digestValue[0], 16) + 1) % 16)[2:] + digestValue[1:]
            authenticatedAttributesTBSTamperedHash = self.buildAuthenticatedAttributes(
                digestValue
            )
            authenticatedAttributesTamperedEncoded = encoder.encode(
                authenticatedAttributesTBSTamperedHash
            )
            # The signerInfo has an attribute with the initial hash
            # But the tampered hash attributes are signed
            signature = self.signingKey.sign(
                authenticatedAttributesTamperedEncoded, pykeyHash
            )
        # signature will be a hexified bit string of the form
        # "'<hex bytes>'H". For some reason that's what BitString wants,
        # but since this is an OCTET STRING, we have to strip off the
        # quotation marks and trailing "H".
        signerInfo["encryptedDigest"] = univ.OctetString(hexValue=signature[1:-2])
        return signerInfo

    def toDER(self):
        contentInfo = rfc2315.ContentInfo()
        contentInfo["contentType"] = rfc2315.signedData

        signedData = rfc2315.SignedData()
        signedData["version"] = rfc2315.Version(1)

        digestAlgorithms = rfc2315.DigestAlgorithmIdentifiers()
        digestAlgorithms[0] = self.pykeyHashToDigestAlgorithm(pykey.HASH_SHA1)
        signedData["digestAlgorithms"] = digestAlgorithms

        dataContentInfo = rfc2315.ContentInfo()
        dataContentInfo["contentType"] = rfc2315.data
        signedData["contentInfo"] = dataContentInfo

        certificates = rfc2315.ExtendedCertificatesAndCertificates().subtype(
            implicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0)
        )
        extendedCertificateOrCertificate = rfc2315.ExtendedCertificateOrCertificate()
        certificate = decoder.decode(
            self.signer.toDER(), asn1Spec=rfc2459.Certificate()
        )[0]
        extendedCertificateOrCertificate["certificate"] = certificate
        certificates[0] = extendedCertificateOrCertificate

        if self.fieldStrip != FieldStrip.CERTIFICATE:
            signedData["certificates"] = certificates

        if self.fieldStrip != FieldStrip.SIGNER_INFO:
            signerInfos = rfc2315.SignerInfos()

            if len(self.sha1) > 0:
                signerInfos[len(signerInfos)] = self.buildSignerInfo(
                    certificate, pykey.HASH_SHA1, self.sha1
                )
            if len(self.sha256) > 0:
                signerInfos[len(signerInfos)] = self.buildSignerInfo(
                    certificate, pykey.HASH_SHA256, self.sha256
                )
            if len(self.md5) > 0:
                signerInfos[len(signerInfos)] = self.buildSignerInfo(
                    certificate, pykey.HASH_MD5, self.md5
                )
            signedData["signerInfos"] = signerInfos

        encoded = encoder.encode(signedData)
        anyTag = univ.Any(encoded).subtype(
            explicitTag=tag.Tag(tag.tagClassContext, tag.tagFormatConstructed, 0)
        )

        contentInfo["content"] = anyTag
        return encoder.encode(contentInfo)

    def toPEM(self):
        output = "-----BEGIN PKCS7-----"
        der = self.toDER()
        b64 = base64.b64encode(der)
        while b64:
            output += "\n" + b64[:64].decode("utf-8")
            b64 = b64[64:]
        output += "\n-----END PKCS7-----\n"
        return output


# The build harness will call this function with an output
# file-like object and a path to a file containing a
# specification. This will read the specification and output
# the cms message as PEM.
def main(output, inputPath):
    with open(inputPath) as configStream:
        output.write(CMS(configStream).toPEM() + "\n")


# When run as a standalone program, this will read a specification from
# stdin and output the cms message as PEM.
if __name__ == "__main__":
    print(CMS(sys.stdin).toPEM())