kms.go

// This package kms implements a simplistic key management system (kms) for
// Quantum Key Distribution Networks (QKDN) which is a simple emulated KMS. x
// It relies on the emulated quantum link out of the quantumlayer package

package kms

import (
	"context"
	"encoding/base64"
	"fmt"
	"io"
	"net"
	"sync"
	"time"

	log "github.com/sirupsen/logrus"
	"google.golang.org/grpc"
	"google.golang.org/grpc/credentials"
	"google.golang.org/grpc/credentials/insecure"
	"google.golang.org/grpc/health"
	healthpb "google.golang.org/grpc/health/grpc_health_v1"
	"google.golang.org/grpc/metadata"

	pbIC "code.fbi.h-da.de/danet/quant/goKMS/api/gen/proto/go/kmsintercom"
	"code.fbi.h-da.de/danet/quant/goKMS/config"
	"code.fbi.h-da.de/danet/quant/goKMS/kms/akms/client"
	"code.fbi.h-da.de/danet/quant/goKMS/kms/akms/server"
	"code.fbi.h-da.de/danet/quant/goKMS/kms/crypto"
	"code.fbi.h-da.de/danet/quant/goKMS/kms/event"
	"code.fbi.h-da.de/danet/quant/goKMS/kms/peers"
	"code.fbi.h-da.de/danet/quant/goKMS/kms/store"
	"code.fbi.h-da.de/danet/quant/goKMS/kms/util"
	pbQS "code.fbi.h-da.de/danet/quipsec/gen/go/quipsec"
	"github.com/google/uuid"
)

type Route struct {
	PathId    uuid.UUID
	Previous  *peers.Peer
	Next      *peers.Peer
	RemoteKMS *RemoteKMS
}

type RemoteKMS struct {
	Id      string
	Address string
	Port    uint16
}

type BitKeyLength string

const (
	BitKeyLen128 BitKeyLength = "128"
	BitKeyLen256 BitKeyLength = "256"
	BitKeyLen512 BitKeyLength = "512"
)

type PlatformKey struct {
	Id        uuid.UUID
	Value     []byte
	ProcessId string
}

// The general emulated KMS.
type KMS struct {
	kmsName        string
	kmsUUID        uuid.UUID
	interComAddr   string
	quantumAddress string
	tlsData        TlsData
	// TODO create a mapping between ids and address
	remoteKMSMapping      map[string]*RemoteKMS
	remoteKMSMappingMutex sync.RWMutex
	quantumModules        map[string]peers.QuantumModule
	quantumModulesMutex   sync.RWMutex
	kmsPeersMutex         sync.Mutex
	// NOTE: There is probably a better way to handle this
	PKStore      map[string]map[uuid.UUID]*PlatformKey
	PKStoreMutex sync.Mutex
	// TODO(maba): find a better name for this
	routingTable      map[uuid.UUID]*Route
	routingTableMutex sync.RWMutex
	KmsPeers          map[string]*peers.Peer
	pbIC.UnimplementedKmsTalkerServer
	supportedKeyLengths map[BitKeyLength]bool
	eventBus            *event.EventBus
	CKMSAkmsClient      client.CkmsAkmsClient
	CKMSAkmsServer      *server.AKMSReceiverServer
}

type TlsData struct {
	TLS      bool
	CaFile   string
	CertFile string
	KeyFile  string
}

// Will keep information about the quantum elements that this EKMS is talking to
// This actually constitutes a quantum element with only a single link

/*
type QuantumElementInterface interface {
	GetQlID() qlElementId
}*/

