handler.go

/*
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/rand"
	"crypto/rsa"
	"crypto/x509"
	"crypto/x509/pkix"
	"fmt"
	"io"
	"log"
	"math/big"
	"strings"
	"time"

	"code.vereign.com/code/viam-apis/versions"
	"github.com/golang/protobuf/proto"

	"code.vereign.com/code/viam-apis/authentication"
	"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"
	"golang.org/x/net/context"
	"google.golang.org/grpc/metadata"
)

// Server represents the gRPC server
type KeyStorageServerImpl struct {
	DataStorageUrl string
	CertPath       string
}

func (s *KeyStorageServerImpl) CreateAuthentication(ctx context.Context) *authentication.Authentication {
	if md, ok := metadata.FromIncomingContext(ctx); ok {
		uuid := strings.Join(md["uuid"], "")
		session := strings.Join(md["session"], "")

		auth := &authentication.Authentication{
			Uuid:    uuid,
			Session: session,
		}

		return auth
	}

	return nil
}

func (s *KeyStorageServerImpl) GetKey(ctx context.Context, in *api.GetKeyRequest) (*api.GetKeyResponse, error) {
	auth := s.CreateAuthentication(ctx)

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

	getKeyResponse := &api.GetKeyResponse{}

	if in.KeyType == api.KeyType_KT_EMPTY {
		getKeyResponse.StatusList = utils.AddStatus(getKeyResponse.StatusList, "400", api.StatusType_ERROR, "KeyType cannot be empty")
	}

	data, _ := client.DoGetDataCall("keys", in.Uuid+"/"+api.KeyType.String(in.KeyType))

	if data.Errors != "" {
		getKeyResponse.Key = nil
		getKeyResponse.StatusList = utils.AddStatus(getKeyResponse.StatusList, "500", api.StatusType_ERROR, data.Errors)
	} else {
		key := &api.Key{}
		proto.Unmarshal(data.Data.Data, key)
		getKeyResponse.Key = key
	}

	return getKeyResponse, nil
}

func (s *KeyStorageServerImpl) SetKey(ctx context.Context, in *api.SetKeyRequest) (*api.SetKeyResponse, error) {
	auth := s.CreateAuthentication(ctx)

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

	setKeyResponse := &api.SetKeyResponse{}

	if in.KeyType == api.KeyType_KT_EMPTY {
		setKeyResponse.StatusList = utils.AddStatus(setKeyResponse.StatusList, "400", api.StatusType_ERROR, "KeyType cannot be empty")
	}

	data, _ := client.DoGetDataCall("keys", in.Uuid+"/"+api.KeyType.String(in.KeyType))

	if data.Errors != "" {
		setKeyResponse.StatusList = utils.AddStatus(setKeyResponse.StatusList, "400", api.StatusType_ERROR, data.Errors)
		return setKeyResponse, nil
	}

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

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

	return setKeyResponse, nil
}

func (s *KeyStorageServerImpl) GenerateKeyPair(ctx context.Context, in *api.GenerateKeyPairRequest) (*api.GenerateKeyPairResponse, error) {
	auth := s.CreateAuthentication(ctx)

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

	uuid, err := newUUID()
	privateKeyBytes, publicKeyBytes, _ := generateKeyPair(int(in.KeySize * 8))

	privateKey := &api.Key{Content: privateKeyBytes}
	publicKey := &api.Key{Content: publicKeyBytes}

	generateKeyPairResponse := &api.GenerateKeyPairResponse{
		Uuid: uuid,
	}

	result, errors, err := client.DoPutDataCall("keys", uuid+"/"+api.KeyType.String(api.KeyType_PRIVATE), privateKey, versions.EntitiesManagementAgentApiVersion)
	generateKeyPairResponse.StatusList = handlePutDataErrors(generateKeyPairResponse.StatusList, errors, err)

	if generateKeyPairResponse.StatusList == nil || len(generateKeyPairResponse.StatusList) == 0 {
		result, errors, err = client.DoPutDataCall("keys", uuid+"/"+api.KeyType.String(api.KeyType_PUBLIC), publicKey, versions.EntitiesManagementAgentApiVersion)
		generateKeyPairResponse.StatusList = handlePutDataErrors(generateKeyPairResponse.StatusList, errors, err)
	}

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

	return generateKeyPairResponse, nil
}

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.CertPath)
	defer client.CloseClient()

	generateCertificateResponse := &api.GenerateCertificateResponse{}

	privateKeyMessage := &api.Key{}
	publicKeyMessage := &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)
	} else {
		proto.Unmarshal(data.Data.Data, privateKeyMessage)
	}

	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)
	} else {
		proto.Unmarshal(data.Data.Data, publicKeyMessage)
	}

	certificateBytes, err := generateCertificate(privateKeyMessage.Content, publicKeyMessage.Content, in.CertificateData)
	if err != nil {
		generateCertificateResponse.StatusList = utils.AddStatus(generateCertificateResponse.StatusList, "400", api.StatusType_ERROR, err.Error())
	}

	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 (s *KeyStorageServerImpl) ReserveKeyUUID(ctx context.Context, in *api.ReserveKeyUUIDRequest) (*api.ReserveKeyUUIDResponse, error) {
	auth := s.CreateAuthentication(ctx)

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

	uuid, err := newUUID()
	emptyKey := &api.Key{
		Content: []byte{},
	}

	reserveKeyUUIDResponse := &api.ReserveKeyUUIDResponse{
		Uuid: uuid,
	}

	result, errors, err := client.DoPutDataCall("keys", uuid+"/"+api.KeyType.String(api.KeyType_PRIVATE), emptyKey, versions.EntitiesManagementAgentApiVersion)
	reserveKeyUUIDResponse.StatusList = handlePutDataErrors(reserveKeyUUIDResponse.StatusList, errors, err)

	if reserveKeyUUIDResponse.StatusList == nil || len(reserveKeyUUIDResponse.StatusList) == 0 {
		result, errors, err = client.DoPutDataCall("keys", uuid+"/"+api.KeyType.String(api.KeyType_PUBLIC), emptyKey, versions.EntitiesManagementAgentApiVersion)
		reserveKeyUUIDResponse.StatusList = handlePutDataErrors(reserveKeyUUIDResponse.StatusList, errors, err)
	}

	if reserveKeyUUIDResponse.StatusList == nil || len(reserveKeyUUIDResponse.StatusList) == 0 {
		result, errors, err = client.DoPutDataCall("keys", uuid+"/"+api.KeyType.String(api.KeyType_CERTIFICATE), emptyKey, versions.EntitiesManagementAgentApiVersion)
		reserveKeyUUIDResponse.StatusList = handlePutDataErrors(reserveKeyUUIDResponse.StatusList, errors, err)
	}

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

	return reserveKeyUUIDResponse, nil
}

func generateKeyPair(bitSize int) ([]byte, []byte, error) {
	privateKey, err := rsa.GenerateKey(rand.Reader, bitSize)
	if err != nil {
		return nil, nil, err
	}

	err = privateKey.Validate()
	if err != nil {
		return nil, nil, err
	}

	publicKey := &privateKey.PublicKey

	privateKeyBytes := x509.MarshalPKCS1PrivateKey(privateKey)
	publicKeyBytes := x509.MarshalPKCS1PublicKey(publicKey)

	return privateKeyBytes, publicKeyBytes, nil
}

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

	privateKey, err := x509.ParsePKCS1PrivateKey(privateKeyBytes)
	if err != nil {
		return nil, err
	}
	publicKey, err := x509.ParsePKCS1PublicKey(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:     true,
		DNSNames: []string{certificateData.Host},
	}

	certificateBytes, err := x509.CreateCertificate(rand.Reader, &template, &template, publicKey, privateKey)
	if err != nil {
		log.Fatalf("Failed to create certificate: %s", err)
	}

	return certificateBytes, nil
}
func newUUID() (string, error) {
	uuid := make([]byte, 16)
	n, err := io.ReadFull(rand.Reader, uuid)
	if n != len(uuid) || err != nil {
		return "", err
	}

	uuid[8] = uuid[8]&^0xc0 | 0x80
	uuid[6] = uuid[6]&^0xf0 | 0x40

	return fmt.Sprintf("%x-%x-%x-%x-%x", uuid[0:4], uuid[4:6], uuid[6:8], uuid[8:10], uuid[10:]), nil
}

func handlePutDataErrors(statusList []*api.Status, errors string, err error) []*api.Status {
	if err != nil {
		statusList = utils.AddStatus(statusList, "500", api.StatusType_ERROR, err.Error())
	} else if errors != "" {
		statusList = utils.AddStatus(statusList, "400", api.StatusType_ERROR, errors)
	}

	return statusList
}