mev-beta/cmd/mev-bot/main.go

package main

import (
	"context"
	"crypto/tls"
	"fmt"
	"math/big"
	"net/url"
	"os"
	"os/signal"
	"strconv"
	"strings"
	"syscall"
	"time"

	"github.com/joho/godotenv"
	"github.com/urfave/cli/v2"

	"github.com/fraktal/mev-beta/internal/config"
	"github.com/fraktal/mev-beta/internal/logger"
	"github.com/fraktal/mev-beta/internal/monitoring"
	"github.com/fraktal/mev-beta/pkg/arbitrage"
	"github.com/fraktal/mev-beta/pkg/metrics"
	"github.com/fraktal/mev-beta/pkg/pools"
	"github.com/fraktal/mev-beta/pkg/security"
	"github.com/fraktal/mev-beta/pkg/tokens"
	"github.com/fraktal/mev-beta/pkg/transport"
)

func main() {
	app := &cli.App{
		Name:  "mev-bot",
		Usage: "An MEV bot that monitors Arbitrum sequencer for swap opportunities",
		Commands: []*cli.Command{
			{
				Name:  "start",
				Usage: "Start the MEV bot",
				Action: func(c *cli.Context) error {
					return startBot()
				},
			},
			{
				Name:  "scan",
				Usage: "Scan for potential arbitrage opportunities",
				Action: func(c *cli.Context) error {
					return scanOpportunities()
				},
			},
		},
	}

	if err := app.Run(os.Args); err != nil {
		fmt.Fprintf(os.Stderr, "Error: %v\n", err)
		os.Exit(1)
	}
}

