package dkms

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

	kms "cloud.google.com/go/kms/apiv1"
	"cloud.google.com/go/kms/apiv1/kmspb"
	"github.com/golang/glog"
	"google.golang.org/protobuf/types/known/wrapperspb"
)

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
}

// ModuleSigner is an interface for signing kernel modules.
type ModuleSigner interface {
	// Sign generates a digital signature for the given module bytes.
	Sign(moduleBytes []byte, hash crypto.Hash) ([]byte, error)
	// retrieveDigestEncryptionAlgorithm returns the algorithm identifier used for
	// encrypting the digest (e.g., RSA).
	retrieveDigestEncryptionAlgorithm() (*pkix.AlgorithmIdentifier, error)
}

// KmsSigner implements the ModuleSigner interface using Google Cloud KMS for signing operations.
type KmsSigner struct {
	ctx     context.Context
	client  *kms.KeyManagementClient
	keyName string
}

// LocalSigner implements the ModuleSigner interface using a local private key for signing.
type LocalSigner struct {
	key crypto.PrivateKey
}

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,
		},
	}
)

func newKmsSigner(ctx context.Context, kmsKey string) (*KmsSigner, error) {
	client, err := kms.NewKeyManagementClient(ctx)
	if err != nil {
		return nil, err
	}
	return &KmsSigner{ctx: ctx, client: client, keyName: kmsKey}, nil
}

func newLocalSigner(keyPath string) (*LocalSigner, error) {
	keyBytes, err := os.ReadFile(keyPath)
	if err != nil {
		return nil, fmt.Errorf("failed to read private key file (%s): %v", keyPath, err)
	}
	key, err := parsePrivateKey(keyBytes)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve private key: %v", err)
	}
	return &LocalSigner{key: key}, nil

}

