Merge pull request #8 from cpp-cyber/development

Development

Author: MaxwellCaron

internal/api/handlers/cloning_handler.go

@@ -11,6 +11,7 @@ import (
 	"github.com/cpp-cyber/proclone/internal/ldap"
 	"github.com/cpp-cyber/proclone/internal/proxmox"
 	"github.com/cpp-cyber/proclone/internal/tools"
+	"github.com/cpp-cyber/proclone/internal/tools/sse"
 	"github.com/gin-contrib/sessions"
 	"github.com/gin-gonic/gin"
 )
@@ -88,16 +89,54 @@ func (ch *CloningHandler) CloneTemplateHandler(c *gin.Context) {
 		return
 	}
 
+	// Check for existing deployments before starting SSE
+	targetPoolName := fmt.Sprintf("%s_%s", req.Template, username)
+	isValid, err := ch.Service.ValidateCloneRequest(targetPoolName, username)
+	if err != nil {
+		log.Printf("Error validating deployment for user %s: %v", username, err)
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error":   "Failed to validate existing deployments",
+			"details": err.Error(),
+		})
+		return
+	}
+	if !isValid {
+		log.Printf("Template %s is already deployed for user %s or they have exceeded deployment limits", req.Template, username)
+		c.JSON(http.StatusConflict, gin.H{
+			"error":   "Deployment not allowed",
+			"details": fmt.Sprintf("Template %s is already deployed for %s or they have exceeded the maximum of 5 deployed pods", req.Template, username),
+		})
+		return
+	}
+
+	// Create new sse object for streaming
+	sseWriter, err := sse.NewWriter(c.Writer)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error":   "Failed to initialize SSE",
+			"details": err.Error(),
+		})
+		return
+	}
+
+	sseWriter.Send(
+		cloning.ProgressMessage{
+			Message:  "Retrieving template information",
+			Progress: 0,
+		},
+	)
+
 	// Create the cloning request using the new format
 	cloneReq := cloning.CloneRequest{
 		Template:                 req.Template,
-		CheckExistingDeployments: true, // Check for existing deployments for single user clones
+		CheckExistingDeployments: false, // Already checked above
 		Targets: []cloning.CloneTarget{
 			{
 				Name:    username,
 				IsGroup: false,
 			},
 		},
+		SSE: sseWriter,
 	}
 
 	if err := ch.Service.CloneTemplate(cloneReq); err != nil {
@@ -144,16 +183,27 @@ func (ch *CloningHandler) AdminCloneTemplateHandler(c *gin.Context) {
 		})
 	}
 
+	// Create new sse object for streaming
+	sseWriter, err := sse.NewWriter(c.Writer)
+	if err != nil {
+		c.JSON(http.StatusInternalServerError, gin.H{
+			"error":   "Failed to initialize SSE",
+			"details": err.Error(),
+		})
+		return
+	}
+
 	// Create clone request
 	cloneReq := cloning.CloneRequest{
 		Template:                 req.Template,
 		Targets:                  targets,
 		CheckExistingDeployments: false,
 		StartingVMID:             req.StartingVMID,
+		SSE:                      sseWriter,
 	}
 
 	// Perform clone operation