func startBot() error {
	// Load environment variables based on GO_ENV
	envMode := strings.ToLower(os.Getenv("GO_ENV"))
	if envMode == "" {
		envMode = "development"
	}

	var envFile string
	if envMode == "development" {
		envFile = ".env"
	} else {
		envFile = fmt.Sprintf(".env.%s", envMode)
	}

	if _, err := os.Stat(envFile); err == nil {
		if err := godotenv.Load(envFile); err != nil {
			fmt.Printf("Warning: failed to load %s: %v\n", envFile, err)
		} else {
			fmt.Printf("Loaded environment variables from %s\n", envFile)
		}
	} else {
		fmt.Printf("Warning: %s not found; proceeding without mode-specific env overrides\n", envFile)
	}

	// Load configuration
	configFile := "config/config.yaml"
	if _, err := os.Stat("config/local.yaml"); err == nil {
		configFile = "config/local.yaml"
	}
	if _, err := os.Stat("config/arbitrum_production.yaml"); err == nil {
		configFile = "config/arbitrum_production.yaml"
	}

	cfg, err := config.Load(configFile)
	if err != nil {
		return fmt.Errorf("failed to load config: %w", err)
	}

	// Initialize logger
	log := logger.New(cfg.Log.Level, cfg.Log.Format, cfg.Log.File)
	log.Info(fmt.Sprintf("Starting MEV bot with Enhanced Security - Config: %s", configFile))

	// Validate RPC endpoints for security
	if err := validateRPCEndpoint(cfg.Arbitrum.RPCEndpoint); err != nil {
		return fmt.Errorf("RPC endpoint validation failed: %w", err)
	}
	if cfg.Arbitrum.WSEndpoint != "" {
		if err := validateRPCEndpoint(cfg.Arbitrum.WSEndpoint); err != nil {
			return fmt.Errorf("WebSocket endpoint validation failed: %w", err)
		}
	}

	log.Debug(fmt.Sprintf("RPC Endpoint: %s", cfg.Arbitrum.RPCEndpoint))
	log.Debug(fmt.Sprintf("WS Endpoint: %s", cfg.Arbitrum.WSEndpoint))
	log.Debug(fmt.Sprintf("Chain ID: %d", cfg.Arbitrum.ChainID))

	// Initialize comprehensive security framework
	securityKeyDir := getEnvOrDefault("MEV_BOT_KEYSTORE_PATH", "keystore")
	securityConfig := &security.SecurityConfig{
		KeyStoreDir:       securityKeyDir,
		EncryptionEnabled: true,
		TransactionRPS:    100,
		RPCRPS:            200,
		MaxBurstSize:      50,
		FailureThreshold:  5,
		RecoveryTimeout:   5 * time.Minute,
		TLSMinVersion:     tls.VersionTLS12, // TLS 1.2 minimum
		EmergencyStopFile: "emergency.stop",
		MaxGasPrice:       "50000000000", // 50 gwei
		AlertWebhookURL:   os.Getenv("SECURITY_WEBHOOK_URL"),
		LogLevel:          cfg.Log.Level,
		RPCURL:            cfg.Arbitrum.RPCEndpoint,
	}

	securityManager, err := security.NewSecurityManager(securityConfig)
	if err != nil {
		return fmt.Errorf("failed to initialize security manager: %w", err)
	}
	defer func() {
		shutdownCtx, cancelShutdown := context.WithTimeout(context.Background(), 15*time.Second)
		defer cancelShutdown()
		if err := securityManager.Shutdown(shutdownCtx); err != nil {
			log.Error("Failed to shutdown security manager", "error", err)
		}
	}()

	log.Info("Security framework initialized successfully")

	// Initialize metrics collector
	metricsCollector := metrics.NewMetricsCollector(log)

	// Start metrics server if enabled
	var metricsServer *metrics.MetricsServer
	if os.Getenv("METRICS_ENABLED") == "true" {
		metricsPort := os.Getenv("METRICS_PORT")
		if metricsPort == "" {
			metricsPort = "9090"
		}
		metricsServer = metrics.NewMetricsServer(metricsCollector, log, metricsPort)
		go func() {
			if err := metricsServer.Start(); err != nil {
				log.Error("Metrics server error: ", err)
			}
		}()
		log.Info(fmt.Sprintf("Metrics server started on port %s", metricsPort))
	}

	// Initialize unified provider manager
	log.Info("Initializing provider manager with separate read-only, execution, and testing pools...")

	// Use existing providers.yaml config file for runtime
	providerConfigPath := "config/providers.yaml"

	providerManager, err := transport.NewUnifiedProviderManager(providerConfigPath)
	if err != nil {
		return fmt.Errorf("failed to initialize provider manager: %w", err)
	}
	defer func() {
		if err := providerManager.Close(); err != nil {
			log.Error("Failed to close provider manager", "error", err)
		}
	}()

	// Get execution client for transaction operations
	executionClient, err := providerManager.GetExecutionHTTPClient()
	if err != nil {
		return fmt.Errorf("failed to get execution client: %w", err)
	}

	// Log provider statistics
	providerStats := providerManager.GetAllStats()
	log.Info(fmt.Sprintf("Provider manager initialized with %d pool(s)", len(providerStats)-1)) // -1 for summary

	// Create key manager for secure transaction signing
	encryptionKey := os.Getenv("MEV_BOT_ENCRYPTION_KEY")
	if encryptionKey == "" {
		return fmt.Errorf("MEV_BOT_ENCRYPTION_KEY environment variable is required for secure operations")
	}

	keystorePath := getEnvOrDefault("MEV_BOT_KEYSTORE_PATH", "keystore")
	if strings.TrimSpace(keystorePath) == "" {
		keystorePath = "keystore"
	}
	fmt.Printf("Using keystore path: %s\n", keystorePath)
	log.Info(fmt.Sprintf("Using keystore path: %s", keystorePath))

	fmt.Printf("DEBUG: Creating KeyManager config...\n")
	keyManagerConfig := &security.KeyManagerConfig{
		KeystorePath:    keystorePath,
		EncryptionKey:   encryptionKey,
		KeyRotationDays: 30,
		MaxSigningRate:  100,
		SessionTimeout:  time.Hour,
		AuditLogPath:    getEnvOrDefault("MEV_BOT_AUDIT_LOG", "logs/audit.log"),
		BackupPath:      getEnvOrDefault("MEV_BOT_BACKUP_PATH", "backups"),
	}
	fmt.Printf("DEBUG: Calling NewKeyManager...\n")
	keyManager, err := security.NewKeyManager(keyManagerConfig, log)
	fmt.Printf("DEBUG: NewKeyManager returned, err=%v\n", err)
	if err != nil {
		return fmt.Errorf("failed to create key manager: %w", err)
	}
	fmt.Printf("DEBUG: KeyManager created successfully\n")

	// Create arbitrage database
	fmt.Printf("DEBUG: Creating arbitrage database at %s...\n", cfg.Database.File)
	arbitrageDB, err := arbitrage.NewSQLiteDatabase(cfg.Database.File, log)
	fmt.Printf("DEBUG: Database creation returned, err=%v\n", err)
	if err != nil {
		return fmt.Errorf("failed to create arbitrage database: %w", err)
	}
	fmt.Printf("DEBUG: Database created successfully\n")
	defer func() {
		if err := arbitrageDB.Close(); err != nil {
			log.Error("Failed to close arbitrage database", "error", err)
		}
	}()

	// Check if arbitrage service is enabled
	if !cfg.Arbitrage.Enabled {
		log.Info("Arbitrage service is disabled in configuration")
		return fmt.Errorf("arbitrage service disabled - enable in config to run")
	}

	// Setup graceful shutdown BEFORE creating services
	ctx, cancel := context.WithCancel(context.Background())
	defer cancel() // Ensure context is canceled on function exit

	// Get read-only provider pool for RPC operations
	readOnlyPool, err := providerManager.GetPoolForMode(transport.ModeReadOnly)
	if err != nil {
		return fmt.Errorf("failed to get read-only provider pool: %w", err)
	}

	// Get RPC client for pool discovery
	rpcClient, err := readOnlyPool.GetRPCClient(false) // Use HTTP for reliability
	if err != nil {
		return fmt.Errorf("failed to get RPC client for pool discovery: %w", err)
	}

	// Initialize Pool Discovery System
	log.Info("Initializing pool discovery system...")
	poolDiscovery := pools.NewPoolDiscovery(rpcClient, log)
	poolCount := poolDiscovery.GetPoolCount()
	log.Info(fmt.Sprintf("✅ Loaded %d pools from discovery system", poolCount))

	// Initialize Token Metadata Cache
	log.Info("Initializing token metadata cache...")
	tokenCache := tokens.NewMetadataCache(log)
	tokenCount := tokenCache.Count()
	log.Info(fmt.Sprintf("✅ Loaded %d tokens from cache", tokenCount))

	// Create arbitrage service with context and pool discovery
	log.Info("Creating arbitrage service...")
	fmt.Printf("DEBUG: Creating arbitrage service...\n")
	arbitrageService, err := arbitrage.NewArbitrageService(
		ctx,
		executionClient,
		log,
		&cfg.Arbitrage,
		keyManager,
		arbitrageDB,
		poolDiscovery,
		tokenCache,
	)
	fmt.Printf("DEBUG: ArbitrageService creation returned, err=%v\n", err)
	if err != nil {
		return fmt.Errorf("failed to create arbitrage service: %w", err)
	}
	log.Info("Arbitrage service created successfully")
	fmt.Printf("DEBUG: ArbitrageService created successfully\n")

	// Initialize data integrity monitoring system
	log.Info("Initializing data integrity monitoring system...")

	// Initialize integrity monitor
	integrityMonitor := monitoring.NewIntegrityMonitor(log)

	// Initialize dashboard server
	dashboardPort := 8080
	if portEnv := os.Getenv("DASHBOARD_PORT"); portEnv != "" {
		if port, err := strconv.Atoi(portEnv); err == nil {
			dashboardPort = port
		}
	}
	dashboardServer := monitoring.NewDashboardServer(log, integrityMonitor, integrityMonitor.GetHealthCheckRunner(), dashboardPort)

	// Start dashboard server
	go func() {
		log.Info(fmt.Sprintf("Starting monitoring dashboard on port %d...", dashboardPort))
		if err := dashboardServer.Start(); err != nil {
			log.Error("Dashboard server error", "error", err)
		}
	}()

	// Start integrity monitoring
	go func() {
		log.Info("Starting integrity monitoring...")
		integrityMonitor.StartHealthCheckRunner(ctx)
	}()

	log.Info("Data integrity monitoring system initialized successfully")
	log.Info(fmt.Sprintf("Dashboard available at http://localhost:%d", dashboardPort))

	sigChan := make(chan os.Signal, 1)
	signal.Notify(sigChan, syscall.SIGINT, syscall.SIGTERM)

	// Handle signals in a goroutine to cancel context immediately
	go func() {
		<-sigChan
		log.Info("Shutdown signal received, canceling context...")
		cancel() // This will cancel the context and stop all operations
	}()

	// Start the arbitrage service with context
	log.Info("Starting arbitrage service...")
	errChan := make(chan error, 1)
	go func() {
		if err := arbitrageService.Start(); err != nil {
			errChan <- fmt.Errorf("arbitrage service error: %w", err)
		}
	}()
	defer func() {
		if err := arbitrageService.Stop(); err != nil {
			log.Error("Failed to stop arbitrage service", "error", err)
		}
	}()
	log.Info("Arbitrage service started successfully")

	log.Info("MEV bot started successfully - monitoring for arbitrage opportunities...")
	log.Info("Press Ctrl+C to stop the bot gracefully...")

	// Wait for context cancellation or error
	select {
	case <-ctx.Done():
		log.Info("Context canceled, stopping MEV bot...")
	case err := <-errChan:
		log.Error("Service error occurred: ", err)
		return err
	}

	// Stop monitoring services
	log.Info("Stopping monitoring services...")
	if err := dashboardServer.Stop(); err != nil {
		log.Error("Failed to stop dashboard server gracefully", "error", err)
	}
	integrityMonitor.StopHealthCheckRunner()

	// Stop metrics server if running
	if metricsServer != nil {
		if err := metricsServer.Stop(); err != nil {
			log.Error("Failed to stop metrics server gracefully", "error", err)
		}
	}

	// Get final stats
	stats := arbitrageService.GetStats()
	log.Info(fmt.Sprintf("Final Statistics - Opportunities: %d, Executions: %d, Successful: %d, Total Profit: %s ETH",
		stats.TotalOpportunitiesDetected,
		stats.TotalOpportunitiesExecuted,
		stats.TotalSuccessfulExecutions,
		formatEther(stats.TotalProfitRealized)))

	log.Info("MEV bot stopped gracefully")
	return nil
}