func (l *LocalSigner) Sign(moduleBytes []byte, hash crypto.Hash) ([]byte, error) {
	// Retrieve the crypto signer if it is implemented for the private key
	signer, ok := l.key.(crypto.Signer)
	if !ok {
		return nil, fmt.Errorf("could not retrieve crypto signer for key type: %T", l.key)
	}
	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 (l *LocalSigner) retrieveDigestEncryptionAlgorithm() (*pkix.AlgorithmIdentifier, error) {
	switch l.key.(type) {
	case *rsa.PrivateKey:
		return &pkix.AlgorithmIdentifier{Algorithm: oidRsaEncryption, Parameters: asn1.NullRawValue}, nil
	default:
		return nil, fmt.Errorf("unsupported digest encryption algorithm: %T", l.key)
	}
}

func (k *KmsSigner) Sign(moduleBytes []byte, hash crypto.Hash) ([]byte, error) {
	h := hash.New()
	h.Write(moduleBytes)
	moduleDigest := h.Sum(nil)
	kmsDigest, err := k.retrieveDigestAlg(hash, moduleDigest)
	if err != nil {
		return nil, fmt.Errorf("failed to retrieve KMS digest algorithm: %v", err)
	}

	// Compute digest's CRC32C
	t := crc32.MakeTable(crc32.Castagnoli)
	digestCRC32C := crc32.Checksum(moduleDigest, t)

	req := &kmspb.AsymmetricSignRequest{
		Name:         k.keyName,
		Digest:       kmsDigest,
		DigestCrc32C: wrapperspb.Int64(int64(digestCRC32C)),
	}

	result, err := k.client.AsymmetricSign(k.ctx, req)
	if err != nil {
		return nil, fmt.Errorf("failed to generate signature using KMS: %v", err)
	}

	// Verify integrity of the result
	if !result.VerifiedDigestCrc32C {
		return nil, fmt.Errorf("AsymmetricSignError: digest corrupted in transit")
	}
	if result.Name != req.Name {
		return nil, fmt.Errorf("AsymmetricSignError: mismatched key name in the response")
	}
	if int64(crc32.Checksum(result.Signature, t)) != result.SignatureCrc32C.Value {
		return nil, fmt.Errorf("AsymmetricSignError: response corrupted in transit")
	}
	return result.Signature, nil
}

func (k *KmsSigner) retrieveDigestEncryptionAlgorithm() (*pkix.AlgorithmIdentifier, error) {
	req := &kmspb.GetCryptoKeyVersionRequest{
		Name: k.keyName,
	}
	keyVersion, err := k.client.GetCryptoKeyVersion(k.ctx, req)
	if err != nil {
		return nil, fmt.Errorf("failed to get key version: %v", err)
	}
	// Verify that key encryption is supported
	encryptAlg := keyVersion.GetAlgorithm().String()
	if strings.HasPrefix(encryptAlg, "RSA_") {
		return &pkix.AlgorithmIdentifier{Algorithm: oidRsaEncryption, Parameters: asn1.NullRawValue}, nil
	}
	return nil, fmt.Errorf("unsupported digest encryption algorithm: %v", encryptAlg)
}

func (k *KmsSigner) retrieveDigestAlg(hash crypto.Hash, moduleDigest []byte) (*kmspb.Digest, error) {
	var digest *kmspb.Digest
	switch hash {
	case crypto.SHA256:
		digest = &kmspb.Digest{
			Digest: &kmspb.Digest_Sha256{
				Sha256: moduleDigest,
			},
		}
	case crypto.SHA384:
		digest = &kmspb.Digest{
			Digest: &kmspb.Digest_Sha384{
				Sha384: moduleDigest,
			},
		}
	case crypto.SHA512:
		digest = &kmspb.Digest{
			Digest: &kmspb.Digest_Sha512{
				Sha512: moduleDigest,
			},
		}
	default:
		return nil, fmt.Errorf("unsupported kms digest algorithm: %s", hash.String())
	}
	return digest, nil
}

// SignModules signs a list of compiled kernel modules using the provided ModuleSigner.
// Currently, only rsa private keys are supported. It creates a pkcs7 signature
// and appends it to the end of a compiled kernel module.
//
// The ModuleSigner handles the specific signing method (e.g., local key or Cloud KMS)
// This is an implementation of https://github.com/torvalds/linux/blob/master/scripts/sign-file.c
func SignModules(modules []Module, certificatePath, hash string, signer ModuleSigner) error {
	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)
	}

	encryptAlg, err := signer.retrieveDigestEncryptionAlgorithm()
	if err != nil {
		return fmt.Errorf("failed to retrieve digest encryption algorithm: %v", err)
	}

	var signedModules, skippedModules []string
	for _, module := range modules {
		moduleName := module.BuiltName
		modulePath := module.BuiltPath()
		glog.Infof("signing module %s", moduleName)
		moduleBytes, err := os.ReadFile(modulePath)
		if err != nil {
			skippedModules = append(skippedModules, moduleName)
			glog.Errorf("failed to read module file (%s) [skipping]: %v", moduleName, err)
			continue
		}
		signedBytes, err := signModule(signer, moduleBytes, digest, cert, encryptAlg)
		if err != nil {
			skippedModules = append(skippedModules, moduleName)
			glog.Errorf("failed to sign module (%s) [skipping]: %v", moduleName, err)
			continue
		}
		if err := os.WriteFile(modulePath, signedBytes, 0644); err != nil {
			skippedModules = append(skippedModules, moduleName)
			glog.Errorf("failed to write signed module (%s) to file (%s) [skipping]: %v", moduleName, modulePath, err)
			continue
		}
		signedModules = append(signedModules, moduleName)
		glog.Infof("successfully signed module '%s'", moduleName)
	}

	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) {
	// Retrieve pem encoded 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(signer ModuleSigner, moduleBytes []byte, digest digestAlgorithm, cert *x509.Certificate, encryptAlg *pkix.AlgorithmIdentifier) ([]byte, error) {
	signature, err := signer.Sign(moduleBytes, digest.hash)
	if err != nil {
		return nil, fmt.Errorf("failed to generate module signature: %v", err)
	}

	digestAlg := &pkix.AlgorithmIdentifier{Algorithm: digest.identifier, Parameters: asn1.NullRawValue}
	sinfo := signerInfo{
		Version: 1,
		IssuerAndSerialNumber: issuerAndSerialNumber{
			Issuer:       asn1.RawValue{FullBytes: cert.RawIssuer},
			SerialNumber: cert.SerialNumber,
		},
		DigestAlgorithm:           *digestAlg,
		DigestEncryptionAlgorithm: *encryptAlg,
		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 nil, fmt.Errorf("failed to marshal signed module pkcs7 blob: %v", err)
	}

	return appendSignature(marshalled, moduleBytes)
}

// 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
}