/*
Copyright 2022 CAcert Inc.
SPDX-License-Identifier: Apache-2.0

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 hsm handles hardware security modules.
package hsm

import (
	"context"
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/asn1"
	"encoding/pem"
	"errors"
	"fmt"
	"log"
	"math/big"
	"os"
	"path"
	"syscall"
	"time"

	"github.com/ThalesIgnite/crypto11"

	"git.cacert.org/cacert-gosigner/pkg/config"
)

var (
	// 1.3.6.1.4.1.18506.2.3.1 Class3 Policy Version 1
	oidCAcertClass3PolicyV1 = []int{1, 3, 6, 1, 4, 1, 18506, 2, 3, 1}

	errKeyGenerationRefused         = errors.New("not in setup mode, refusing key generation")
	errCertificateGenerationRefused = errors.New("not in setup mode, refusing certificate generation")
)

type caFile struct {
	sc    *config.SignerConfig
	label string
}

func (c *caFile) buildCertificatePath(ctx context.Context) (string, error) {
	fileName := c.sc.CertificateFileName(c.label)

	caDir := ctx.Value(ctxCADirectory)

	if caDir != nil {
		caPath, ok := caDir.(string)
		if !ok {
			return "", errors.New("context object CA directory is not a string")
		}

		return path.Join(caPath, fileName), nil
	}

	return fileName, nil
}

func (c *caFile) loadCertificate(ctx context.Context) (*x509.Certificate, error) {
	certFile, err := c.buildCertificatePath(ctx)
	if err != nil {
		return nil, err
	}

	certFileInfo, err := os.Stat(certFile)
	if err != nil {
		if errors.Is(err, syscall.ENOENT) {
			return nil, nil
		}

		return nil, fmt.Errorf("could not get info for %s: %w", certFile, err)
	}

	if !certFileInfo.Mode().IsRegular() {
		return nil, fmt.Errorf("certificate file %s is not a regular file", certFile)
	}

	certData, err := os.ReadFile(certFile)
	if err != nil {
		return nil, fmt.Errorf("could not read %s: %w", certFile, err)
	}

	pemData, _ := pem.Decode(certData)
	if pemData == nil {
		return nil, fmt.Errorf("no PEM data in %s", certFile)
	}

	if pemData.Type != "CERTIFICATE" {
		return nil, fmt.Errorf("no certificate found in %s", certFile)
	}

	certificate, err := x509.ParseCertificate(pemData.Bytes)
	if err != nil {
		return nil, fmt.Errorf("could not parse certificate from %s: %w", certFile, err)
	}

	return certificate, nil
}

func (c *caFile) storeCertificate(ctx context.Context, certificate []byte) error {
	certFile, err := c.buildCertificatePath(ctx)
	if err != nil {
		return err
	}

	err = os.WriteFile(certFile, certificate, 0o600)
	if err != nil {
		return fmt.Errorf("could not write certificate file %s: %w", certFile, err)
	}

	return nil
}

func GetRootCACertificate(ctx context.Context, label string) (*x509.Certificate, error) {
	var (
		certificate *x509.Certificate
		keyPair     crypto.Signer
	)

	sc := GetSignerConfig(ctx)

	caCert, err := sc.GetCADefinition(label)
	if err != nil {
		return nil, fmt.Errorf("could not get CA definition for label %s: %w", label, err)
	}

	if !caCert.IsRoot() {
		return nil, fmt.Errorf("CA definition %s is not a root CA definition", label)
	}

	caFile := &caFile{sc: sc, label: label}

	certificate, err = caFile.loadCertificate(ctx)
	if err != nil {
		return nil, err
	}

	if certificate != nil && !IsSetupMode(ctx) {
		caCert.Certificate = certificate

		return certificate, nil
	}

	keyPair, err = getKeyPair(ctx, label, caCert.KeyInfo)
	if err != nil {
		return nil, err
	}

	if certificate != nil && certificateMatches(certificate, keyPair) {
		caCert.Certificate, caCert.KeyPair = certificate, keyPair

		return certificate, nil
	}

	if !IsSetupMode(ctx) {
		return nil, errCertificateGenerationRefused
	}

	notBefore, notAfter := sc.CalculateValidity(caCert, time.Now())
	subject := sc.CalculateSubject(caCert)

	certificate, err = generateRootCACertificate(
		ctx,
		caFile,
		keyPair,
		&x509.Certificate{
			Subject:               subject,
			NotBefore:             notBefore,
			NotAfter:              notAfter,
			MaxPathLen:            0,
			MaxPathLenZero:        false,
			BasicConstraintsValid: true,
			KeyUsage:              x509.KeyUsageCRLSign | x509.KeyUsageCertSign,
			IsCA:                  true,
		},
	)

	if err != nil {
		return nil, err
	}

	p11Context, err := GetP11Context(ctx, caCert)
	if err != nil {
		return nil, err
	}

	err = addCertificate(p11Context, label, certificate)
	if err != nil {
		return nil, err
	}

	caCert.Certificate, caCert.KeyPair = certificate, keyPair

	return certificate, nil
}

func GetIntermediaryCACertificate(ctx context.Context, certLabel string) (*x509.Certificate, error) {
	var (
		certificate *x509.Certificate
		keyPair     crypto.Signer
	)

	sc := GetSignerConfig(ctx)

	caCert, err := sc.GetCADefinition(certLabel)
	if err != nil {
		return nil, fmt.Errorf("could not get CA definition for label %s: %w", certLabel, err)
	}

	if caCert.IsRoot() {
		return nil, fmt.Errorf(
			"CA definition %s is a root CA definition, intermediary expected",
			certLabel,
		)
	}

	keyPair, err = getKeyPair(ctx, certLabel, caCert.KeyInfo)
	if err != nil {
		return nil, err
	}

	certFile := &caFile{sc: sc, label: certLabel}

	certificate, err = certFile.loadCertificate(ctx)
	if err != nil {
		return nil, err
	}

	if certificate != nil && certificateMatches(certificate, keyPair) {
		caCert.Certificate, caCert.KeyPair = certificate, keyPair

		return certificate, nil
	}

	if !IsSetupMode(ctx) {
		return nil, errCertificateGenerationRefused
	}

	notBefore, notAfter := sc.CalculateValidity(caCert, time.Now())
	subject := sc.CalculateSubject(caCert)

	certificate, err = generateIntermediaryCACertificate(
		ctx,
		certFile,
		sc,
		certLabel,
		keyPair.Public(),
		&x509.Certificate{
			Subject:               subject,
			NotBefore:             notBefore,
			NotAfter:              notAfter,
			MaxPathLen:            caCert.MaxPathLen,
			MaxPathLenZero:        true,
			BasicConstraintsValid: true,
			IsCA:                  true,
			KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign | x509.KeyUsageCRLSign,
			ExtKeyUsage:           caCert.ExtKeyUsage,
			IssuingCertificateURL: []string{sc.BuildIssuerURL(caCert)},
			OCSPServer:            []string{sc.BuildOCSPURL(caCert)},
			CRLDistributionPoints: []string{sc.BuildCRLUrl(caCert)},
			PolicyIdentifiers: []asn1.ObjectIdentifier{
				// use policy identifiers from http://wiki.cacert.org/OidAllocation
				oidCAcertClass3PolicyV1,
			},
		},
	)

	if err != nil {
		return nil, err
	}

	p11Context, err := GetP11Context(ctx, caCert)
	if err != nil {
		return nil, err
	}

	err = addCertificate(p11Context, certLabel, certificate)
	if err != nil {
		return nil, err
	}

	caCert.Certificate, caCert.KeyPair = certificate, keyPair

	return certificate, nil
}

func generateIntermediaryCACertificate(
	ctx context.Context,
	certFile *caFile,
	config *config.SignerConfig,
	certLabel string,
	publicKey crypto.PublicKey,
	template *x509.Certificate,
) (*x509.Certificate, error) {
	parent, err := config.GetParentCA(certLabel)
	if err != nil {
		return nil, fmt.Errorf("could not get parent CA for label %s: %w", certLabel, err)
	}

	serial, err := randomSerialNumber()
	if err != nil {
		return nil, err
	}

	template.SerialNumber = serial
	template.SignatureAlgorithm, err = determineSignatureAlgorithm(parent.KeyPair)

	if err != nil {
		return nil, err
	}

	certBytes, err := x509.CreateCertificate(
		rand.Reader,
		template,
		parent.Certificate,
		publicKey,
		parent.KeyPair,
	)
	if err != nil {
		return nil, fmt.Errorf("could not create intermediary CA certificate: %w", err)
	}

	certBlock := &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	}

	err = certFile.storeCertificate(ctx, pem.EncodeToMemory(certBlock))
	if err != nil {
		return nil, err
	}

	certificate, err := x509.ParseCertificate(certBytes)
	if err != nil {
		return nil, fmt.Errorf("could not parse generated certificate: %w", err)
	}

	return certificate, nil
}

func addCertificate(p11Context *crypto11.Context, label string, certificate *x509.Certificate) error {
	objectID, err := randomObjectID()
	if err != nil {
		return err
	}

	err = p11Context.ImportCertificateWithLabel(objectID, []byte(label), certificate)
	if err != nil {
		return fmt.Errorf("could not import certificate into token: %w", err)
	}

	return nil
}

func getKeyPair(ctx context.Context, label string, keyInfo *config.PrivateKeyInfo) (crypto.Signer, error) {
	sc := GetSignerConfig(ctx)

	cert, err := sc.GetCADefinition(label)
	if err != nil {
		return nil, fmt.Errorf("could not get CA definition for label %s: %w", label, err)
	}

	if cert.KeyPair != nil {
		return cert.KeyPair, nil
	}

	p11Context, err := GetP11Context(ctx, cert)
	if err != nil {
		return nil, err
	}

	keyPair, err := p11Context.FindKeyPair(nil, []byte(label))
	if err != nil {
		return nil, fmt.Errorf("could not find requested key pair: %w", err)
	}

	if keyPair != nil {
		return keyPair, nil
	}

	if !IsSetupMode(ctx) {
		return nil, errKeyGenerationRefused
	}

	switch keyInfo.Algorithm {
	case x509.RSA:
		keyPair, err = generateRSAKeyPair(p11Context, label, keyInfo)
		if err != nil {
			return nil, err
		}
	case x509.ECDSA:
		keyPair, err = generateECDSAKeyPair(p11Context, label, keyInfo)
		if err != nil {
			return nil, err
		}
	default:
		return nil, fmt.Errorf(
			"could not generate private key with label %s with unsupported key algorithm %s",
			label,
			keyInfo.Algorithm,
		)
	}

	return keyPair, nil
}

func generateECDSAKeyPair(
	p11Context *crypto11.Context,
	label string,
	keyInfo *config.PrivateKeyInfo,
) (crypto11.Signer, error) {
	var (
		err         error
		newObjectID []byte
		signer      crypto11.Signer
	)

	newObjectID, err = randomObjectID()
	if err != nil {
		return nil, err
	}

	signer, err = p11Context.GenerateECDSAKeyPairWithLabel(newObjectID, []byte(label), keyInfo.EccCurve)
	if err != nil {
		return nil, fmt.Errorf("could not generate ECDSA key pair: %w", err)
	}

	return signer, nil
}

func generateRSAKeyPair(
	p11Context *crypto11.Context,
	label string,
	keyInfo *config.PrivateKeyInfo,
) (crypto11.Signer, error) {
	var (
		err         error
		newObjectID []byte
		signer      crypto11.Signer
	)

	newObjectID, err = randomObjectID()
	if err != nil {
		return nil, err
	}

	signer, err = p11Context.GenerateRSAKeyPairWithLabel(newObjectID, []byte(label), keyInfo.RSABits)
	if err != nil {
		return nil, fmt.Errorf("could not generate RSA key pair: %w", err)
	}

	return signer, nil
}

func randomObjectID() ([]byte, error) {
	const objectIDSize = 20

	result := make([]byte, objectIDSize)

	_, err := rand.Read(result)
	if err != nil {
		return nil, fmt.Errorf("could not create new random object id: %w", err)
	}

	return result, nil
}

func generateRootCACertificate(
	ctx context.Context,
	certFile *caFile,
	keyPair crypto.Signer,
	template *x509.Certificate,
) (*x509.Certificate, error) {
	serial, err := randomSerialNumber()
	if err != nil {
		return nil, err
	}

	template.SerialNumber = serial
	template.SignatureAlgorithm, err = determineSignatureAlgorithm(keyPair)

	if err != nil {
		return nil, err
	}

	certBytes, err := x509.CreateCertificate(
		rand.Reader,
		template,
		template,
		keyPair.Public(),
		keyPair,
	)
	if err != nil {
		return nil, fmt.Errorf("could not create root certificate: %w", err)
	}

	certBlock := &pem.Block{
		Type:  "CERTIFICATE",
		Bytes: certBytes,
	}

	if err = certFile.storeCertificate(ctx, pem.EncodeToMemory(certBlock)); err != nil {
		return nil, err
	}

	certificate, err := x509.ParseCertificate(certBytes)
	if err != nil {
		return nil, fmt.Errorf("could not parse generated certificate: %w", err)
	}

	return certificate, nil
}

func determineSignatureAlgorithm(keyPair crypto.Signer) (x509.SignatureAlgorithm, error) {
	switch keyPair.Public().(type) {
	case *ecdsa.PublicKey:
		return x509.ECDSAWithSHA256, nil
	case *rsa.PublicKey:
		return x509.SHA256WithRSA, nil
	default:
		return x509.UnknownSignatureAlgorithm,
			fmt.Errorf("could not determine signature algorithm for key of type %T", keyPair)
	}
}

func certificateMatches(certificate *x509.Certificate, key crypto.Signer) bool {
	switch v := certificate.PublicKey.(type) {
	case *ecdsa.PublicKey:
		if pub, ok := key.Public().(*ecdsa.PublicKey); ok {
			if v.Equal(pub) {
				return true
			}
		}
	case *rsa.PublicKey:
		if pub, ok := key.Public().(*rsa.PublicKey); ok {
			if v.Equal(pub) {
				return true
			}
		}
	default:
		log.Printf("unsupported public key %v", v)
	}

	log.Printf(
		"public key from certificate does not match private key: %s != %s",
		certificate.PublicKey,
		key.Public(),
	)

	return false
}

func randomSerialNumber() (*big.Int, error) {
	serialNumber, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), 128))
	if err != nil {
		return nil, fmt.Errorf("could not generate serial number: %w", err)
	}

	return serialNumber, nil
}