/*
Copyright (c) 2018 Vereign AG [https://www.vereign.com]

This is free software: you can redistribute it and/or modify
it under the terms of the GNU Affero General Public License as
published by the Free Software Foundation, either version 3 of the
License, or (at your option) any later version.

This program is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
GNU Affero General Public License for more details.

You should have received a copy of the GNU Affero General Public License
along with this program. If not, see <http://www.gnu.org/licenses/>.
*/

package handler

import (
	"crypto/aes"
	"crypto/cipher"
	"crypto/rand"
	"crypto/rsa"
	"crypto/sha256"
	"crypto/x509"
	"crypto/x509/pkix"
	"encoding/pem"
	"io/ioutil"
	"math/big"
	"time"

	"code.vereign.com/code/viam-apis/data-storage-agent/client"
	"code.vereign.com/code/viam-apis/key-storage-agent/api"
	"code.vereign.com/code/viam-apis/utils"
	"code.vereign.com/code/viam-apis/versions"
	"github.com/golang/protobuf/proto"
	"golang.org/x/net/context"
)

func (s *KeyStorageServerImpl) GenerateCertificate(ctx context.Context, in *api.GenerateCertificateRequest) (*api.GenerateCertificateResponse, error) {
	auth := s.CreateAuthentication(ctx)

	client := &client.DataStorageClientImpl{}
	client.SetUpClient(auth, s.DataStorageUrl, s.CertFilePath)
	defer client.CloseClient()

	generateCertificateResponse := &api.GenerateCertificateResponse{}

	aesKeyBytes, err := rsaDecryptWithServerKey(s.VereignPrivateKeyFilePath, in.EncryptedAesKey, []byte("aeskeys"))
	if err != nil {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList,
			"400", api.StatusType_ERROR, err.Error())
		return generateCertificateResponse, nil
	}

	// Get and decrypt rsa private key
	encryptedPrivateKeyMessage := &api.Key{}
	data, _ := client.DoGetDataCall("keys", in.Uuid+"/"+api.KeyType.String(api.KeyType_PRIVATE))
	if data.Errors != "" {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList,
			"400", api.StatusType_ERROR, data.Errors)
		return generateCertificateResponse, nil
	}
	proto.Unmarshal(data.Data.Data, encryptedPrivateKeyMessage)

	privateKeyBytes, err := aesDecrypt(aesKeyBytes, in.PrivateKeyNonce, encryptedPrivateKeyMessage.Content)
	if err != nil {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList,
			"400", api.StatusType_ERROR, err.Error())
		return generateCertificateResponse, nil
	}

	// Get and decrypt rsa public key
	encryptedPublicKeyMessage := &api.Key{}
	data, _ = client.DoGetDataCall("keys", in.Uuid+"/"+api.KeyType.String(api.KeyType_PUBLIC))
	if data.Errors != "" {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList,
			"400", api.StatusType_ERROR, data.Errors)
		return generateCertificateResponse, nil
	}
	proto.Unmarshal(data.Data.Data, encryptedPublicKeyMessage)

	publicKeyBytes, err := aesDecrypt(aesKeyBytes, in.PublicKeyNonce, encryptedPublicKeyMessage.Content)
	if err != nil {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList,
			"400", api.StatusType_ERROR, err.Error())
		return generateCertificateResponse, nil
	}

	certificateBytes, err := generateCertificate(privateKeyBytes, publicKeyBytes, s.VereignCertFilePath, in.CertificateData)
	if err != nil {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList,
			"400", api.StatusType_ERROR, err.Error())
		return generateCertificateResponse, nil
	}

	certificateMessage := &api.Key{
		Content: certificateBytes,
	}

	result, errors, err := client.DoPutDataCall("keys", in.Uuid+"/"+api.KeyType.String(api.KeyType_CERTIFICATE), certificateMessage, versions.EntitiesManagementAgentApiVersion)
	generateCertificateResponse.StatusList = handlePutDataErrors(generateCertificateResponse.StatusList, errors, err)

	if generateCertificateResponse.StatusList == nil || len(generateCertificateResponse.StatusList) == 0 {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList,
			"200", api.StatusType_INFO, result)
	}

	return generateCertificateResponse, nil
}

func generateCertificate(privateKeyBytes []byte, publicKeyBytes []byte, caCertFilePath string,
	certificateData *api.GenerateCertificateRequest_CertificateData) ([]byte, error) {

	privateKey, err := x509.ParsePKCS8PrivateKey(privateKeyBytes)
	if err != nil {
		return nil, err
	}
	publicKey, err := x509.ParsePKIXPublicKey(publicKeyBytes)
	if err != nil {
		return nil, err
	}

	notBeforeTime := time.Unix(certificateData.NotBefore.Seconds, int64(certificateData.NotBefore.Nanos)).UTC()
	notAfterTime := time.Unix(certificateData.NotAfter.Seconds, int64(certificateData.NotAfter.Nanos)).UTC()

	template := x509.Certificate{
		SerialNumber: big.NewInt(1),
		Subject: pkix.Name{
			Country:            []string{certificateData.Country},
			Organization:       []string{certificateData.Organization},
			OrganizationalUnit: []string{certificateData.OrganizationalUnit},
			CommonName:         certificateData.CommonName,
		},
		NotBefore:             notBeforeTime,
		NotAfter:              notAfterTime,
		KeyUsage:              x509.KeyUsageKeyEncipherment | x509.KeyUsageDigitalSignature | x509.KeyUsageCertSign,
		ExtKeyUsage:           []x509.ExtKeyUsage{x509.ExtKeyUsageServerAuth},
		BasicConstraintsValid: true,
		IsCA:     false,
		DNSNames: []string{certificateData.Host},
	}

	caCertificate, err := readCertificateFromFile(caCertFilePath)
	if err != nil {
		return nil, err
	}

	certificateBytes, err := x509.CreateCertificate(rand.Reader, &template, caCertificate, publicKey, privateKey)
	if err != nil {
		return nil, err
	}

	return certificateBytes, nil
}

func readPrivateKeyFromFile(fileName string) (*rsa.PrivateKey, error) {
	privateKeyPemBlock, err := readPemBlockFromFile(fileName)
	if err != nil {
		return nil, err
	}

	parseResult, err := x509.ParsePKCS8PrivateKey(privateKeyPemBlock.Bytes)
	if err != nil {
		return nil, err
	}
	privateKey := parseResult.(*rsa.PrivateKey)

	return privateKey, nil
}

func readPemBlockFromFile(fileName string) (*pem.Block, error) {
	fileBytes, err := ioutil.ReadFile(fileName)
	if err != nil {
		return nil, err
	}

	certificatePemBlock, _ := pem.Decode(fileBytes)

	return certificatePemBlock, nil
}

func rsaDecryptWithServerKey(privateKeyFilePath string, encryptedMessage []byte, label []byte) ([]byte, error) {
	serverPrivateKey, err := readPrivateKeyFromFile(privateKeyFilePath)
	if err != nil {
		return nil, err
	}

	message, err := rsa.DecryptOAEP(sha256.New(), rand.Reader, serverPrivateKey, encryptedMessage, label)
	if err != nil {
		return nil, err
	}

	return message, nil
}

func aesDecrypt(aesKey []byte, nonce []byte, encryptedMessage []byte) ([]byte, error) {
	block, err := aes.NewCipher(aesKey)
	if err != nil {
		return nil, err
	}

	aesgcm, err := cipher.NewGCM(block)
	if err != nil {
		return nil, err
	}

	message, err := aesgcm.Open(nil, nonce, encryptedMessage, nil)
	if err != nil {
		return nil, err
	}

	return message, nil
}