keymanager/server/internal/services/servers.go

package services

import (
	"context"
	"crypto/rand"
	"crypto/sha256"
	"encoding/hex"
	"fmt"
	"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"
	"go.mongodb.org/mongo-driver/v2/mongo/options"
)

func generateToken(n int) (string, error) {
	b := make([]byte, n)
	if _, err := rand.Read(b); err != nil {
		return "", err
	}
	return hex.EncodeToString(b), nil
}

func HashToken(token string) string {
	sum := sha256.Sum256([]byte(token))
	return hex.EncodeToString(sum[:])
}

func CreateServer() (*models.Server, string, error) {
	token, err := generateToken(32)
	if err != nil {
		return nil, "", err
	}
	expires := time.Now().Add(time.Hour)
	s := &models.Server{
		ServerID:      uuid.NewString(),
		PreRegToken:   token,
		PreRegExpires: &expires,
		Status:        "pending",
		CreatedAt:     time.Now(),
	}
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	_, err = db.Col("servers").InsertOne(ctx, s)
	if err != nil {
		return nil, "", err
	}
	return s, token, nil
}

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

	var s models.Server
	err := db.Col("servers").FindOne(ctx, bson.M{"server_id": serverID}).Decode(&s)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func GetServerByPreRegToken(token string) (*models.Server, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var s models.Server
	err := db.Col("servers").FindOne(ctx, bson.M{
		"pre_reg_token":   token,
		"pre_reg_expires": bson.M{"$gt": time.Now()},
	}).Decode(&s)
	if err != nil {
		return nil, err
	}
	return &s, nil
}

func RegisterServer(serverID, preRegToken, hostname, ipAddress, osInfo string) (string, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var s models.Server
	err := db.Col("servers").FindOne(ctx, bson.M{
		"server_id":       serverID,
		"pre_reg_token":   preRegToken,
		"pre_reg_expires": bson.M{"$gt": time.Now()},
	}).Decode(&s)
	if err != nil {
		return "", fmt.Errorf("invalid or expired pre-registration token")
	}

	agentToken, err := generateToken(32)
	if err != nil {
		return "", err
	}
	tokenHash := HashToken(agentToken)
	now := time.Now()

	_, err = db.Col("servers").UpdateOne(ctx,
		bson.M{"server_id": serverID},
		bson.M{"$set": bson.M{
			"hostname":         hostname,
			"ip_address":       ipAddress,
			"os_info":          osInfo,
			"agent_token_hash": tokenHash,
			"status":           "active",
			"last_seen":        now,
			"pre_reg_token":    "",
			"pre_reg_expires":  nil,
		}},
	)
	if err != nil {
		return "", err
	}
	return agentToken, nil
}

func ValidateAgentToken(serverID, agentToken string) (*models.Server, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	tokenHash := HashToken(agentToken)
	var s models.Server
	err := db.Col("servers").FindOne(ctx, bson.M{
		"server_id":        serverID,
		"agent_token_hash": tokenHash,
	}).Decode(&s)
	if err != nil {
		return nil, fmt.Errorf("invalid agent token")
	}
	return &s, nil
}

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

	now := time.Now()
	_, err := db.Col("servers").UpdateOne(ctx,
		bson.M{"server_id": serverID},
		bson.M{"$set": bson.M{"last_seen": now, "status": "active"}},
	)
	return err
}

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

	opts := options.Find().SetSort(bson.D{{Key: "created_at", Value: -1}})
	cursor, err := db.Col("servers").Find(ctx, bson.M{}, opts)
	if err != nil {
		return nil, err
	}
	defer cursor.Close(ctx)

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

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

	_, err := db.Col("servers").DeleteOne(ctx, bson.M{"server_id": serverID})
	if err != nil {
		return err
	}
	// Also remove assignments
	_, err = db.Col("assignments").DeleteMany(ctx, bson.M{"server_id": serverID})
	return err
}

func MarkOfflineServers(threshold time.Duration) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	cutoff := time.Now().Add(-threshold)
	_, err := db.Col("servers").UpdateMany(ctx,
		bson.M{
			"status":    "active",
			"last_seen": bson.M{"$lt": cutoff},
		},
		bson.M{"$set": bson.M{"status": "offline"}},
	)
	return err
}