// formatEther formats wei amount to ether string
func formatEther(wei *big.Int) string {
	if wei == nil {
		return "0.000000"
	}
	eth := new(big.Float).SetInt(wei)
	eth.Quo(eth, big.NewFloat(1e18))
	return fmt.Sprintf("%.6f", eth)
}

// getEnvOrDefault returns environment variable value or default if not set
func getEnvOrDefault(key, defaultValue string) string {
	if value := os.Getenv(key); value != "" {
		return value
	}
	return defaultValue
}

// validateRPCEndpoint validates RPC endpoint URL for security
func validateRPCEndpoint(endpoint string) error {
	if endpoint == "" {
		return fmt.Errorf("RPC endpoint cannot be empty")
	}

	u, err := url.Parse(endpoint)
	if err != nil {
		return fmt.Errorf("invalid RPC endpoint URL: %w", err)
	}

	// Check for valid schemes
	switch u.Scheme {
	case "http", "https", "ws", "wss":
		// Valid schemes
	default:
		return fmt.Errorf("invalid RPC scheme: %s (must be http, https, ws, or wss)", u.Scheme)
	}

	// Check for localhost/private networks in production
	if strings.Contains(u.Hostname(), "localhost") || strings.Contains(u.Hostname(), "127.0.0.1") {
		// Allow localhost only if explicitly enabled
		if os.Getenv("MEV_BOT_ALLOW_LOCALHOST") != "true" {
			return fmt.Errorf("localhost RPC endpoints not allowed in production (set MEV_BOT_ALLOW_LOCALHOST=true to override)")
		}
	}

	// Validate hostname is not empty
	if u.Hostname() == "" {
		return fmt.Errorf("RPC endpoint must have a valid hostname")
	}

	return nil
}

