pkiapi/internal/storage/mongo.go

package storage

import (
	"context"
	"crypto/rsa"
	"encoding/base64"
	"time"

	"github.com/stef/pkiapi/internal/pki"
	"go.mongodb.org/mongo-driver/bson"
	"go.mongodb.org/mongo-driver/mongo"
	"go.mongodb.org/mongo-driver/mongo/options"
)

// MongoStore gère le stockage des certificats dans MongoDB
type MongoStore struct {
	client     *mongo.Client
	collection *mongo.Collection
}

// CertificateDoc représente un certificat stocké dans MongoDB
type CertificateDoc struct {
	ID         string `bson:"_id"`
	Subject    string `bson:"subject"`
	Issuer     string `bson:"issuer"`
	NotBefore  time.Time `bson:"not_before"`
	NotAfter   time.Time `bson:"not_after"`
	IsCA       bool `bson:"is_ca"`
	Revoked    bool `bson:"revoked"`
	Cert       string `bson:"cert"` // Base64 encoded certificate
	PrivateKey string `bson:"private_key"` // Base64 encoded private key (for all certificates)
	CreatedAt  time.Time `bson:"created_at"`
}

// NewMongoStore crée une connexion MongoDB et retourne un MongoStore
func NewMongoStore(mongoURI string, dbName string) (*MongoStore, error) {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	client, err := mongo.Connect(ctx, options.Client().ApplyURI(mongoURI))
	if err != nil {
		return nil, err
	}

	// Tester la connexion
	ctx, cancel = context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()
	if err := client.Ping(ctx, nil); err != nil {
		return nil, err
	}

	collection := client.Database(dbName).Collection("certificates")

	// Créer un index sur l'ID
	indexModel := mongo.IndexModel{
		Keys: bson.D{{Key: "_id", Value: 1}},
	}
	_, err = collection.Indexes().CreateOne(ctx, indexModel)
	if err != nil {
		return nil, err
	}

	return &MongoStore{
		client:     client,
		collection: collection,
	}, nil
}

// SaveCertificate sauvegarde un certificat dans MongoDB
func (m *MongoStore) SaveCertificate(id string, cert *pki.Certificate) error {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	doc := CertificateDoc{
		ID:        id,
		Subject:   cert.Subject,
		Issuer:    cert.Issuer,
		NotBefore: cert.NotBefore,
		NotAfter:  cert.NotAfter,
		IsCA:      cert.IsCA,
		Revoked:   cert.Revoked,
		CreatedAt: time.Now(),
	}

	// Encoder le certificat en base64
	if cert.Cert != nil {
		doc.Cert = base64.StdEncoding.EncodeToString(cert.Cert.Raw)
	}

	// Encoder la clé privée en base64 (pour tous les certificats)
	if cert.PrivateKey != nil {
		privKeyBytes, err := marshalPrivateKey(cert.PrivateKey)
		if err != nil {
			return err
		}
		doc.PrivateKey = base64.StdEncoding.EncodeToString(privKeyBytes)
	}

	_, err := m.collection.UpdateOne(
		ctx,
		bson.M{"_id": id},
		bson.M{"$set": doc},
		options.Update().SetUpsert(true),
	)
	return err
}

// GetCertificate récupère un certificat depuis MongoDB
func (m *MongoStore) GetCertificate(id string) (*pki.Certificate,error) {
	ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
	defer cancel()

	var doc CertificateDoc
	err := m.collection.FindOne(ctx, bson.M{"_id": id}).Decode(&doc)
	if err == mongo.ErrNoDocuments {
		return nil, ErrNotFound
	}
	if err != nil {
		return nil, err
	}

	// Décoder le certificat
	cert := &pki.Certificate{
		ID:        doc.ID,
		Subject:   doc.Subject,
		Issuer:    doc.Issuer,
		NotBefore: doc.NotBefore,
		NotAfter:  doc.NotAfter,
		IsCA:      doc.IsCA,
		Revoked:   doc.Revoked,
	}

	// Décoder le certificat X.509
	if doc.Cert != "" {
		certBytes, err := base64.StdEncoding.DecodeString(doc.Cert)
		if err != nil {
			return nil, err
		}
		parsedCert, err := parseCertificate(certBytes)
		if err != nil {
			return nil, err
		}
		cert.Cert = parsedCert
	}

	// Décoder la clé privée
	if doc.PrivateKey != "" {
		privKeyBytes, err := base64.StdEncoding.DecodeString(doc.PrivateKey)
		if err != nil {
			return nil, err
		}
		privKey, err := unmarshalPrivateKey(privKeyBytes)
		if err != nil {
			return nil, err
		}
		// Type assertion pour convertir interface{} en *rsa.PrivateKey
		rsaKey, ok := privKey.(*rsa.PrivateKey)
		if !ok {
			// Si ce n'est pas une clé RSA, on la laisse nil (CAs sans clé privée)
			cert.PrivateKey = nil
		} else {
			cert.PrivateKey = rsaKey
		}
	}

	return cert, nil
}

// ListCertificates retourne tous les certificats depuis MongoDB
func (m *MongoStore) ListCertificates() []*pki.Certificate {
	ctx, cancel := context.WithTimeout(context.Background(), 10*time.Second)
	defer cancel()

	cursor, err := m.collection.Find(ctx, bson.M{})
	if err != nil {
		return []*pki.Certificate{}
	}
	defer cursor.Close(ctx)

	var results []*pki.Certificate
	for cursor.Next(ctx) {
		var doc CertificateDoc
		if err := cursor.Decode(&doc); err != nil {
			continue
		}

		cert := &pki.Certificate{
			ID:        doc.ID,
			Subject:   doc.Subject,
			Issuer:    doc.Issuer,
			NotBefore: doc.NotBefore,
			NotAfter:  doc.NotAfter,
			IsCA:      doc.IsCA,
			Revoked:   doc.Revoked,
		}

		// Décoder le certificat X.509
		if doc.Cert != "" {
			if certBytes, err := base64.StdEncoding.DecodeString(doc.Cert); err == nil {
				if parsedCert, err := parseCertificate(certBytes); err == nil && parsedCert != nil {
					cert.Cert = parsedCert
				}
			}
		}

		results = append(results, cert)
	}

	return results
}

// Close ferme la connexion MongoDB
func (m *MongoStore) Close() error {
	if m.client != nil {
		ctx, cancel := context.WithTimeout(context.Background(), 5*time.Second)
		defer cancel()
		return m.client.Disconnect(ctx)
	}
	return nil
}