mev-beta/pkg/security/transaction_security.go

package security

import (
	"context"
	"fmt"
	"math/big"
	"time"

	"github.com/ethereum/go-ethereum/common"
	"github.com/ethereum/go-ethereum/core/types"
	"github.com/ethereum/go-ethereum/ethclient"

	"github.com/fraktal/mev-beta/internal/logger"
)

// TransactionSecurity provides comprehensive transaction security checks
type TransactionSecurity struct {
	logger         *logger.Logger
	inputValidator *InputValidator
	safeMath       *SafeMath
	client         *ethclient.Client
	chainID        uint64

	// Security thresholds
	maxTransactionValue *big.Int
	maxGasPrice         *big.Int
	maxSlippageBps      uint64

	// Blacklisted addresses
	blacklistedAddresses map[common.Address]bool

	// Rate limiting per address
	transactionCounts map[common.Address]int
	lastReset         time.Time
	maxTxPerAddress   int
}

// TransactionSecurityResult contains the security analysis result
type TransactionSecurityResult struct {
	Approved        bool                   `json:"approved"`
	RiskLevel       string                 `json:"risk_level"` // LOW, MEDIUM, HIGH, CRITICAL
	SecurityChecks  map[string]bool        `json:"security_checks"`
	Warnings        []string               `json:"warnings"`
	Errors          []string               `json:"errors"`
	RecommendedGas  *big.Int               `json:"recommended_gas,omitempty"`
	MaxSlippage     uint64                 `json:"max_slippage_bps,omitempty"`
	EstimatedProfit *big.Int               `json:"estimated_profit,omitempty"`
	Metadata        map[string]interface{} `json:"metadata,omitempty"`
}

// MEVTransactionRequest represents an MEV transaction request
type MEVTransactionRequest struct {
	Transaction    *types.Transaction `json:"transaction"`
	ExpectedProfit *big.Int           `json:"expected_profit"`
	MaxSlippage    uint64             `json:"max_slippage_bps"`
	Deadline       time.Time          `json:"deadline"`
	Priority       string             `json:"priority"` // LOW, MEDIUM, HIGH
	Source         string             `json:"source"`   // Origin of the transaction
}

// NewTransactionSecurity creates a new transaction security checker
func NewTransactionSecurity(client *ethclient.Client, logger *logger.Logger, chainID uint64) *TransactionSecurity {
	return &TransactionSecurity{
		logger:               logger,
		inputValidator:       NewInputValidator(chainID),
		safeMath:             NewSafeMath(),
		client:               client,
		chainID:              chainID,
		maxTransactionValue:  new(big.Int).Mul(big.NewInt(1000), big.NewInt(1e18)), // 1000 ETH
		maxGasPrice:          new(big.Int).Mul(big.NewInt(10000), big.NewInt(1e9)), // 10000 Gwei
		maxSlippageBps:       1000,                                                 // 10%
		blacklistedAddresses: make(map[common.Address]bool),
		transactionCounts:    make(map[common.Address]int),
		lastReset:            time.Now(),
		maxTxPerAddress:      100, // Max 100 transactions per address per hour
	}
}

// AnalyzeMEVTransaction performs comprehensive security analysis on an MEV transaction
func (ts *TransactionSecurity) AnalyzeMEVTransaction(ctx context.Context, req *MEVTransactionRequest) (*TransactionSecurityResult, error) {
	result := &TransactionSecurityResult{
		Approved:       true,
		RiskLevel:      "LOW",
		SecurityChecks: make(map[string]bool),
		Warnings:       []string{},
		Errors:         []string{},
		Metadata:       make(map[string]interface{}),
	}

	// Basic transaction validation
	if err := ts.basicTransactionChecks(req.Transaction, result); err != nil {
		return result, fmt.Errorf("basic transaction checks failed: %w", err)
	}

	// MEV-specific checks
	if err := ts.mevSpecificChecks(ctx, req, result); err != nil {
		return result, fmt.Errorf("MEV specific checks failed: %w", err)
	}

	// Gas price and limit validation
	if err := ts.gasValidation(req.Transaction, result); err != nil {
		return result, fmt.Errorf("gas validation failed: %w", err)
	}

	// Profit validation
	if err := ts.profitValidation(req, result); err != nil {
		return result, fmt.Errorf("profit validation failed: %w", err)
	}

	// Front-running protection checks
	if err := ts.frontRunningProtection(ctx, req, result); err != nil {
		return result, fmt.Errorf("front-running protection failed: %w", err)
	}

	// Rate limiting checks
	if err := ts.rateLimitingChecks(req.Transaction, result); err != nil {
		return result, fmt.Errorf("rate limiting checks failed: %w", err)
	}

	// Calculate final risk level
	ts.calculateRiskLevel(result)

	return result, nil
}