-	err := ch.Service.CloneTemplate(cloneReq)
+	err = ch.Service.CloneTemplate(cloneReq)
 	if err != nil {
 		log.Printf("Admin %s encountered error while bulk cloning template: %v", username, err)
 		c.JSON(http.StatusInternalServerError, gin.H{
@@ -264,7 +314,7 @@ func (ch *CloningHandler) GetPodsHandler(c *gin.Context) {
 
 	// Loop through the user's deployed pods and add template information
 	for i := range pods {
-		templateName := strings.Replace(pods[i].Name[5:], fmt.Sprintf("_%s", username), "", 1)
+		templateName := strings.Replace(strings.ToLower(pods[i].Name[5:]), fmt.Sprintf("_%s", strings.ToLower(username)), "", 1)
 		templateInfo, err := ch.Service.DatabaseService.GetTemplateInfo(templateName)
 		if err != nil {
 			log.Printf("Error retrieving template info for pod %s: %v", pods[i].Name, err)

internal/api/handlers/dashboard_handler.go

@@ -91,7 +91,7 @@ func (dh *DashboardHandler) GetUserDashboardStatsHandler(c *gin.Context) {
 
 	// Loop through the user's deployed pods and add template information
 	for i := range pods {
-		templateName := strings.Replace(pods[i].Name[5:], fmt.Sprintf("_%s", username), "", 1)
+		templateName := strings.Replace(strings.ToLower(pods[i].Name[5:]), fmt.Sprintf("_%s", strings.ToLower(username)), "", 1)
 		templateInfo, err := dh.cloningHandler.Service.DatabaseService.GetTemplateInfo(templateName)
 		if err != nil {
 			log.Printf("Error retrieving template info for pod %s: %v", pods[i].Name, err)

internal/api/middleware/authorization.go

@@ -63,12 +63,14 @@ func Logout(c *gin.Context) {
 
 func CORSMiddleware(fqdn string) gin.HandlerFunc {
 	return func(c *gin.Context) {
-		c.Writer.Header().Set("Content-Type", "application/json")
+		c.Writer.Header().Set("Content-Type", "application/json; text/event-stream")
 		c.Writer.Header().Set("Access-Control-Allow-Origin", fqdn)
 		c.Writer.Header().Set("Access-Control-Max-Age", "86400")
 		c.Writer.Header().Set("Access-Control-Allow-Credentials", "true")
 		c.Writer.Header().Set("Access-Control-Allow-Methods", "POST, GET, OPTIONS, PUT, DELETE")
 		c.Writer.Header().Set("Access-Control-Allow-Headers", "Content-Type, Content-Length, Accept-Encoding, X-CSRF-Token, Authorization, Origin")
+		c.Writer.Header().Set("Cache-Control", "no-cache")
+		c.Writer.Header().Set("Connection", "keep-alive")
 
 		if c.Request.Method == "OPTIONS" {
 			c.AbortWithStatus(200)

internal/cloning/cloning_service.go

@@ -5,7 +5,7 @@ import (
 	"fmt"
 	"log"
 	"os"
-	"strings"
+	"regexp"
 	"time"
 
 	"github.com/cpp-cyber/proclone/internal/ldap"
@@ -85,11 +85,11 @@ func (cs *CloningService) CloneTemplate(req CloneRequest) error {
 	// 3. Identify router and other VMs
 	var router *proxmox.VM
 	var templateVMs []proxmox.VM
+	routerPattern := regexp.MustCompile(`(?i)(router|pfsense|vyos)`)
 
 	for _, vm := range templatePool {
 		// Check to see if this VM is the router
-		lowerVMName := strings.ToLower(vm.Name)
-		if strings.Contains(lowerVMName, "router") || strings.Contains(lowerVMName, "pfsense") {
+		if routerPattern.MatchString(vm.Name) {
 			router = &proxmox.VM{
 				Name: vm.Name,
 				Node: vm.NodeName,
@@ -166,6 +166,13 @@ func (cs *CloningService) CloneTemplate(req CloneRequest) error {
 	}
 
 	// 7. Clone targets to proxmox
+	req.SSE.Send(
+		ProgressMessage{
+			Message:  "Cloning VMs",
+			Progress: 10,
+		},
+	)
+
 	for _, target := range req.Targets {
 		// Find best node per target
 		bestNode, err := cs.ProxmoxService.FindBestNode()
@@ -186,9 +193,16 @@ func (cs *CloningService) CloneTemplate(req CloneRequest) error {
 		if err != nil {
 			errors = append(errors, fmt.Sprintf("failed to clone router VM for %s: %v", target.Name, err))
 		} else {
+			// Determine router type
+			routerType, err := cs.getRouterType(*router)
+			if err != nil {
+				errors = append(errors, fmt.Sprintf("failed to get router type for %s: %v", target.Name, err))
+			}
+
 			// Store router info for later operations
 			clonedRouters = append(clonedRouters, RouterInfo{
 				TargetName: target.Name,
+				RouterType: routerType,
 				PodNumber:  target.PodNumber,
 				Node:       bestNode,
 				VMID:       target.VMIDs[0],
@@ -251,6 +265,12 @@ func (cs *CloningService) CloneTemplate(req CloneRequest) error {
 	}
 
 	// 10. Start all routers and wait for them to be running
+	req.SSE.Send(
+		ProgressMessage{
+			Message:  "Starting routers",
+			Progress: 25,
+		},
+	)
 	log.Printf("Starting %d routers", len(clonedRouters))
 	for _, routerInfo := range clonedRouters {
 		// Wait for router disk to be available
@@ -278,6 +298,13 @@ func (cs *CloningService) CloneTemplate(req CloneRequest) error {
 	}
 
 	// 11. Configure all pod routers (separate step after all routers are running)
+	req.SSE.Send(
+		ProgressMessage{
+			Message:  "Configuring pod routers",
+			Progress: 33,
+		},
+	)
+
 	log.Printf("Configuring %d pod routers", len(clonedRouters))
 	for _, routerInfo := range clonedRouters {
 		// Double-check that router is still running before configuration
@@ -288,12 +315,20 @@ func (cs *CloningService) CloneTemplate(req CloneRequest) error {
 		}
 
 		log.Printf("Configuring pod router for %s (Pod: %d, VMID: %d)", routerInfo.TargetName, routerInfo.PodNumber, routerInfo.VMID)
-		err = cs.configurePodRouter(routerInfo.PodNumber, routerInfo.Node, routerInfo.VMID)
+		err = cs.configurePodRouter(routerInfo.PodNumber, routerInfo.Node, routerInfo.VMID, routerInfo.RouterType)
 		if err != nil {
 			errors = append(errors, fmt.Sprintf("failed to configure pod router for %s: %v", routerInfo.TargetName, err))
 		}
 	}
 
+	// Router configuration complete - update progress
+	req.SSE.Send(
+		ProgressMessage{
+			Message:  "Finalizing deployment",
+			Progress: 90,
+		},
+	)
+
 	// 12. Set permissions on the pool to the user/group
 	for _, target := range req.Targets {
 		err = cs.ProxmoxService.SetPoolPermission(target.PoolName, target.Name, target.IsGroup)
@@ -308,6 +343,14 @@ func (cs *CloningService) CloneTemplate(req CloneRequest) error {
 		errors = append(errors, fmt.Sprintf("failed to increment template deployments for %s: %v", req.Template, err))
 	}
 
+	// Final completion message
+	req.SSE.Send(
+		ProgressMessage{
+			Message:  "Template cloning completed!",
+			Progress: 100,
+		},
+	)
+
 	// Handle errors and cleanup if necessary
 	if len(errors) > 0 {
 		cs.cleanupFailedClones(createdPools)
@@ -357,9 +400,7 @@ func (cs *CloningService) DeletePod(pod string) error {
 					VMID: vm.VmId,
 				})
 				stoppedCount++
-			} else {
 			}
-		} else {
 		}
 	}