tm_back/pkg/security/csp.go

package security

import (
	"net/http"
	"strings"

	"github.com/labstack/echo/v4"
)

// CSPConfig holds CSP configuration
type CSPConfig struct {
	DefaultSrc []string
	ScriptSrc  []string
	StyleSrc   []string
	ImgSrc     []string
	ConnectSrc []string
	FontSrc    []string
	ObjectSrc  []string
	MediaSrc   []string
	FrameSrc   []string
	ReportURI  string
	ReportOnly bool
}

// DefaultCSPConfig returns a secure default CSP configuration
func DefaultCSPConfig() *CSPConfig {
	return &CSPConfig{
		DefaultSrc: []string{"'self'"},
		ScriptSrc:  []string{"'self'"},
		StyleSrc:   []string{"'self'"},
		ImgSrc:     []string{"'self'", "data:", "https:", "http:"},
		ConnectSrc: []string{"'self'", "https://api.opplens.com"},
		FontSrc:    []string{"'self'", "https://fonts.gstatic.com"},
		ObjectSrc:  []string{"'none'"},
		MediaSrc:   []string{"'self'"},
		FrameSrc:   []string{"'none'"},
		ReportURI:  "/api/v1/csp-report",
		ReportOnly: false,
	}
}

// CSPSource represents a CSP directive source
type CSPSource string

// CSPDirective represents a CSP directive
type CSPDirective struct {
	Name    string
	Sources []CSPSource
}

// CSPMiddleware creates Echo middleware for Content Security Policy
func CSPMiddleware(config *CSPConfig) echo.MiddlewareFunc {
	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			// Build CSP header
			directives := []string{}

			if len(config.DefaultSrc) > 0 {
				directives = append(directives, "default-src "+strings.Join(toStringSlice(config.DefaultSrc), " "))
			}

			if len(config.ScriptSrc) > 0 {
				directives = append(directives, "script-src "+strings.Join(toStringSlice(config.ScriptSrc), " "))
			}

			if len(config.StyleSrc) > 0 {
				directives = append(directives, "style-src "+strings.Join(toStringSlice(config.StyleSrc), " "))
			}

			if len(config.ImgSrc) > 0 {
				directives = append(directives, "img-src "+strings.Join(toStringSlice(config.ImgSrc), " "))
			}

			if len(config.ConnectSrc) > 0 {
				directives = append(directives, "connect-src "+strings.Join(toStringSlice(config.ConnectSrc), " "))
			}

			if len(config.FontSrc) > 0 {
				directives = append(directives, "font-src "+strings.Join(toStringSlice(config.FontSrc), " "))
			}

			if len(config.ObjectSrc) > 0 {
				directives = append(directives, "object-src "+strings.Join(toStringSlice(config.ObjectSrc), " "))
			}

			if len(config.MediaSrc) > 0 {
				directives = append(directives, "media-src "+strings.Join(toStringSlice(config.MediaSrc), " "))
			}

			if len(config.FrameSrc) > 0 {
				directives = append(directives, "frame-src "+strings.Join(toStringSlice(config.FrameSrc), " "))
			}

			if config.ReportURI != "" {
				directives = append(directives, "report-uri "+config.ReportURI)
			}

			cspHeader := strings.Join(directives, "; ")

			if config.ReportOnly {
				c.Response().Header().Set("Content-Security-Policy-Report-Only", cspHeader)
			} else {
				c.Response().Header().Set("Content-Security-Policy", cspHeader)
			}

			return next(c)
		}
	}
}

// CSPMiddlewareForPaths applies CSP middleware based on request path prefix.
// Requests matching no prefix proceed without a CSP header.
func CSPMiddlewareForPaths(adminPrefix string, adminConfig *CSPConfig, publicPrefix string, publicConfig *CSPConfig) echo.MiddlewareFunc {
	adminMW := CSPMiddleware(adminConfig)
	publicMW := CSPMiddleware(publicConfig)

	return func(next echo.HandlerFunc) echo.HandlerFunc {
		return func(c echo.Context) error {
			path := c.Request().URL.Path
			switch {
			case strings.HasPrefix(path, adminPrefix):
				return adminMW(next)(c)
			case strings.HasPrefix(path, publicPrefix):
				return publicMW(next)(c)
			default:
				return next(c)
			}
		}
	}
}

// CSPReportHandler handles CSP violation reports
func CSPReportHandler(c echo.Context) error {
	var report struct {
		CSPReport struct {
			DocumentURI        string `json:"document-uri"`
			ViolatedDirective  string `json:"violated-directive"`
			EffectiveDirective string `json:"effective-directive"`
			OriginalPolicy     string `json:"original-policy"`
			BlockedURI         string `json:"blocked-uri"`
			StatusCode         int    `json:"status-code"`
		} `json:"csp-report"`
	}

	if err := c.Bind(&report); err != nil {
		return c.JSON(http.StatusBadRequest, map[string]string{"error": "Invalid CSP report"})
	}

	// Log CSP violation (implement proper logging here)
	// logger.Warn("CSP Violation", map[string]interface{}{
	//     "document_uri":        report.CSPReport.DocumentURI,
	//     "violated_directive":  report.CSPReport.ViolatedDirective,
	//     "blocked_uri":         report.CSPReport.BlockedURI,
	// })

	return c.JSON(http.StatusOK, map[string]string{"status": "reported"})
}

func toStringSlice(sources []string) []string {
	result := make([]string, len(sources))
	for i, src := range sources {
		result[i] = string(src)
	}
	return result
}