// basicTransactionChecks performs basic transaction security checks
func (ts *TransactionSecurity) basicTransactionChecks(tx *types.Transaction, result *TransactionSecurityResult) error {
	// Validate transaction using input validator
	validationResult := ts.inputValidator.ValidateTransaction(tx)
	if !validationResult.Valid {
		result.Approved = false
		result.Errors = append(result.Errors, validationResult.Errors...)
		result.SecurityChecks["basic_validation"] = false
		return fmt.Errorf("transaction failed basic validation")
	}
	result.SecurityChecks["basic_validation"] = true
	result.Warnings = append(result.Warnings, validationResult.Warnings...)

	// Check against blacklisted addresses
	if tx.To() != nil {
		if ts.blacklistedAddresses[*tx.To()] {
			result.Approved = false
			result.Errors = append(result.Errors, "transaction recipient is blacklisted")
			result.SecurityChecks["blacklist_check"] = false
			return fmt.Errorf("blacklisted recipient address")
		}
	}
	result.SecurityChecks["blacklist_check"] = true

	// Check transaction size
	if tx.Size() > 128*1024 { // 128KB limit
		result.Approved = false
		result.Errors = append(result.Errors, "transaction size exceeds limit")
		result.SecurityChecks["size_check"] = false
		return fmt.Errorf("transaction too large")
	}
	result.SecurityChecks["size_check"] = true

	return nil
}

// mevSpecificChecks performs MEV-specific security validations
func (ts *TransactionSecurity) mevSpecificChecks(ctx context.Context, req *MEVTransactionRequest, result *TransactionSecurityResult) error {
	// Check deadline
	if req.Deadline.Before(time.Now()) {
		result.Approved = false
		result.Errors = append(result.Errors, "transaction deadline has passed")
		result.SecurityChecks["deadline_check"] = false
		return fmt.Errorf("deadline expired")
	}

	// Warn if deadline is too far in the future
	if req.Deadline.After(time.Now().Add(1 * time.Hour)) {
		result.Warnings = append(result.Warnings, "deadline is more than 1 hour in the future")
	}
	result.SecurityChecks["deadline_check"] = true

	// Validate slippage
	if req.MaxSlippage > ts.maxSlippageBps {
		result.Approved = false
		result.Errors = append(result.Errors, fmt.Sprintf("slippage %d bps exceeds maximum %d bps", req.MaxSlippage, ts.maxSlippageBps))
		result.SecurityChecks["slippage_check"] = false
		return fmt.Errorf("excessive slippage")
	}

	if req.MaxSlippage > 500 { // Warn if > 5%
		result.Warnings = append(result.Warnings, fmt.Sprintf("high slippage detected: %d bps", req.MaxSlippage))
	}
	result.SecurityChecks["slippage_check"] = true

	// Check transaction priority vs gas price
	if err := ts.validatePriorityVsGasPrice(req, result); err != nil {
		return err
	}

	return nil
}

// gasValidation performs gas-related security checks
func (ts *TransactionSecurity) gasValidation(tx *types.Transaction, result *TransactionSecurityResult) error {
	// Calculate minimum required gas
	minGas := uint64(21000) // Base transaction gas
	if len(tx.Data()) > 0 {
		// Add gas for contract call
		minGas += uint64(len(tx.Data())) * 16 // 16 gas per non-zero byte
	}

	if tx.Gas() < minGas {
		result.Approved = false
		result.Errors = append(result.Errors, fmt.Sprintf("gas limit %d below minimum required %d", tx.Gas(), minGas))
		result.SecurityChecks["gas_limit_check"] = false
		return fmt.Errorf("insufficient gas limit")
	}

	// Recommend optimal gas limit (add 20% buffer)
	recommendedGas := new(big.Int).SetUint64(minGas * 120 / 100)
	result.RecommendedGas = recommendedGas
	result.SecurityChecks["gas_limit_check"] = true

	// Validate gas price
	if tx.GasPrice() != nil {
		if err := ts.safeMath.ValidateGasPrice(tx.GasPrice()); err != nil {
			result.Approved = false
			result.Errors = append(result.Errors, fmt.Sprintf("invalid gas price: %v", err))
			result.SecurityChecks["gas_price_check"] = false
			return fmt.Errorf("invalid gas price")
		}

		// Check if gas price is suspiciously high
		highGasThreshold := new(big.Int).Mul(big.NewInt(1000), big.NewInt(1e9)) // 1000 Gwei
		if tx.GasPrice().Cmp(highGasThreshold) > 0 {
			result.Warnings = append(result.Warnings, fmt.Sprintf("high gas price detected: %s Gwei",
				new(big.Int).Div(tx.GasPrice(), big.NewInt(1e9)).String()))
		}
	}
	result.SecurityChecks["gas_price_check"] = true

	return nil
}

