/*
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 (
	"crypto"
	"crypto/ecdsa"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"encoding/asn1"
	"errors"
	"fmt"
	"math/big"
	"os"
	"path"
	"syscall"
	"time"

	"github.com/ThalesIgnite/crypto11"
	"github.com/sirupsen/logrus"

	"git.cacert.org/cacert-gosigner/internal/cainfo"
	"git.cacert.org/cacert-gosigner/internal/config"
	"git.cacert.org/cacert-gosigner/internal/health"
	"git.cacert.org/cacert-gosigner/pkg/messages"
)

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(caDirectory string) string {
	fileName := c.sc.CertificateFileName(c.label)

	if caDirectory == "" {
		return fileName
	}

	return path.Join(caDirectory, fileName)
}

func (c *caFile) loadCertificate(caDirectory string) (*x509.Certificate, error) {
	certFile := c.buildCertificatePath(caDirectory)

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

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

	return certificate, nil
}

func (c *caFile) storeCertificate(caDirectory string, certificate []byte) error {
	certFile := c.buildCertificatePath(caDirectory)

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

	return nil
}

type Access struct {
	logger       *logrus.Logger
	caDirectory  string
	signerConfig *config.SignerConfig
	p11Contexts  map[string]*crypto11.Context
	setupMode    bool
	verbose      bool
}

func (a *Access) Healthy() (*health.Info, error) {
	healthy := true

	moreInfo := make(map[string]string)

	for _, ca := range a.signerConfig.RootCAs() {
		_, err := a.GetRootCACertificate(ca)
		if err != nil {
			healthy = false

			moreInfo[ca] = string(messages.CertStatusFailed)

			continue
		}

		moreInfo[ca] = string(messages.CertStatusOk)
	}

	for _, ca := range a.signerConfig.SubordinateCAs() {
		_, err := a.GetSubordinateCACertificate(ca)
		if err != nil {
			healthy = false

			moreInfo[ca] = string(messages.CertStatusFailed)

			continue
		}

		def, err := a.signerConfig.GetCADefinition(ca)
		if err != nil {
			healthy = false

			moreInfo[ca] = string(messages.CertStatusFailed)

			continue
		}

		_, err = a.getKeyPair(ca, def.KeyInfo)
		if err != nil {
			healthy = false

			moreInfo[ca] = string(messages.CertStatusFailed)

			continue
		}

		moreInfo[ca] = string(messages.CertStatusOk)
	}

	return &health.Info{
		Healthy:  healthy,
		Source:   "HSM",
		MoreInfo: moreInfo,
	}, nil
}

func NewAccess(logger *logrus.Logger, options ...ConfigOption) (*Access, error) {
	access := &Access{logger: logger}
	access.setupContext(options...)

	return access, nil
}

func (a *Access) GetCAInfo(name string) (*cainfo.Result, error) {
	sc := a.GetSignerConfig()

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

	def, err := sc.GetCADefinition(name)
	if err != nil {
		return nil, fmt.Errorf("certificate %s unknown: %w", name, err)
	}

	certificate, err := caFile.loadCertificate(a.caDirectory)
	if err != nil {
		return nil, fmt.Errorf("certificate %s not available: %w", name, err)
	}

	return &cainfo.Result{
		Certificate: certificate.Raw,
		Profiles:    cainfo.MapProfiles(def.Profiles),
	}, nil
}

func (a *Access) GetRootCACertificate(label string) (*x509.Certificate, error) {
	var (
		certificate *x509.Certificate
		keyPair     crypto.Signer
	)

	sc := a.GetSignerConfig()

	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(a.caDirectory)
	if err != nil {
		return nil, err
	}

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

		return certificate, nil
	}

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

	if certificate != nil {
		err := certificateMatches(certificate, keyPair)
		if err != nil {
			return nil, err
		}

		caCert.Certificate, caCert.KeyPair = certificate, keyPair

		return certificate, nil
	}

	if !a.IsSetupMode() {
		return nil, errCertificateGenerationRefused
	}

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

	certificate, err = a.generateRootCACertificate(
		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 := a.GetP11Context(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 (a *Access) GetSubordinateCACertificate(certLabel string) (*x509.Certificate, error) {
	var (
		certificate *x509.Certificate
		keyPair     crypto.Signer
	)

	sc := a.GetSignerConfig()

	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, subordinate expected",
			certLabel,
		)
	}

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

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

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

	if certificate != nil {
		err := certificateMatches(certificate, keyPair)
		if err != nil {
			return nil, err
		}

		caCert.Certificate, caCert.KeyPair = certificate, keyPair

		return certificate, nil
	}

	if !a.IsSetupMode() {
		return nil, errCertificateGenerationRefused
	}

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

	certificate, err = a.generateSubordinateCACertificate(
		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(certLabel)},
			PolicyIdentifiers: []asn1.ObjectIdentifier{
				// use policy identifiers from http://wiki.cacert.org/OidAllocation
				oidCAcertClass3PolicyV1,
			},
		},
	)

	if err != nil {
		return nil, err
	}

	p11Context, err := a.GetP11Context(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 (a *Access) generateSubordinateCACertificate(
	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 subordinate CA certificate %s: %w", certLabel, err)
	}

	err = certFile.storeCertificate(a.caDirectory, certBytes)
	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 (a *Access) getKeyPair(label string, keyInfo *config.PrivateKeyInfo) (crypto.Signer, error) {
	sc := a.GetSignerConfig()

	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 := a.GetP11Context(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 !a.IsSetupMode() {
		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 (a *Access) generateRootCACertificate(
	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)
	}

	if err = certFile.storeCertificate(a.caDirectory, certBytes); 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) error {
	switch v := certificate.PublicKey.(type) {
	case *ecdsa.PublicKey:
		if pub, ok := key.Public().(*ecdsa.PublicKey); ok {
			if v.Equal(pub) {
				return nil
			}
		}
	case *rsa.PublicKey:
		if pub, ok := key.Public().(*rsa.PublicKey); ok {
			if v.Equal(pub) {
				return nil
			}
		}
	default:
		return fmt.Errorf("unsupported public key %v", v)
	}

	return fmt.Errorf(
		"public key from certificate does not match private key: %s != %s",
		certificate.PublicKey,
		key.Public(),
	)
}

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
}