chore: cleanup firewalls in cleanup script

Author: nadim-az

cmd/cleanup/main.go

@@ -37,19 +37,37 @@ func main() {
 
 	doClient := digitalocean.NewGodoClient(*token)
 
-	logger.Info("Starting droplet cleanup",
+	logger.Info("Starting resource cleanup",
 		zap.String("prefix", *namePrefix),
 		zap.String("long_running_tag", *longRunning),
 		zap.Bool("dry_run", *dryRun))
 
-	if err := cleanupDroplets(ctx, doClient, logger); err != nil {
+	skippedPrefixes, err := getSkippedDropletPrefixes(ctx, doClient, logger)
+	if err != nil {
+		logger.Fatal("Failed to get skipped droplet prefixes", zap.Error(err))
+	}
+
+	if err := cleanupDroplets(ctx, doClient, logger, skippedPrefixes); err != nil {
 		logger.Fatal("Failed to cleanup droplets", zap.Error(err))
 	}
 
-	logger.Info("Droplet cleanup completed successfully")
+	if err := cleanupFirewalls(ctx, doClient, logger, skippedPrefixes); err != nil {
+		logger.Fatal("Failed to cleanup firewalls", zap.Error(err))
+	}
+
+	logger.Info("Resource cleanup completed successfully")
+}
+
+func extractPrefix(dropletName string) string {
+	// Extract prefix pattern like "petri-ib-XXXXXX" from "petri-ib-XXXXXX-validator-0"
+	parts := strings.Split(dropletName, "-")
+	if len(parts) >= 3 {
+		return strings.Join(parts[:3], "-")
+	}
+	return dropletName
 }
 
-func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *zap.Logger) error {
+func getSkippedDropletPrefixes(ctx context.Context, client digitalocean.DoClient, logger *zap.Logger) (map[string]bool, error) {
 	opts := &godo.ListOptions{
 		Page:    1,
 		PerPage: 200,
@@ -59,7 +77,7 @@ func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *
 	for {
 		droplets, err := client.ListDroplets(ctx, opts)
 		if err != nil {
-			return fmt.Errorf("failed to list droplets: %w", err)
+			return nil, fmt.Errorf("failed to list droplets: %w", err)
 		}
 
 		allDroplets = append(allDroplets, droplets...)
@@ -71,15 +89,21 @@ func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *
 		opts.Page++
 	}
 
-	logger.Info("Retrieved droplets", zap.Int("total_count", len(allDroplets)))
-
-	var dropletsToDelete []godo.Droplet
+	skippedPrefixes := make(map[string]bool)
 	now := time.Now()
+
 	for _, droplet := range allDroplets {
+		// Skip non-petri droplets
 		if !strings.HasPrefix(droplet.Name, *namePrefix) {
 			continue
 		}
 
+		prefix := extractPrefix(droplet.Name)
+		if skippedPrefixes[prefix] {
+			continue
+		}
+
+		// Skip droplets with the LONG_RUNNING tag
 		hasLongRunningTag := false
 		for _, tag := range droplet.Tags {
 			if tag == *longRunning {
@@ -89,29 +113,69 @@ func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *
 		}
 
 		if hasLongRunningTag {
-			logger.Info("Skipping droplet with long-running tag",
-				zap.String("name", droplet.Name),
-				zap.Int("id", droplet.ID))
+			prefix := extractPrefix(droplet.Name)
+			skippedPrefixes[prefix] = true
+			logger.Info("Marking droplet as skipped (long-running)",
+				zap.String("droplet_name", droplet.Name))
 			continue
 		}
 
-		// Parse the creation time and skip if created in the last 30 minutes
+		// Skip if droplet was created in the last 30 minutes
 		createdAt, err := time.Parse(time.RFC3339, droplet.Created)
 		if err != nil {
-			logger.Warn("Failed to parse droplet creation time, skipping",
+			logger.Error("Failed to parse droplet creation time, marking as skipped",
 				zap.String("name", droplet.Name),
-				zap.Int("id", droplet.ID),
 				zap.String("created_at", droplet.Created),
 				zap.Error(err))
+			prefix := extractPrefix(droplet.Name)
+			skippedPrefixes[prefix] = true
 			continue
 		}
 
 		if now.Sub(createdAt) < 30*time.Minute {
-			logger.Info("Skipping recently created droplet",
-				zap.String("name", droplet.Name),
-				zap.Int("id", droplet.ID),
-				zap.String("created_at", droplet.Created),
-				zap.Duration("age", now.Sub(createdAt)))
+			prefix := extractPrefix(droplet.Name)
+			skippedPrefixes[prefix] = true
+			logger.Info("Marking droplet as skipped (too recent)",
+				zap.String("droplet_name", droplet.Name),
+				zap.String("created_at", droplet.Created))
+		}
+	}
+
+	return skippedPrefixes, nil
+}
+
+func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *zap.Logger, skippedPrefixes map[string]bool) error {
+	opts := &godo.ListOptions{
+		Page:    1,
+		PerPage: 200,
+	}
+
+	var allDroplets []godo.Droplet
+	for {
+		droplets, err := client.ListDroplets(ctx, opts)
+		if err != nil {
+			return fmt.Errorf("failed to list droplets: %w", err)
+		}
+
+		allDroplets = append(allDroplets, droplets...)
+
+		if len(droplets) < opts.PerPage {
+			break
+		}
+
+		opts.Page++
+	}
+
+	logger.Info("Retrieved droplets", zap.Int("total_count", len(allDroplets)))
+
+	var dropletsToDelete []godo.Droplet
+	for _, droplet := range allDroplets {
+		if !strings.HasPrefix(droplet.Name, *namePrefix) {
+			continue
+		}
+
+		if shouldSkipResource(droplet.Name, skippedPrefixes) {
+			logger.Info("Skipping droplet", zap.String("name", droplet.Name))
 			continue
 		}
 
@@ -121,14 +185,9 @@ func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *
 	logger.Info("Found droplets to delete", zap.Int("count", len(dropletsToDelete)))
 
 	if *dryRun {
-		logger.Info("Dry run mode - would delete the following droplets:")
-		for _, droplet := range dropletsToDelete {
-			logger.Info("Would delete droplet",
-				zap.String("name", droplet.Name),
-				zap.Int("id", droplet.ID),
-				zap.Strings("tags", droplet.Tags),
-				zap.String("created_at", droplet.Created))
-		}
+		logger.Info("Dry run mode - would delete droplets",
+			zap.Int("count", len(dropletsToDelete)),
+			zap.Any("droplets", dropletsToDelete))
 		return nil
 	}
 
@@ -150,8 +209,6 @@ func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *
 		logger.Info("Successfully deleted droplet",
 			zap.String("name", droplet.Name),
 			zap.Int("id", droplet.ID))
-
-		time.Sleep(100 * time.Millisecond)
 	}
 
 	logger.Info("Cleanup completed",
@@ -160,3 +217,78 @@ func cleanupDroplets(ctx context.Context, client digitalocean.DoClient, logger *
 
 	return nil
 }
+
+func shouldSkipResource(resourceName string, skippedPrefixes map[string]bool) bool {
+	prefix := extractPrefix(resourceName)
+	return skippedPrefixes[prefix]
+}
+
+func cleanupFirewalls(ctx context.Context, client digitalocean.DoClient, logger *zap.Logger, skippedPrefixes map[string]bool) error {
+	opts := &godo.ListOptions{
+		Page:    1,
+		PerPage: 200,
+	}
+
+	var allFirewalls []godo.Firewall
+	for {
+		firewalls, err := client.ListFirewalls(ctx, opts)
+		if err != nil {
+			return fmt.Errorf("failed to list firewalls: %w", err)
+		}
+
+		allFirewalls = append(allFirewalls, firewalls...)
+
+		if len(firewalls) < opts.PerPage {
+			break
+		}
+
+		opts.Page++
+	}
+
+	logger.Info("Retrieved firewalls", zap.Int("total_count", len(allFirewalls)))
+
+	var firewallsToDelete []godo.Firewall
+	for _, firewall := range allFirewalls {
+		if !strings.HasPrefix(firewall.Name, *namePrefix) {
+			continue
+		}
+
+		if shouldSkipResource(firewall.Name, skippedPrefixes) {
+			logger.Info("Skipping firewall ",
+				zap.String("name", firewall.Name))
+			continue
+		}
+
+		firewallsToDelete = append(firewallsToDelete, firewall)
+	}
+
+	logger.Info("Found firewalls to delete", zap.Int("count", len(firewallsToDelete)))
+
+	if *dryRun {
+		logger.Info("Dry run mode - would delete firewalls",
+			zap.Int("count", len(firewallsToDelete)),
+			zap.Any("firewalls", firewallsToDelete))
+		return nil
+	}
+
+	var deletedCount int
+	for _, firewall := range firewallsToDelete {
+		logger.Info("Deleting firewall", zap.String("name", firewall.Name))
+
+		if err := client.DeleteFirewall(ctx, firewall.ID); err != nil {
+			logger.Error("Failed to delete firewall",
+				zap.String("name", firewall.Name),
+				zap.Error(err))
+			continue
+		}
+
+		deletedCount++
+		logger.Info("Successfully deleted firewall",
+			zap.String("name", firewall.Name))
+	}
+
+	logger.Info("Firewall cleanup completed",
+		zap.Int("deleted_count", deletedCount))
+
+	return nil
+}

petri/core/provider/digitalocean/client.go

@@ -20,6 +20,7 @@ type DoClient interface {
 	// Firewall operations
 	CreateFirewall(ctx context.Context, req *godo.FirewallRequest) (*godo.Firewall, error)
 	GetFirewall(ctx context.Context, firewallID string) (*godo.Firewall, error)
+	ListFirewalls(ctx context.Context, opts *godo.ListOptions) ([]godo.Firewall, error)
 	DeleteFirewall(ctx context.Context, firewallID string) error
 
 	// Domain operations
@@ -34,6 +35,7 @@ type DoClient interface {
 
 	// Tag operations
 	CreateTag(ctx context.Context, req *godo.TagCreateRequest) (*godo.Tag, error)
+	ListTags(ctx context.Context, opts *godo.ListOptions) ([]godo.Tag, error)
 	DeleteTag(ctx context.Context, tag string) error
 }
 
@@ -115,6 +117,14 @@ func (c *godoClient) CreateFirewall(ctx context.Context, req *godo.FirewallReque
 	return firewall, nil
 }
 
+func (c *godoClient) ListFirewalls(ctx context.Context, opts *godo.ListOptions) ([]godo.Firewall, error) {
+	firewalls, res, err := c.Firewalls.List(ctx, opts)
+	if err := checkResponse(res, err); err != nil {
+		return nil, err
+	}
+	return firewalls, nil
+}
+
 func (c *godoClient) DeleteFirewall(ctx context.Context, firewallID string) error {
 	res, err := c.Firewalls.Delete(ctx, firewallID)
 	return checkResponse(res, err)
@@ -183,6 +193,14 @@ func (c *godoClient) CreateTag(ctx context.Context, req *godo.TagCreateRequest)
 	return tag, nil
 }
 
+func (c *godoClient) ListTags(ctx context.Context, opts *godo.ListOptions) ([]godo.Tag, error) {
+	tags, res, err := c.Tags.List(ctx, opts)
+	if err := checkResponse(res, err); err != nil {
+		return nil, err
+	}
+	return tags, nil
+}
+
 func (c *godoClient) DeleteTag(ctx context.Context, tag string) error {
 	res, err := c.Tags.Delete(ctx, tag)
 	return checkResponse(res, err)