// profitValidation validates expected profit and ensures it covers costs
func (ts *TransactionSecurity) profitValidation(req *MEVTransactionRequest, result *TransactionSecurityResult) error {
	if req.ExpectedProfit == nil || req.ExpectedProfit.Sign() <= 0 {
		result.Approved = false
		result.Errors = append(result.Errors, "expected profit must be positive")
		result.SecurityChecks["profit_check"] = false
		return fmt.Errorf("invalid expected profit")
	}

	// Calculate transaction cost
	if req.Transaction.GasPrice() != nil {
		gasInt64, err := SafeUint64ToInt64(req.Transaction.Gas())
		if err != nil {
			ts.logger.Error("Transaction gas exceeds int64 maximum", "gas", req.Transaction.Gas(), "error", err)
			result.Approved = false
			result.Errors = append(result.Errors, fmt.Sprintf("gas value exceeds maximum allowed: %v", err))
			result.SecurityChecks["profit_check"] = false
			return fmt.Errorf("gas value exceeds maximum allowed: %w", err)
		}
		gasCost := new(big.Int).Mul(req.Transaction.GasPrice(), big.NewInt(gasInt64))

		// Ensure profit exceeds gas cost by at least 50%
		minProfit := new(big.Int).Mul(gasCost, big.NewInt(150))
		minProfit.Div(minProfit, big.NewInt(100))

		if req.ExpectedProfit.Cmp(minProfit) < 0 {
			result.Approved = false
			result.Errors = append(result.Errors, "expected profit does not cover transaction costs with adequate margin")
			result.SecurityChecks["profit_check"] = false
			return fmt.Errorf("insufficient profit margin")
		}

		result.EstimatedProfit = req.ExpectedProfit
		result.Metadata["gas_cost"] = gasCost.String()
		result.Metadata["profit_margin"] = new(big.Int).Div(
			new(big.Int).Mul(req.ExpectedProfit, big.NewInt(100)),
			gasCost,
		).String() + "%"
	}

	result.SecurityChecks["profit_check"] = true
	return nil
}

// frontRunningProtection implements front-running protection measures
func (ts *TransactionSecurity) frontRunningProtection(ctx context.Context, req *MEVTransactionRequest, result *TransactionSecurityResult) error {
	// Check if transaction might be front-runnable
	if req.Transaction.GasPrice() != nil {
		// Get current network gas price
		networkGasPrice, err := ts.client.SuggestGasPrice(ctx)
		if err != nil {
			result.Warnings = append(result.Warnings, "could not fetch network gas price for front-running analysis")
		} else {
			// If our gas price is significantly higher, we might be front-runnable
			threshold := new(big.Int).Mul(networkGasPrice, big.NewInt(150)) // 50% above network
			threshold.Div(threshold, big.NewInt(100))

			if req.Transaction.GasPrice().Cmp(threshold) > 0 {
				result.Warnings = append(result.Warnings, "transaction gas price significantly above network average - vulnerable to front-running")
				result.Metadata["front_running_risk"] = "HIGH"
			} else {
				result.Metadata["front_running_risk"] = "LOW"
			}
		}
	}

	// Recommend using private mempool for high-value transactions
	if req.Transaction.Value() != nil {
		highValueThreshold := new(big.Int).Mul(big.NewInt(10), big.NewInt(1e18)) // 10 ETH
		if req.Transaction.Value().Cmp(highValueThreshold) > 0 {
			result.Warnings = append(result.Warnings, "high-value transaction should consider private mempool")
		}
	}

	result.SecurityChecks["front_running_protection"] = true
	return nil
}