func NewKMS(kmsUUID uuid.UUID, logOutput io.Writer, logLevel log.Level, logInJson bool, config *config.Config) (newKMS *KMS) {
	/*
	 * Setup logging
	 */
	//What level
	log.SetLevel(logLevel)
	// Where to send log out put
	log.SetOutput(logOutput)
	// and plain-text (standard) or json
	if !logInJson {
		log.SetFormatter(&log.TextFormatter{})
	} else {
		log.SetFormatter(&log.JSONFormatter{})
	}
	// print code function if level is set to Trace
	if logLevel == log.TraceLevel {
		log.SetReportCaller(true)
	} else {
		log.SetReportCaller(false)
	}

	tlsData := TlsData{
		TLS:      config.TLS,
		CaFile:   config.CAFile,
		CertFile: config.CertFile,
		KeyFile:  config.KeyFile,
	}

	ckmsAkmsClient := client.NewCkmsAkmsClient(config.AkmsURL)

	createdEKMS := &KMS{
		kmsName:             config.Name,
		kmsUUID:             kmsUUID,
		interComAddr:        config.InterComAddr,
		quantumAddress:      config.QuantumAddr,
		tlsData:             tlsData,
		remoteKMSMapping:    make(map[string]*RemoteKMS),
		quantumModules:      make(map[string]peers.QuantumModule),
		routingTable:        make(map[uuid.UUID]*Route),
		PKStore:             make(map[string]map[uuid.UUID]*PlatformKey),
		KmsPeers:            make(map[string]*peers.Peer),
		supportedKeyLengths: make(map[BitKeyLength]bool),
		eventBus:            event.NewEventBus(),
		CKMSAkmsClient:      ckmsAkmsClient,
	}

	createdEKMS.supportedKeyLengths[BitKeyLen256] = true

	// start the inter communication gRPC server
	go createdEKMS.startGRPC(config.InterComAddr, config.QuantumAddr, tlsData)

	// initialize from config
	err := initializePeers(createdEKMS, config)
	if err != nil {
		log.Fatalf("Failed to initialize peers: %s", err)
	}

	// Start the akmsCkmsReceiverServer
	if config.AkmsCkmsServerPort != "" {
		createdEKMS.CKMSAkmsServer = server.NewAKMSReceiver(config.AkmsCkmsServerPort, createdEKMS.eventBus, createdEKMS.GenerateAndSendKSAKey)
		log.Infof("Starting AKMS receiver server on port: %s", config.AkmsCkmsServerPort)
		go createdEKMS.CKMSAkmsServer.Serve()
	}

	return createdEKMS
}

func initializePeers(kms *KMS, config *config.Config) error {
	var qm peers.QuantumModule
	var err error
	for _, peer := range config.Peers {
		pqm := peer.QuantumModule
		switch qmt := peer.QuantumModule.QmType; qmt {
		case "emulated":
			qm = peers.NewEmulatedQuantumModule(pqm.Address, config.Id)
		case "etsi":
			qm, err = peers.NewETSI014HTTPQuantumModule(pqm.Address, config.Id, pqm.SlaveSAEID, pqm.MasterSAEID, config.CAFile, config.CertFile, config.KeyFile, config.ETSI14TLS, pqm.MasterMode)
			if err != nil {
				log.Fatalf("Failed to create ETSI QKD module: %s", err)
				return nil
			}
		default:
			log.Fatalf("Unknown type: %s for quantum module", qmt)
			return nil
		}

		err := kms.AddQuantumElement(qm)
		if err != nil {
			log.Fatalf("Failed to add quantum element: %s", err)
			return nil
		}

		var gRPCTransportCreds credentials.TransportCredentials
		if config.TLS {
			gRPCTransportCreds, err = util.GenerateGRPCClientTransportCredsWithTLS(config.CAFile, config.CertFile, config.KeyFile)
			if err != nil {
				log.Error(err)
				return nil
			}
		} else {
			gRPCTransportCreds = insecure.NewCredentials()
		}

		newPeerConn, err := grpc.Dial(peer.PeerInterComAddr, grpc.WithTransportCredentials(gRPCTransportCreds))
		if err != nil {
			return nil
		}

		client := &peers.GRPCClient{}
		switch pt := peer.Type; pt {
		case "danet":
			client.KmsTalkerClient = pbIC.NewKmsTalkerClient(newPeerConn)
		}

		p, err := kms.AddPeer(peer.PeerId, peer.PeerInterComAddr, qm, client)
		if err != nil {
			log.Fatalf("Failed to create a peer: %s", err)
			return nil
		}

		// TODO: check again; we might want to use this based on the quantum
		// module in use.
		if peer.Sync {
			go func() {
				time.Sleep(time.Second * 32)
				if err := p.SyncBulkKeys(); err != nil {
					log.Info("SYNC ERROR: ", err)
				} else {
					log.Info("SYNCED successfully!")
				}
			}()
		}
	}
	return nil
}

