package dkms

import (
	"bytes"
	"crypto"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/asn1"
	"encoding/binary"
	"encoding/pem"
	"fmt"
	"math/big"
	"os"
	"strings"

	"github.com/golang/glog"
)

const (
	// PKEYIDPKCS7 is the constant PKEY_ID_PKCS7 defined in https://github.com/torvalds/linux/blob/master/scripts/sign-file.c
	PKEYIDPKCS7 = byte(2)
	// magicNumber is the constant magic_number defined in https://github.com/torvalds/linux/blob/master/scripts/sign-file.c
	magicNumber = "~Module signature appended~\n"
)

// https://www.rfc-editor.org/rfc/rfc2315#section-9.1
//  SignedData ::= SEQUENCE {
//  version Version,
//  digestAlgorithms DigestAlgorithmIdentifiers,
//  contentInfo ContentInfo,
//  certificates
//     [0] IMPLICIT ExtendedCertificatesAndCertificates
//       OPTIONAL,
//  crls
//    [1] IMPLICIT CertificateRevocationLists OPTIONAL,
//  signerInfos SignerInfos }

type signedData struct {
	Version                    int
	DigestAlgorithmIdentifiers []pkix.AlgorithmIdentifier `asn1:"set"`
	ContentInfo                contentInfo
	SignerInfos                []signerInfo `asn1:"set"`
}

// https://www.rfc-editor.org/rfc/rfc2315#section-7
// ContentInfo ::= SEQUENCE {
//
//	 contentType ContentType,
//	 content  // NOTE: not needed for detached signature
//	 [0] EXPLICIT ANY DEFINED BY contentType OPTIONAL }
//
//	ContentType ::= OBJECT IDENTIFIER
type contentInfo struct {
	ContentType asn1.ObjectIdentifier
}

// https://www.rfc-editor.org/rfc/rfc2315#section-9.2
// SignerInfo ::= SEQUENCE {
//
//	version Version,
//	issuerAndSerialNumber IssuerAndSerialNumber,
//	digestAlgorithm DigestAlgorithmIdentifier,
//	authenticatedAttributes  // NOTE: not used in Linux kernel module signature
//	  [0] IMPLICIT Attributes OPTIONAL,
//	digestEncryptionAlgorithm
//	  DigestEncryptionAlgorithmIdentifier,
//	encryptedDigest EncryptedDigest,
//	unauthenticatedAttributes // NOTE: not used in Linux kernel module signature
//	  [1] IMPLICIT Attributes OPTIONAL }
//
// EncryptedDigest ::= OCTET STRING
type signerInfo struct {
	Version                   int
	IssuerAndSerialNumber     issuerAndSerialNumber
	DigestAlgorithm           pkix.AlgorithmIdentifier
	DigestEncryptionAlgorithm pkix.AlgorithmIdentifier
	EncryptedDigest           []byte
}

// https://www.rfc-editor.org/rfc/rfc2315#section-6.7
// IssuerAndSerialNumber ::= SEQUENCE {
//
//	issuer Name,
//	serialNumber CertificateSerialNumber }
type issuerAndSerialNumber struct {
	Issuer       asn1.RawValue
	SerialNumber *big.Int
}

type pkcs7Blob struct {
	Oid        asn1.ObjectIdentifier
	SignedData signedData `asn1:"tag:0,explicit"`
}

type digestAlgorithm struct {
	identifier asn1.ObjectIdentifier
	hash       crypto.Hash
}

var (
	oidData          = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 1}
	oidRsaEncryption = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 1, 1}
	oidSignedData    = asn1.ObjectIdentifier{1, 2, 840, 113549, 1, 7, 2}

	digests = map[string]digestAlgorithm{
		"sha256": {
			identifier: asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 1},
			hash:       crypto.SHA256,
		},
		"sha384": {
			identifier: asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 2},
			hash:       crypto.SHA384,
		},
		"sha512": {
			identifier: asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 3},
			hash:       crypto.SHA512,
		},
		"sha3-256": {
			identifier: asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 8},
			hash:       crypto.SHA3_256,
		},
		"sha3-384": {
			identifier: asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 9},
			hash:       crypto.SHA3_384,
		},
		"sha3-512": {
			identifier: asn1.ObjectIdentifier{2, 16, 840, 1, 101, 3, 4, 2, 10},
			hash:       crypto.SHA3_512,
		},
	}
)