// rateLimitingChecks implements per-address rate limiting
func (ts *TransactionSecurity) rateLimitingChecks(tx *types.Transaction, result *TransactionSecurityResult) error {
	// Reset counters if more than an hour has passed
	if time.Since(ts.lastReset) > time.Hour {
		ts.transactionCounts = make(map[common.Address]int)
		ts.lastReset = time.Now()
	}

	// Get sender address via signature recovery
	signer := types.LatestSignerForChainID(tx.ChainId())
	addr, err := types.Sender(signer, tx)
	if err != nil {
		// If signature recovery fails, use zero address
		// Note: In production, this should be logged to a centralized logging system
		addr = common.Address{}
	}

	// Increment counter
	ts.transactionCounts[addr]++

	// Check if limit exceeded
	if ts.transactionCounts[addr] > ts.maxTxPerAddress {
		result.Approved = false
		result.Errors = append(result.Errors, fmt.Sprintf("rate limit exceeded for address %s", addr.Hex()))
		result.SecurityChecks["rate_limiting"] = false
		return fmt.Errorf("rate limit exceeded")
	}

	// Warn if approaching limit
	if ts.transactionCounts[addr] > ts.maxTxPerAddress*8/10 {
		result.Warnings = append(result.Warnings, "approaching rate limit for this address")
	}

	result.SecurityChecks["rate_limiting"] = true
	result.Metadata["transaction_count"] = ts.transactionCounts[addr]
	result.Metadata["rate_limit"] = ts.maxTxPerAddress

	return nil
}

// validatePriorityVsGasPrice ensures gas price matches declared priority
func (ts *TransactionSecurity) validatePriorityVsGasPrice(req *MEVTransactionRequest, result *TransactionSecurityResult) error {
	if req.Transaction.GasPrice() == nil {
		return nil
	}

	gasPrice := req.Transaction.GasPrice()
	gasPriceGwei := new(big.Int).Div(gasPrice, big.NewInt(1e9))

	switch req.Priority {
	case "LOW":
		if gasPriceGwei.Cmp(big.NewInt(100)) > 0 { // > 100 Gwei
			result.Warnings = append(result.Warnings, "gas price seems high for LOW priority transaction")
		}
	case "MEDIUM":
		if gasPriceGwei.Cmp(big.NewInt(500)) > 0 { // > 500 Gwei
			result.Warnings = append(result.Warnings, "gas price seems high for MEDIUM priority transaction")
		}
	case "HIGH":
		if gasPriceGwei.Cmp(big.NewInt(50)) < 0 { // < 50 Gwei
			result.Warnings = append(result.Warnings, "gas price seems low for HIGH priority transaction")
		}
	}

	result.SecurityChecks["priority_gas_alignment"] = true
	return nil
}

// calculateRiskLevel calculates the overall risk level based on checks and warnings
func (ts *TransactionSecurity) calculateRiskLevel(result *TransactionSecurityResult) {
	if !result.Approved {
		result.RiskLevel = "CRITICAL"
		return
	}

	// Count failed checks
	failedChecks := 0
	for _, passed := range result.SecurityChecks {
		if !passed {
			failedChecks++
		}
	}

	// Determine risk level
	if failedChecks > 0 {
		result.RiskLevel = "HIGH"
	} else if len(result.Warnings) > 3 {
		result.RiskLevel = "MEDIUM"
	} else if len(result.Warnings) > 0 {
		result.RiskLevel = "LOW"
	} else {
		result.RiskLevel = "MINIMAL"
	}
}

// AddBlacklistedAddress adds an address to the blacklist
func (ts *TransactionSecurity) AddBlacklistedAddress(addr common.Address) {
	ts.blacklistedAddresses[addr] = true
}

// RemoveBlacklistedAddress removes an address from the blacklist
func (ts *TransactionSecurity) RemoveBlacklistedAddress(addr common.Address) {
	delete(ts.blacklistedAddresses, addr)
}

// GetSecurityMetrics returns current security metrics
func (ts *TransactionSecurity) GetSecurityMetrics() map[string]interface{} {
	return map[string]interface{}{
		"blacklisted_addresses_count":  len(ts.blacklistedAddresses),
		"active_address_count":         len(ts.transactionCounts),
		"max_transactions_per_address": ts.maxTxPerAddress,
		"max_transaction_value":        ts.maxTransactionValue.String(),
		"max_gas_price":                ts.maxGasPrice.String(),
		"max_slippage_bps":             ts.maxSlippageBps,
		"last_reset":                   ts.lastReset.Format(time.RFC3339),
	}
}