func (kms *KMS) startGRPC(interComAddr string, quantumAddress string, tlsData TlsData) {
	interKMSLis, err := net.Listen("tcp", interComAddr)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}

	var gRPCTransportCreds credentials.TransportCredentials
	if tlsData.TLS {
		gRPCTransportCreds, err = util.GenerateGRPCServerTransportCredsWithTLS(tlsData.CaFile, tlsData.CertFile, tlsData.KeyFile)
		if err != nil {
			log.Fatalf("unable to generate TLS creds: %v", err)
		}
	} else {
		gRPCTransportCreds = insecure.NewCredentials()
	}

	interKMSServer := grpc.NewServer(grpc.Creds(gRPCTransportCreds))
	healthCheck := health.NewServer()

	healthpb.RegisterHealthServer(interKMSServer, healthCheck)
	pbIC.RegisterKmsTalkerServer(interKMSServer, &kmsTalkerServer{
		keyNegotiationMap: make(map[uuid.UUID]*store.KmsKSElement),
		KMS:               kms,
	})

	quantumLis, err := net.Listen("tcp", quantumAddress)
	if err != nil {
		log.Fatalf("failed to listen: %v", err)
	}
	quantumServ := grpc.NewServer()
	pbQS.RegisterKmsQkdmCommunicationServiceServer(quantumServ, &quipSecServer{
		KMS: kms,
	})

	go func() {
		// set status to serving
		healthCheck.SetServingStatus(pbIC.KmsTalker_ServiceDesc.ServiceName, healthpb.HealthCheckResponse_SERVING)
		// TODO: add logic for adjusting health status based on operating status of
		// the services
		//    for{}
	}()

	log.Infof("quantum server listening at %v", quantumLis.Addr())
	go quantumServ.Serve(quantumLis) //nolint:errcheck

	log.Infof("inter KMS server listening at %v", interKMSLis.Addr())
	if err := interKMSServer.Serve(interKMSLis); err != nil {
		log.Fatalf("failed to serve: %v", err)
	}
}

func (kms *KMS) AddQuantumElement(qm peers.QuantumModule) error {
	kms.quantumModulesMutex.Lock()
	defer kms.quantumModulesMutex.Unlock()
	log.Infof("quantum module address: %s ", qm.Address())
	kms.quantumModules[qm.Address()] = qm
	return nil
}

func (kms *KMS) AddPeer(peerKmsId string, kmsPeerSocket string, servingQLE peers.QuantumModule, client *peers.GRPCClient) (*peers.Peer, error) {
	// check if peer exists
	if _, there := kms.KmsPeers[peerKmsId]; there {
		log.Errorf("Trying to add existing peer %s, with KMS ID %s", kmsPeerSocket, peerKmsId)
		return nil, fmt.Errorf("trying to add existing peer %s, with KMS ID %s", kmsPeerSocket, peerKmsId)
	}
	peer, err := peers.NewKmsPeer(peerKmsId, servingQLE, kmsPeerSocket, kms.interComAddr, client, kms.eventBus)
	if err != nil {
		return nil, err
	}
	peer.TcpSocketStr = kmsPeerSocket

	kms.kmsPeersMutex.Lock()
	defer kms.kmsPeersMutex.Unlock()
	kms.KmsPeers[peerKmsId] = peer

	// go peer.PeerHandler(kms.kmsName)
	return peer, nil
}

