// pkcs7: RFC2315 PKCS#7 CMS SignedData functions
// author: Stonesea, openck@aliyun.com
//
// Copyright (C) 2023 CyberKylin
// https://github.com/nfedera/rdpsign
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//    http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package pkcs7

import (
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha1"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/pem"
	"errors"
	"io/ioutil"
	"math/big"
)

type SignedDataObject struct {
	ID         asn1.ObjectIdentifier //signedData oid: 1.2.840.113549.1.7.2
	SignedData SignedData            `asn1:"explicit,tag:0"`
}

/*
RFC 2315: 9.1 SignedData type
The signed-data content type shall have ASN.1 type SignedData:

    SignedData ::= SEQUENCE {
     version Version,
     digestAlgorithms DigestAlgorithmIdentifiers,
     contentInfo ContentInfo,
     certificates
        [0] IMPLICIT ExtendedCertificatesAndCertificates
          OPTIONAL,
     crls
       [1] IMPLICIT CertificateRevocationLists OPTIONAL,
     signerInfos SignerInfos }
*/

// why the field Certificate is not Certificates(set) ?
// process Certificate in functions
// Ignore crls(CertificateRevocationLists), hsh 2017.11.15
type SignedData struct {
	Version          version
	DigestAlgorithms digestAlgorithmIdentifiers `asn1:"set"`
	ContentInfo      contentInfo                //data oid: 1.2.840.113549.1.7.1
	Certificate      asn1.RawValue              `asn1:"optional,explicit,tag:0"`
	SignerInfos      signerInfos                `asn1:"set"`
}

type version int
type digestAlgorithmIdentifiers []pkix.AlgorithmIdentifier

type AlgorithmIdentifier struct {
	ID asn1.ObjectIdentifier
}

/*
	    ContentInfo ::= SEQUENCE {
	     contentType ContentType,
	     content
	       [0] EXPLICIT ANY DEFINED BY contentType OPTIONAL }

		ContentType ::= OBJECT IDENTIFIER
*/
type contentInfo struct {
	ContentType asn1.ObjectIdentifier
	Content     asn1.RawValue `asn1:"optional,explicit,tag:0"`
}

/*
SignerInfos ::= SET OF SignerInfo

	SignerInfo ::= SEQUENCE {
	  version Version,
	  issuerAndSerialNumber IssuerAndSerialNumber,
	  digestAlgorithm DigestAlgorithmIdentifier,
	  authenticatedAttributes
	    [0] IMPLICIT Attributes OPTIONAL,
	  digestEncryptionAlgorithm
	    DigestEncryptionAlgorithmIdentifier,
	  encryptedDigest EncryptedDigest,
	  unauthenticatedAttributes
	    [1] IMPLICIT Attributes OPTIONAL }
*/
type signerInfos []signerInfo

type signerInfo struct {
	Version                   version
	IssuerAndSerialNumber     IssuerAndSerialNumber
	DigestAlgorithm           pkix.AlgorithmIdentifier
	DigestEncryptionAlgorithm pkix.AlgorithmIdentifier
	EncryptedDigest           encryptedDigest
}

type IssuerAndSerialNumber struct {
	Issuer       asn1.RawValue
	SerialNumber *big.Int
}

type encryptedDigest []byte

// Sign the data with SHA-1 hash algorithm as defined in RFC 3174.
func SignDataSha1(certFile, keyFile string, data []byte) (sd []byte, err error) {
	// parse cert
	b, err := ioutil.ReadFile(certFile)
	if err != nil {
		return
	}
	cblock, _ := pem.Decode(b)
	rsaPub, err := x509.ParseCertificate(cblock.Bytes)
	if err != nil {
		err = errors.New("parse certfile: " + err.Error())
		return
	}
	// parse key
	b, err = ioutil.ReadFile(keyFile)
	if err != nil {
		return
	}
	kblock, _ := pem.Decode(b)
	rsaPriv, err := x509.ParsePKCS1PrivateKey(kblock.Bytes)
	if err != nil {
		err = errors.New("parse keyfile: " + err.Error())
		return
	}

	hashed := sha1.Sum(data)
	// fmt.Println(hashed)
	// sign data
	signature, err := rsa.SignPKCS1v15(rand.Reader, rsaPriv, crypto.SHA1, hashed[:])

	// prepare SignedData struct
	digestAlgo := pkix.AlgorithmIdentifier{
		Algorithm:  asn1.ObjectIdentifier{1, 3, 14, 3, 2, 26}, //sha1
		Parameters: asn1.NullRawValue,
	}

	contentInfo := contentInfo{
		ContentType: asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 1}, //data
	}

	issuerSN := IssuerAndSerialNumber{
		Issuer:       asn1.RawValue{FullBytes: rsaPub.RawIssuer},
		SerialNumber: rsaPub.SerialNumber,
	}

	signer := signerInfo{
		Version:               2,
		IssuerAndSerialNumber: issuerSN,
		DigestAlgorithm: pkix.AlgorithmIdentifier{
			Algorithm:  asn1.ObjectIdentifier{1, 3, 14, 3, 2, 26},
			Parameters: asn1.NullRawValue},
		DigestEncryptionAlgorithm: pkix.AlgorithmIdentifier{
			Algorithm:  asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 1},
			Parameters: asn1.NullRawValue},
		EncryptedDigest: signature,
	}
	//SignerInfos
	signedData := SignedData{
		Version:          1,
		DigestAlgorithms: digestAlgorithmIdentifiers{digestAlgo},
		ContentInfo:      contentInfo,
		Certificate:      asn1.RawValue{Class: asn1.ClassContextSpecific, Tag: 0, IsCompound: true, Bytes: cblock.Bytes},
		SignerInfos:      signerInfos{signer},
	}
	sobj := SignedDataObject{
		ID:         asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 2},
		SignedData: signedData,
	}
	sd, err = asn1.Marshal(sobj)
	return
}

