keymanager/server/internal/services/keys.go

package services

import (
	"context"
	"crypto/md5"
	"encoding/base64"
	"fmt"
	"strings"
	"time"

	"github.com/google/uuid"
	"github.com/mrhid6/keymanager/server/internal/db"
	"github.com/mrhid6/keymanager/server/internal/models"
	"go.mongodb.org/mongo-driver/v2/bson"
)

func computeFingerprint(pubKey string) string {
	parts := strings.Fields(pubKey)
	if len(parts) < 2 {
		return ""
	}
	raw, err := base64.StdEncoding.DecodeString(parts[1])
	if err != nil {
		return ""
	}
	sum := md5.Sum(raw)
	var pairs []string
	for _, b := range sum {
		pairs = append(pairs, fmt.Sprintf("%02x", b))
	}
	return "MD5:" + strings.Join(pairs, ":")
}

func CreateKey(label, publicKey, source, generatedByServerID string) (*models.Key, error) {
	key := &models.Key{
		KeyID:               uuid.NewString(),
		Label:               label,
		PublicKey:           publicKey,
		Fingerprint:         computeFingerprint(publicKey),
		Source:              source,
		GeneratedByServerID: generatedByServerID,
		CreatedAt:           time.Now(),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err := db.Col("keys").InsertOne(ctx, key)
	if err != nil {
		return nil, err
	}
	return key, nil
}

func GetKey(keyID string) (*models.Key, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var key models.Key
	err := db.Col("keys").FindOne(ctx, bson.M{"key_id": keyID}).Decode(&key)
	if err != nil {
		return nil, err
	}
	return &key, nil
}

type KeyWithCount struct {
	models.Key    `bson:",inline"`
	AssignedCount int `bson:"-" json:"assigned_count"`
}

func ListKeys() ([]KeyWithCount, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("keys").Find(ctx, bson.M{})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var keys []models.Key
	if err := cursor.All(ctx, &keys); err != nil {
		return nil, err
	}

	result := make([]KeyWithCount, 0, len(keys))
	for _, k := range keys {
		count, _ := db.Col("assignments").CountDocuments(ctx, bson.M{
			"key_id":     k.KeyID,
			"revoked_at": nil,
		})
		result = append(result, KeyWithCount{Key: k, AssignedCount: int(count)})
	}
	return result, nil
}

func DeleteKey(keyID string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	if _, err := db.Col("keys").DeleteOne(ctx, bson.M{"key_id": keyID}); err != nil {
		return err
	}
	_, err := db.Col("assignments").DeleteMany(ctx, bson.M{"key_id": keyID})
	return err
}

func AssignKey(keyID, serverID string) (*models.Assignment, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	// Check if already assigned and active
	var existing models.Assignment
	err := db.Col("assignments").FindOne(ctx, bson.M{
		"key_id":     keyID,
		"server_id":  serverID,
		"revoked_at": nil,
	}).Decode(&existing)
	if err == nil {
		return &existing, nil
	}

	a := &models.Assignment{
		KeyID:      keyID,
		ServerID:   serverID,
		AssignedAt: time.Now(),
	}
	_, err = db.Col("assignments").InsertOne(ctx, a)
	if err != nil {
		return nil, err
	}
	return a, nil
}

func RevokeAssignment(keyID, serverID string) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	now := time.Now()
	_, err := db.Col("assignments").UpdateOne(ctx,
		bson.M{"key_id": keyID, "server_id": serverID, "revoked_at": nil},
		bson.M{"$set": bson.M{"revoked_at": now}},
	)
	return err
}

func GetAssignmentsForKey(keyID string) ([]models.Assignment, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("assignments").Find(ctx, bson.M{"key_id": keyID, "revoked_at": nil})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var assignments []models.Assignment
	if err := cursor.All(ctx, &assignments); err != nil {
		return nil, err
	}
	return assignments, nil
}

type AssignmentWithServer struct {
	models.Assignment `bson:",inline"`
	Server            *models.Server `json:"server,omitempty"`
}

func GetAssignmentsWithServers(keyID string) ([]AssignmentWithServer, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("assignments").Find(ctx, bson.M{"key_id": keyID})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var assignments []models.Assignment
	if err := cursor.All(ctx, &assignments); err != nil {
		return nil, err
	}

	result := make([]AssignmentWithServer, 0, len(assignments))
	for _, a := range assignments {
		item := AssignmentWithServer{Assignment: a}
		var srv models.Server
		if err := db.Col("servers").FindOne(ctx, bson.M{"server_id": a.ServerID}).Decode(&srv); err == nil {
			item.Server = &srv
		}
		result = append(result, item)
	}
	return result, nil
}

type AssignmentWithKey struct {
	models.Assignment `bson:",inline"`
	Key               *models.Key `json:"key,omitempty"`
}

func GetAssignmentsWithKeysForServer(serverID string) ([]AssignmentWithKey, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cursor, err := db.Col("assignments").Find(ctx, bson.M{"server_id": serverID})
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

	var assignments []models.Assignment
	if err := cursor.All(ctx, &assignments); err != nil {
		return nil, err
	}

	result := make([]AssignmentWithKey, 0, len(assignments))
	for _, a := range assignments {
		var key models.Key
		if err := db.Col("keys").FindOne(ctx, bson.M{"key_id": a.KeyID}).Decode(&key); err != nil {
			continue
		}
		result = append(result, AssignmentWithKey{Assignment: a, Key: &key})
	}
	return result, nil
}