func (kms *KMS) AssignForwardingRoute(pId, pHop, nHop string, remoteKMS *RemoteKMS) error {
	pathId, err := uuid.Parse(pId)
	if err != nil {
		return fmt.Errorf("the given path id %s is no uuid; err = %w", pathId, err)
	}

	var previousHop *peers.Peer
	var nextHop *peers.Peer
	var ok bool
	if pHop != "" {
		previousHop, ok = kms.KmsPeers[pHop]
		if !ok {
			return fmt.Errorf("no peer found for %s", pHop)
		}
	}
	if nHop != "" {
		nextHop, ok = kms.KmsPeers[nHop]

		if !ok {
			return fmt.Errorf("no peer found for %s", nHop)
		}
	}

	tmpRoute := &Route{
		PathId:    pathId,
		Previous:  previousHop,
		Next:      nextHop,
		RemoteKMS: remoteKMS,
	}

	kms.routingTableMutex.Lock()
	// set the route within routing table
	kms.routingTable[pathId] = tmpRoute
	kms.routingTableMutex.Unlock()

	if tmpRoute.RemoteKMS != nil {
		kms.remoteKMSMappingMutex.Lock()
		if _, ok := kms.remoteKMSMapping[tmpRoute.RemoteKMS.Id]; !ok {
			kms.remoteKMSMapping[tmpRoute.RemoteKMS.Id] = tmpRoute.RemoteKMS
		}
		kms.remoteKMSMappingMutex.Unlock()
	}

	if tmpRoute.Previous == nil && tmpRoute.Next != nil && tmpRoute.RemoteKMS != nil {
		// generate pk key
		pk, err := crypto.Random256BitKey()
		if err != nil {
			return err
		}

		// generate process id
		processId := uuid.New()

		// update PKStore
		kms.PKStoreMutex.Lock()
		keys, ok := kms.PKStore[tmpRoute.RemoteKMS.Id]
		if !ok {
			kms.PKStore[tmpRoute.RemoteKMS.Id] = map[uuid.UUID]*PlatformKey{
				pathId: {
					Id:        pk.ID,
					Value:     pk.Key,
					ProcessId: processId.String(),
				},
			}
		} else {
			keys[pathId] = &PlatformKey{
				Id:        pk.ID,
				Value:     pk.Key,
				ProcessId: processId.String(),
			}
		}
		log.Debug("Current PKSTORE: ", kms.PKStore)
		kms.PKStoreMutex.Unlock()

		err = tmpRoute.Next.SendInitialPayloadBasedOnGRPCClient(pk, tmpRoute.PathId, processId, kms.kmsUUID.String(), remoteKMS.Address)
		if err != nil {
			log.Error(err)
			return err
		}
	}

	return nil
}

func (kms *KMS) GetSpecificPK(remoteKMSId string, keyId uuid.UUID) (*PlatformKey, error) {
	kms.PKStoreMutex.Lock()
	defer kms.PKStoreMutex.Unlock()

	keyIDs, ok := kms.PKStore[remoteKMSId]
	if !ok {
		return nil, fmt.Errorf("no entry for given KMS id: %s", remoteKMSId)
	}
	pk, ok := keyIDs[keyId]
	if !ok {
		return nil, fmt.Errorf("no key ready to use for given KMS id: %s", remoteKMSId)
	}
	delete(keyIDs, keyId)

	return pk, nil
}

func (kms *KMS) GetRandomItemFromPKStore(remoteKMSId string) (uuid.UUID, *PlatformKey, error) {
	kms.PKStoreMutex.Lock()
	defer kms.PKStoreMutex.Unlock()

	keyIds, ok := kms.PKStore[remoteKMSId]
	if !ok {
		// TODO: return proper error
		log.Errorf("error with he following pk store: %v", kms.PKStore)
		return uuid.Nil, nil, fmt.Errorf("path not found for: %s", remoteKMSId)
	}
	return util.RandomItemFromMapAndRemove(keyIds)
}

func (kms *KMS) GetRemoteKMS(remoteKMSId string) (*RemoteKMS, error) {
	kms.remoteKMSMappingMutex.RLock()
	defer kms.remoteKMSMappingMutex.RUnlock()

	remoteKMS, ok := kms.remoteKMSMapping[remoteKMSId]
	if !ok {
		return nil, fmt.Errorf("Address for remoteKMS with id %s not found", remoteKMSId)
	}
	return remoteKMS, nil
}