// SignModules signs compiled kernel modules with a private key and certificate.
//
// Currently, only rsa private keys are supported. It creates a pksc7 signature
// and appends it to the end of a compiled kernel module.
// This is an implementation of https://github.com/torvalds/linux/blob/master/scripts/sign-file.c
func SignModules(modules []Module, keyPath, certificatePath, hash string) error {
	keyBytes, err := os.ReadFile(keyPath)
	if err != nil {
		return fmt.Errorf("failed to read private key file (%s): %v", keyPath, err)
	}
	key, err := parsePrivateKey(keyBytes)
	if err != nil {
		return fmt.Errorf("failed to retrieve private key: %v", err)
	}

	certBytes, err := os.ReadFile(certificatePath)
	if err != nil {
		return fmt.Errorf("failed to read certificate key file (%s): %v", certificatePath, err)
	}
	cert, err := x509.ParseCertificate(certBytes)
	if err != nil {
		return fmt.Errorf("failed to retrieve certificate (%s): %v", certificatePath, err)
	}
	digest, ok := digests[strings.ToLower(hash)]
	if !ok {
		return fmt.Errorf("unsupported digest algorithm: %s", hash)
	}

	var signedModules, skippedModules []string
	for _, module := range modules {
		glog.Infof("signing module %s", module.BuiltName)
		if err := signModule(module, key, cert, digest); err != nil {
			skippedModules = append(skippedModules, module.BuiltName)
			glog.Errorf("failed to sign module (%s) [skipping]: %v", module.BuiltName, err)
			continue
		}
		signedModules = append(signedModules, module.BuiltName)
		glog.Infof("successfully signed module '%s'", module.BuiltName)
	}

	glog.Infof("\nSummary:")
	if len(signedModules) > 0 {
		glog.Infof("Successfully signed module(s): %v", signedModules)
	}
	if len(skippedModules) > 0 {
		glog.Errorf("Failed to sign module(s): %v", skippedModules)
	}

	return nil
}

func parsePrivateKey(contents []byte) (crypto.PrivateKey, error) {
	// Retrive pem encodded block.
	block, _ := pem.Decode(contents)
	if block == nil {
		return nil, fmt.Errorf("failed to decode pem-formatted key or invalid key type")
	}

	var privateKey crypto.PrivateKey
	var err error
	switch block.Type {
	case "RSA PRIVATE KEY":
		privateKey, err = x509.ParsePKCS1PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("failed to parse rsa private key: %v", err)
		}
	case "PRIVATE KEY":
		key, err := x509.ParsePKCS8PrivateKey(block.Bytes)
		if err != nil {
			return nil, fmt.Errorf("failed to parse private key: %v", err)
		}
		switch k := key.(type) {
		case *rsa.PrivateKey:
			privateKey = k
		default:
			return nil, fmt.Errorf("unsupported key type: %T", key)
		}
	default:
		return nil, fmt.Errorf("unsupported block type: %s", block.Type)
	}
	return privateKey, nil
}