func scanOpportunities() error {
	fmt.Println("Scanning for arbitrage opportunities...")

	// Load configuration
	configFile := "config/config.yaml"
	if _, err := os.Stat("config/local.yaml"); err == nil {
		configFile = "config/local.yaml"
	}
	if _, err := os.Stat("config/arbitrum_production.yaml"); err == nil {
		configFile = "config/arbitrum_production.yaml"
	}

	cfg, err := config.Load(configFile)
	if err != nil {
		return fmt.Errorf("failed to load config: %w", err)
	}

	// Initialize logger
	log := logger.New(cfg.Log.Level, cfg.Log.Format, cfg.Log.File)
	log.Info("Starting one-time arbitrage opportunity scan...")

	// Initialize provider manager for scanning
	providerManager, err := transport.NewUnifiedProviderManager("config/providers.yaml")
	if err != nil {
		return fmt.Errorf("failed to initialize provider manager: %w", err)
	}
	defer func() {
		if err := providerManager.Close(); err != nil {
			log.Error("Failed to close provider manager in scan mode", "error", err)
		}
	}()

	// Get read-only client for scanning (more efficient)
	client, err := providerManager.GetReadOnlyHTTPClient()
	if err != nil {
		return fmt.Errorf("failed to get read-only client: %w", err)
	}

	// Create key manager (not used for scanning but needed for service)
	encryptionKey := os.Getenv("MEV_BOT_ENCRYPTION_KEY")
	if encryptionKey == "" {
		return fmt.Errorf("MEV_BOT_ENCRYPTION_KEY environment variable is required")
	}

	keyManagerConfig := &security.KeyManagerConfig{
		KeystorePath:    getEnvOrDefault("MEV_BOT_KEYSTORE_PATH", "keystore"),
		EncryptionKey:   encryptionKey,
		KeyRotationDays: 30,
		MaxSigningRate:  100,
		SessionTimeout:  time.Hour,
		AuditLogPath:    getEnvOrDefault("MEV_BOT_AUDIT_LOG", "logs/audit.log"),
		BackupPath:      getEnvOrDefault("MEV_BOT_BACKUP_PATH", "backups"),
	}
	keyManager, err := security.NewKeyManager(keyManagerConfig, log)
	if err != nil {
		return fmt.Errorf("failed to create key manager: %w", err)
	}

	// Create arbitrage database
	arbitrageDB, err := arbitrage.NewSQLiteDatabase(cfg.Database.File, log)
	if err != nil {
		return fmt.Errorf("failed to create arbitrage database: %w", err)
	}
	defer func() {
		if err := arbitrageDB.Close(); err != nil {
			log.Error("Failed to close arbitrage database in scan mode", "error", err)
		}
	}()

	// Get read-only provider pool for RPC operations in scan mode
	readOnlyPool, err := providerManager.GetPoolForMode(transport.ModeReadOnly)
	if err != nil {
		return fmt.Errorf("failed to get read-only provider pool in scan mode: %w", err)
	}

	// Get RPC client for pool discovery in scan mode
	rpcClient, err := readOnlyPool.GetRPCClient(false)
	if err != nil {
		return fmt.Errorf("failed to get RPC client for pool discovery in scan mode: %w", err)
	}

	// Initialize pool discovery and token cache for scan mode
	poolDiscovery := pools.NewPoolDiscovery(rpcClient, log)
	tokenCache := tokens.NewMetadataCache(log)

	// Create arbitrage service with scanning enabled but execution disabled
	scanConfig := cfg.Arbitrage
	scanConfig.MaxConcurrentExecutions = 0 // Disable execution for scan mode

	arbitrageService, err := arbitrage.NewArbitrageService(
		context.Background(),
		client,
		log,
		&scanConfig,
		keyManager,
		arbitrageDB,
		poolDiscovery,
		tokenCache,
	)
	if err != nil {
		return fmt.Errorf("failed to create arbitrage service: %w", err)
	}

	// Start the service in scan mode
	if err := arbitrageService.Start(); err != nil {
		return fmt.Errorf("failed to start arbitrage service: %w", err)
	}
	defer func() {
		if err := arbitrageService.Stop(); err != nil {
			log.Error("Failed to stop arbitrage service in scan mode", "error", err)
		}
	}()

	// Create context with timeout for scanning
	ctx, cancel := context.WithTimeout(context.Background(), 30*time.Second)
	defer cancel()

	log.Info("Scanning for 30 seconds...")

	// Wait for scan duration
	<-ctx.Done()

	// Get and display results
	stats := arbitrageService.GetStats()
	log.Info(fmt.Sprintf("Scan Results - Opportunities Detected: %d", stats.TotalOpportunitiesDetected))

	// Get recent opportunities from database
	history, err := arbitrageDB.GetExecutionHistory(ctx, 10)
	if err == nil && len(history) > 0 {
		log.Info(fmt.Sprintf("Found %d recent opportunities in database", len(history)))
	}

	fmt.Printf("Scan completed. Found %d arbitrage opportunities.\n", stats.TotalOpportunitiesDetected)
	fmt.Println("Check the database and logs for detailed results.")

	return nil
}