// NOTE: address/remoteid still have to decide.
func (kms *KMS) GenerateAndSendKSAKey(remoteKMSId string, pathId uuid.UUID, requestId string, number int) error {
	if number < 1 {
		log.Errorf("number must be positive and at least 1, provided: %d\n", number)
		return fmt.Errorf("number must be positive and at least 1, provided: %d", number)
	}

	log.Infof("KMS store: %v", kms.remoteKMSMapping)
	remoteKMS, err := kms.GetRemoteKMS(remoteKMSId)
	if err != nil {
		log.Error(err)
		return err
	}

	pk, err := kms.GetSpecificPK(remoteKMSId, pathId)
	if err != nil {
		log.Error(err)
		return err
	}

	keys := make([]*pbIC.Key, number)
	akmsKSAKeys := make([]client.KSAKey, number)
	cryptoAlgo := crypto.NewAES()
	for i := 0; i < number; i++ {
		// generate ksa key
		ksaKeyId := uuid.New()
		ksaKey, err := crypto.Random256BitKey()
		if err != nil {
			log.Error(err)
			return err
		}

		// encrypt the key
		nonce, encryptedKSAKey, err := cryptoAlgo.Encrypt(ksaKey.Key, pk.Value)
		if err != nil {
			log.Error(err)
			return err
		}

		ksaKeyAsString := base64.StdEncoding.EncodeToString(ksaKey.Key)
		encryptedKSAKeyAsString := base64.StdEncoding.EncodeToString(encryptedKSAKey)
		nonceAsString := base64.StdEncoding.EncodeToString(nonce)

		keys[i] = &pbIC.Key{
			Id:    ksaKeyId.String(),
			Nonce: nonceAsString,
			Key:   encryptedKSAKeyAsString,
		}

		akmsKSAKeys[i] = client.KSAKey{
			KeyID: ksaKeyId.String(),
			Key:   ksaKeyAsString,
		}
	}

	// TODO: move this somewhere else!
	// send to remote
	var gRPCTransportCreds credentials.TransportCredentials
	if kms.tlsData.TLS {
		gRPCTransportCreds, err = util.GenerateGRPCClientTransportCredsWithTLS(kms.tlsData.CaFile, kms.tlsData.CertFile, kms.tlsData.KeyFile)
		if err != nil {
			log.Fatalf("unable to generate TLS creds: %v", err)
		}
	} else {
		gRPCTransportCreds = insecure.NewCredentials()
	}

	remoteConn, err := grpc.Dial(remoteKMS.Address, grpc.WithTransportCredentials(gRPCTransportCreds))
	if err != nil {
		log.Error(err)
		return err
	}
	remoteClient := pbIC.NewKmsTalkerClient(remoteConn)
	// Start the negotiation process for a transport key

	ctx, cancel := context.WithTimeout(context.Background(), time.Second)
	// create a new context with some metadata
	md := metadata.Pairs("hostname", kms.kmsName)
	ctx = metadata.NewOutgoingContext(ctx, md)
	defer cancel()
	_, err = remoteClient.KeyDelivery(ctx, &pbIC.KeyDeliveryRequest{
		KeyId:     pk.Id.String(),
		RequestId: requestId,
		KmsId:     kms.kmsUUID.String(),
		// TODO: change if multiple keys are generated
		Key: keys,
	})
	if err != nil {
		log.Error(err)
		return err
	}

	// Use the real processID when we know what it is
	err = kms.CKMSAkmsClient.SendKSAKeys(requestId, pk.ProcessId, akmsKSAKeys)
	if err != nil {
		log.Error(err)
		return err
	}

	return nil
}

func (kms *KMS) EventBus() *event.EventBus {
	return kms.eventBus
}

// TODO/XXX error handling.
func (kms *KMS) RemovePeer(kmsPeerSocket string) {
	if _, there := kms.KmsPeers[kmsPeerSocket]; there {
		// peer.quit <- true
		delete(kms.KmsPeers, kmsPeerSocket)
		return
	}
	log.Errorf("%s: Can not find a peer with socket: %s", kms.kmsName, kmsPeerSocket)
}

func (kms *KMS) FindPeerUuid(lookup uuid.UUID) (peer *peers.Peer) {
	if kms.KmsPeers != nil {
		for _, peer = range kms.KmsPeers {
			if peer.Id() == lookup {
				return peer
			}
		}
	}

	return nil
}

func (kms *KMS) RoutingTableDeepCopy() map[uuid.UUID]*Route {
	routingTableCopy := make(map[uuid.UUID]*Route, len(kms.KmsPeers))

	kms.routingTableMutex.Lock()
	defer kms.routingTableMutex.Unlock()
	for k, v := range kms.routingTable {
		routingTableCopy[k] = v
	}

	return routingTableCopy
}

func (kms *KMS) PeersDeepCopy() map[string]*peers.Peer {
	peersCopy := make(map[string]*peers.Peer, len(kms.KmsPeers))

	kms.kmsPeersMutex.Lock()
	defer kms.kmsPeersMutex.Unlock()
	for k, v := range kms.KmsPeers {
		peersCopy[k] = v
	}

	return peersCopy
}