func signModule(module Module, key crypto.PrivateKey, certificate *x509.Certificate, digest digestAlgorithm) error {
	moduleName := module.BuiltName
	modulePath := module.BuiltPath()

	moduleBytes, err := os.ReadFile(modulePath)
	if err != nil {
		return fmt.Errorf("failed to read module file (%s): %v", modulePath, err)
	}
	signature, err := generateModuleSignature(key, moduleBytes, digest.hash)
	if err != nil {
		return fmt.Errorf("failed to generate signature for module (%s): %v", moduleName, err)
	}

	digestAlg := &pkix.AlgorithmIdentifier{Algorithm: digest.identifier, Parameters: asn1.NullRawValue}
	encrypAlg, err := retrieveDigestEncryptionAlgo(key)
	if err != nil {
		return fmt.Errorf("failed to retrieve digest encryption algorithm for module (%s): %v", moduleName, err)
	}
	sinfo := signerInfo{
		Version: 1,
		IssuerAndSerialNumber: issuerAndSerialNumber{
			Issuer:       asn1.RawValue{FullBytes: certificate.RawIssuer},
			SerialNumber: certificate.SerialNumber,
		},
		DigestAlgorithm:           *digestAlg,
		DigestEncryptionAlgorithm: *encrypAlg,
		EncryptedDigest:           signature,
	}

	blob := pkcs7Blob{
		Oid: oidSignedData,
		SignedData: signedData{
			Version:                    1,
			DigestAlgorithmIdentifiers: []pkix.AlgorithmIdentifier{*digestAlg},
			ContentInfo:                contentInfo{ContentType: oidData},
			SignerInfos:                []signerInfo{sinfo},
		},
	}
	marshalled, err := asn1.Marshal(blob)
	if err != nil {
		return fmt.Errorf("failed to marshal pkcs7 blob for module (%s): %v", moduleName, err)
	}

	signedBytes, err := appendSignature(marshalled, moduleBytes)
	if err != nil {
		return fmt.Errorf("failed to append signature to module (%s): %v", moduleName, err)
	}

	// Write the buffer content to the output file, overwriting it.
	if err := os.WriteFile(modulePath, signedBytes, 0644); err != nil {
		return fmt.Errorf("failed to write signed module to %s: %v", modulePath, err)
	}
	return nil
}

func generateModuleSignature(privateKey crypto.PrivateKey, moduleBytes []byte, hash crypto.Hash) ([]byte, error) {
	// Retrieve the crypto signer if it is implemented for the private key
	signer, ok := privateKey.(crypto.Signer)
	if !ok {
		return nil, fmt.Errorf("could not retrieve crypto signer for key type: %T", privateKey)
	}
	h := hash.New()
	h.Write(moduleBytes)
	moduleDigest := h.Sum(nil)
	signature, err := signer.Sign(rand.Reader, moduleDigest, hash)
	if err != nil {
		return nil, fmt.Errorf("failed to sign the module digest: %v", err)
	}
	return signature, nil
}

func retrieveDigestEncryptionAlgo(key crypto.PrivateKey) (*pkix.AlgorithmIdentifier, error) {
	switch key.(type) {
	case *rsa.PrivateKey:
		return &pkix.AlgorithmIdentifier{Algorithm: oidRsaEncryption, Parameters: asn1.NullRawValue}, nil
	default:
		return nil, fmt.Errorf("unsupported digest encryption algorithm: %T", key)
	}
}

// appendSignature appends a raw PKCS#7 signature to the end of a given kernel module.
func appendSignature(signature []byte, moduleBytes []byte) ([]byte, error) {
	var buf bytes.Buffer
	// Write bytes of kernel module.
	if _, err := buf.Write(moduleBytes); err != nil {
		return nil, fmt.Errorf("failed to write module bytes to buffer: %v", err)
	}
	// Copy bytes of module signature.
	sigSize, err := buf.Write(signature)
	if err != nil {
		return nil, fmt.Errorf("failed to write signature to buffer: %v", err)
	}
	// Append the marker and the PKCS#7 message.
	// signatureInfo is the struct module_signature defined in
	// https://github.com/torvalds/linux/blob/master/scripts/sign-file.c
	signatureInfo := [12]byte{}
	// signatureInfo[2] is the id_type of struct module_signature
	signatureInfo[2] = PKEYIDPKCS7
	// signatureInfo[8:12] is the sig_len of struct module_signature.
	// Using BigEndian as the sig_len should be in network byte order.
	binary.BigEndian.PutUint32(signatureInfo[8:12], uint32(sigSize))
	if _, err := buf.Write(signatureInfo[:]); err != nil {
		return nil, fmt.Errorf("failed to write module signature struct to buffer: %v", err)
	}

	if _, err := buf.WriteString(magicNumber); err != nil {
		return nil, fmt.Errorf("failed to write magic number to buffer: %v", err)
	}

	return buf.Bytes(), nil
}