/*
   Copyright The ocicrypt Authors.

   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 utils

import (
	"crypto/ecdsa"
	"crypto/elliptic"
	"crypto/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"math/big"
	"time"

	"github.com/pkg/errors"
)

// CreateRSAKey creates an RSA key
func CreateRSAKey(bits int) (*rsa.PrivateKey, error) {
	key, err := rsa.GenerateKey(rand.Reader, bits)
	if err != nil {
		return nil, errors.Wrap(err, "rsa.GenerateKey failed")
	}
	return key, nil
}

// CreateRSATestKey creates an RSA key of the given size and returns
// the public and private key in PEM or DER format
func CreateRSATestKey(bits int, password []byte, pemencode bool) ([]byte, []byte, error) {
	key, err := CreateRSAKey(bits)
	if err != nil {
		return nil, nil, err
	}

	pubData, err := x509.MarshalPKIXPublicKey(&key.PublicKey)
	if err != nil {
		return nil, nil, errors.Wrap(err, "x509.MarshalPKIXPublicKey failed")
	}
	privData := x509.MarshalPKCS1PrivateKey(key)

	// no more encoding needed for DER
	if !pemencode {
		return pubData, privData, nil
	}

	publicKey := pem.EncodeToMemory(&pem.Block{
		Type:  "PUBLIC KEY",
		Bytes: pubData,
	})

	var block *pem.Block

	typ := "RSA PRIVATE KEY"
	if len(password) > 0 {
		block, err = x509.EncryptPEMBlock(rand.Reader, typ, privData, password, x509.PEMCipherAES256)
		if err != nil {
			return nil, nil, errors.Wrap(err, "x509.EncryptPEMBlock failed")
		}
	} else {
		block = &pem.Block{
			Type:  typ,
			Bytes: privData,
		}
	}

	privateKey := pem.EncodeToMemory(block)

	return publicKey, privateKey, nil
}

// CreateECDSATestKey creates and elliptic curve key for the given curve and returns
// the public and private key in DER format
func CreateECDSATestKey(curve elliptic.Curve) ([]byte, []byte, error) {
	key, err := ecdsa.GenerateKey(curve, rand.Reader)
	if err != nil {
		return nil, nil, errors.Wrapf(err, "ecdsa.GenerateKey failed")
	}

	pubData, err := x509.MarshalPKIXPublicKey(&key.PublicKey)
	if err != nil {
		return nil, nil, errors.Wrapf(err, "x509.MarshalPKIXPublicKey failed")
	}

	privData, err := x509.MarshalECPrivateKey(key)
	if err != nil {
		return nil, nil, errors.Wrapf(err, "x509.MarshalECPrivateKey failed")
	}

	return pubData, privData, nil
}

// CreateTestCA creates a root CA for testing
func CreateTestCA() (*rsa.PrivateKey, *x509.Certificate, error) {
	key, err := rsa.GenerateKey(rand.Reader, 2048)
	if err != nil {
		return nil, nil, errors.Wrap(err, "rsa.GenerateKey failed")
	}

	ca := &x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			CommonName: "test-ca",
		},
		NotBefore:             time.Now(),
		NotAfter:              time.Now().AddDate(1, 0, 0),
		IsCA:                  true,
		KeyUsage:              x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		BasicConstraintsValid: true,
	}
	caCert, err := certifyKey(&key.PublicKey, ca, key, ca)

	return key, caCert, err
}

// CertifyKey certifies a public key using the given CA's private key and cert;
// The certificate template for the public key is optional
func CertifyKey(pubbytes []byte, template *x509.Certificate, caKey *rsa.PrivateKey, caCert *x509.Certificate) (*x509.Certificate, error) {
	pubKey, err := ParsePublicKey(pubbytes, "CertifyKey")
	if err != nil {
		return nil, err
	}
	return certifyKey(pubKey, template, caKey, caCert)
}

func certifyKey(pub interface{}, template *x509.Certificate, caKey *rsa.PrivateKey, caCert *x509.Certificate) (*x509.Certificate, error) {
	if template == nil {
		template = &x509.Certificate{
			SerialNumber: big.NewInt(1),
			Subject: pkix.Name{
				CommonName: "testkey",
			},
			NotBefore:             time.Now(),
			NotAfter:              time.Now().Add(time.Hour),
			IsCA:                  false,
			KeyUsage:              x509.KeyUsageDigitalSignature,
			BasicConstraintsValid: true,
		}
	}

	certDER, err := x509.CreateCertificate(rand.Reader, template, caCert, pub, caKey)
	if err != nil {
		return nil, errors.Wrap(err, "x509.CreateCertificate failed")
	}

	cert, err := x509.ParseCertificate(certDER)
	if err != nil {
		return nil, errors.Wrap(err, "x509.ParseCertificate failed")
	}

	return cert, nil
}