mail-migrator/main.go

package main

import (
	"context"
	"encoding/csv"
	"fmt"
	"io"
	"log"
	"os"
	"time"

	"github.com/sirupsen/logrus"
	"github.com/urfave/cli/v2"
	bolt "go.etcd.io/bbolt"
)

func main() {
	app := &cli.App{
		Name:  "mail-migrator",
		Usage: "Migrate IMAP mailboxes between servers",
		Flags: []cli.Flag{
			&cli.StringFlag{Name: "src", Usage: "Source IMAP URL (user:pass@host:port)"},
			&cli.StringFlag{Name: "dst", Usage: "Destination IMAP URL (user:pass@host:port)"},
			&cli.StringFlag{Name: "batch", Usage: "CSV file with src,dst entries for batch migration"},
			&cli.BoolFlag{Name: "insecure", Usage: "Allow insecure TLS (trust self-signed certs)"},
			&cli.BoolFlag{Name: "resume", Usage: "Resume incomplete migrations"},
			&cli.StringFlag{Name: "state", Value: "state.db", Usage: "Path to state DB (Bolt)"},
			&cli.StringFlag{Name: "log", Usage: "Path to log file (optional)"},
		},
		Action: func(c *cli.Context) error {
			setupLogging(c.String("log"))
			insecure := c.Bool("insecure")
			resume := c.Bool("resume")
			statePath := c.String("state")

			if batch := c.String("batch"); batch != "" {
				return runBatch(batch, insecure, resume, statePath)
			}

			src := c.String("src")
			dst := c.String("dst")
			if src == "" || dst == "" {
				return fmt.Errorf("either src/dst or batch must be provided")
			}
			return migrateOne(context.Background(), src, dst, insecure, resume, statePath)
		},
	}

	if err := app.Run(os.Args); err != nil {
		logrus.Fatal(err)
	}
}

func setupLogging(file string) {
	logrus.SetFormatter(&logrus.TextFormatter{FullTimestamp: true})
	logrus.SetLevel(logrus.DebugLevel) // Changed from InfoLevel to DebugLevel for troubleshooting
	if file != "" {
		f, err := os.OpenFile(file, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644)
		if err != nil {
			log.Fatalf("unable to open log file: %v", err)
		}
		logrus.SetOutput(io.MultiWriter(os.Stdout, f))
	}
}

func runBatch(path string, insecure, resume bool, statePath string) error {
	f, err := os.Open(path)
	if err != nil {
		return err
	}
	defer f.Close()
	reader := csv.NewReader(f)
	for {
		rec, err := reader.Read()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		if len(rec) < 2 {
			logrus.Warn("invalid record, skipping")
			continue
		}
		src, dst := rec[0], rec[1]
		ctx, cancel := context.WithTimeout(context.Background(), 30*time.Minute)
		if err := migrateOne(ctx, src, dst, insecure, resume, statePath); err != nil {
			logrus.Errorf("migration failed for %s -> %s: %v", src, dst, err)
		}
		cancel()
	}
	return nil
}

// migrateOne performs migration for single account
func migrateOne(ctx context.Context, src, dst string, insecure, resume bool, statePath string) error {
	logrus.Infof("Starting migration %s -> %s (insecure=%v, resume=%v)", src, dst, insecure, resume)

	// Parse IMAP URLs
	srcCfg, err := parseIMAPURL(src)
	if err != nil {
		return fmt.Errorf("parse source URL: %w", err)
	}

	dstCfg, err := parseIMAPURL(dst)
	if err != nil {
		return fmt.Errorf("parse destination URL: %w", err)
	}

	// Open state database if resume is enabled
	var db *bolt.DB
	if resume {
		db, err = openStateDB(statePath)
		if err != nil {
			return fmt.Errorf("open state DB: %w", err)
		}
		defer db.Close()
	}

	// Connect to source
	srcClient, err := connectIMAP(srcCfg, insecure)
	if err != nil {
		return fmt.Errorf("connect to source: %w", err)
	}
	defer srcClient.Logout()

	// Connect to destination
	dstClient, err := connectIMAP(dstCfg, insecure)
	if err != nil {
		return fmt.Errorf("connect to destination: %w", err)
	}
	defer dstClient.Close()

	// List mailboxes from source
	mailboxes, err := listMailboxes(srcClient)
	if err != nil {
		return fmt.Errorf("list mailboxes: %w", err)
	}

	logrus.Infof("Found %d mailboxes to migrate", len(mailboxes))

	// Migrate each mailbox
	for _, mbox := range mailboxes {
		// Skip certain system mailboxes
		if shouldSkipMailbox(mbox.Name) {
			logrus.Debugf("Skipping mailbox: %s", mbox.Name)
			continue
		}

		stateKey := formatStateKey(src, dst, mbox.Name)

		logrus.Infof("Migrating mailbox: %s", mbox.Name)
		if err := copyMessages(srcClient, dstClient, mbox.Name, mbox.Name, stateKey, db); err != nil {
			logrus.Errorf("Failed to migrate mailbox %s: %v", mbox.Name, err)
			continue
		}
	}

	logrus.Infof("Migration completed successfully")
	return nil
}

// shouldSkipMailbox determines if a mailbox should be skipped
func shouldSkipMailbox(name string) bool {
	// Skip common system/virtual mailboxes
	skipList := []string{
		"[Gmail]/All Mail",
		"[Gmail]/Important",
		"[Gmail]/Starred",
		"[Gmail]/Chats",
	}

	for _, skip := range skipList {
		if name == skip {
			return true
		}
	}
	return false
}