// Sign the data with SHA256 hash algorithms as defined in FIPS 180-4.
func SignDataSha256(certFile, keyFile string, data []byte) (sd []byte, err error) {
	// parse cert
	b, err := ioutil.ReadFile(certFile)
	if err != nil {
		return
	}
	cblock, _ := pem.Decode(b)
	rsaPub, err := x509.ParseCertificate(cblock.Bytes)
	if err != nil {
		err = errors.New("parse certfile: " + err.Error())
		return
	}
	// parse key
	b, err = ioutil.ReadFile(keyFile)
	if err != nil {
		return
	}
	kblock, _ := pem.Decode(b)
	rsaPriv, err := x509.ParsePKCS1PrivateKey(kblock.Bytes)
	if err != nil {
		err = errors.New("parse keyfile: " + err.Error())
		return
	}

	hashed := sha256.Sum256(data)
	// fmt.Println(hashed)
	// sign data
	signature, err := rsa.SignPKCS1v15(rand.Reader, rsaPriv, crypto.SHA256, hashed[:])

	// prepare SignedData struct
	digestAlgo := pkix.AlgorithmIdentifier{
		Algorithm:  asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 1}, //SHA256
		Parameters: asn1.NullRawValue,
	}

	contentInfo := contentInfo{
		ContentType: asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 1}, //data
	}

	issuerSN := IssuerAndSerialNumber{
		Issuer:       asn1.RawValue{FullBytes: rsaPub.RawIssuer},
		SerialNumber: rsaPub.SerialNumber,
	}

	signer := signerInfo{
		Version:               2,
		IssuerAndSerialNumber: issuerSN,
		DigestAlgorithm: pkix.AlgorithmIdentifier{
			Algorithm:  asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 1},
			Parameters: asn1.NullRawValue},
		DigestEncryptionAlgorithm: pkix.AlgorithmIdentifier{
			Algorithm:  asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 1},
			Parameters: asn1.NullRawValue},
		EncryptedDigest: signature,
	}
	//SignerInfos
	signedData := SignedData{
		Version:          1,
		DigestAlgorithms: digestAlgorithmIdentifiers{digestAlgo},
		ContentInfo:      contentInfo,
		Certificate:      asn1.RawValue{Class: asn1.ClassContextSpecific, Tag: 0, IsCompound: true, Bytes: cblock.Bytes},
		SignerInfos:      signerInfos{signer},
	}
	sobj := SignedDataObject{
		ID:         asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 2},
		SignedData: signedData,
	}
	sd, err = asn1.Marshal(sobj)
	return
}

// parse and verify ans.1 SignedDataObject
func parseSignedDataObject(derBytes []byte, hashAlgorithms crypto.Hash, hashed []byte) (sd SignedDataObject, err error) {
	var rest []byte
	if rest, err = asn1.Unmarshal(derBytes, &sd); err != nil {
		return
	} else if len(rest) != 0 {
		//fmt.Println(rest)
		err = errors.New("ParseRdpSign: trailing data after ASN.1 data")
	}

	// test cert
	if len(sd.SignedData.Certificate.FullBytes) > 0 {
		var cert *x509.Certificate
		cert, err = x509.ParseCertificate(sd.SignedData.Certificate.FullBytes)
		if err != nil {
			return
		}

		switch pub := cert.PublicKey.(type) {
		case *rsa.PublicKey:
			sig := sd.SignedData.SignerInfos[0].EncryptedDigest
			err = rsa.VerifyPKCS1v15(pub, hashAlgorithms, hashed, sig)
			if err != nil {
				return
			}
			//fmt.Println(plain)
		}

		if err != nil {
			return
		}
	}
	return
}

func ParseSignedDataObjectSha1(derBytes []byte, hashed []byte) (sd SignedDataObject, err error) {
	return parseSignedDataObject(derBytes, crypto.SHA1, hashed)
}

func ParseSignedDataObjectSha256(derBytes []byte, hashed []byte) (sd SignedDataObject, err error) {
	return parseSignedDataObject(derBytes, crypto.SHA256, hashed)
}