package middleware

import (
	"fmt"
	"net/http"
	"strings"
)

// inspired from Helmet.js
// https://github.com/helmetjs/helmet/tree/main

type (
	HelmetOption struct {
		ContentSecurityPolicy CSP

		StrictTransportSecurity *TransportSecurity

		// "require-corp" will be the default policy
		CrossOriginEmbedderPolicy Embedder

		// "same-origin" will be the default policy
		CrossOriginOpenerPolicy Opener

		// "same-origin" will be the default policy
		CrossOriginResourcePolicy Resource

		// "no-referrer" will be the default policy
		ReferrerPolicy []Referrer

		OriginAgentCluster bool

		// set true to remove header "X-Content-Type-Options"
		DisableSniffMimeType bool

		// set true for header "X-DNS-Prefetch-Control: off"
		//
		// default is "X-DNS-Prefetch-Control: on"
		DisableDNSPrefetch bool

		// set true to remove header "X-Download-Options: noopen"
		DisableXDownload bool

		// X-Frame-Options
		XFrameOption XFrame

		// X-Permitted-Cross-Domain-Policies
		//
		// default value will be "none"
		CrossDomainPolicies CDP

		// X-XSS-Protection
		//
		// default is off
		XssProtection bool
	}

	// CSP  is Content-Security-Policy settings
	//
	// https://developer.mozilla.org/en-US/docs/Web/HTTP/Headers/Content-Security-Policy/Sources
	CSP struct {
		// default-src, default value will be 'self'
		DefaultSrc []string
		// script-src,  default value will be 'self'
		ScriptSrc []string
		// script-src-attr, default value will be 'none'
		ScriptSrcAttr []string
		// style-src, default value will be 'self' https: 'unsafe-inline'
		StyleSrc []string
		// img-src, default value will be 'self' data:
		ImgSrc []string
		// object-src, default value will be 'none'
		ObjectSrc []string
		// base-uri, default value will be 'self'
		BaseUri []string
		// font-src, default value will be 'self' https: data:
		FontSrc []string
		// form-action, default value will be 'self'
		FormAction []string
		// frame-ancestors, default value will be 'self'
		FrameAncestors []string

		UpgradeInsecureRequests bool
	}

	TransportSecurity struct {
		// Age in seconts
		MaxAge            uint
		IncludeSubDomains bool
		Preload           bool
	}

	Embedder string
	Opener   string
	Resource string
	Referrer string
	// CDP Cross-Domain-Policy
	CDP    string
	XFrame string
)

const (
	YearDuration = 365 * 24 * 60 * 60

	// EmbedderDefault default value will be "require-corp"
	EmbedderDefault        Embedder = ""
	EmbedderRequireCorp    Embedder = "require-corp"
	EmbedderCredentialLess Embedder = "credentialless"
	EmbedderUnsafeNone     Embedder = "unsafe-none"

	// OpenerDefault default value will be "same-origin"
	OpenerDefault               Opener = ""
	OpenerSameOrigin            Opener = "same-origin"
	OpenerSameOriginAllowPopups Opener = "same-origin-allow-popups"
	OpenerUnsafeNone            Opener = "unsafe-none"

	// ResourceDefault default value will be "same-origin"
	ResourceDefault     Resource = ""
	ResourceSameOrigin  Resource = "same-origin"
	ResourceSameSite    Resource = "same-site"
	ResourceCrossOrigin Resource = "cross-origin"

	NoReferrer                  Referrer = "no-referrer"
	NoReferrerWhenDowngrade     Referrer = "no-referrer-when-downgrade"
	SameOrigin                  Referrer = "same-origin"
	Origin                      Referrer = "origin"
	StrictOrigin                Referrer = "strict-origin"
	OriginWhenCrossOrigin       Referrer = "origin-when-cross-origin"
	StrictOriginWhenCrossOrigin Referrer = "strict-origin-when-cross-origin"
	UnsafeUrl                   Referrer = "unsafe-url"

	// CDPDefault default value is  "none"
	CDPDefault       CDP = ""
	CDPNone          CDP = "none"
	CDPMasterOnly    CDP = "master-only"
	CDPByContentType CDP = "by-content-type"
	CDPAll           CDP = "all"

	// XFrameDefault default value will be "sameorigin"
	XFrameDefault    XFrame = ""
	XFrameSameOrigin XFrame = "sameorigin"
	XFrameDeny       XFrame = "deny"
)

// Helmet headers to secure server response
func Helmet(opt HelmetOption) func(http.Handler) http.Handler {
	return func(h http.Handler) http.Handler {
		return http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
			w.Header().Add("Content-Security-Policy", opt.ContentSecurityPolicy.value())

			// Cross-Origin-Embedder-Policy, if nil set default
			if opt.CrossOriginEmbedderPolicy == EmbedderDefault {
				w.Header().Add("Cross-Origin-Embedder-Policy", string(EmbedderRequireCorp))
			} else {
				w.Header().Add("Cross-Origin-Embedder-Policy", string(opt.CrossOriginEmbedderPolicy))
			}

			// Cross-Origin-Opener-Policy, if nil set default
			if opt.CrossOriginOpenerPolicy == OpenerDefault {
				w.Header().Add("Cross-Origin-Opener-Policy", string(OpenerSameOrigin))
			} else {
				w.Header().Add("Cross-Origin-Opener-Policy", string(opt.CrossOriginOpenerPolicy))
			}

			// Cross-Origin-Resource-Policy, if nil set default
			if opt.CrossOriginResourcePolicy == ResourceDefault {
				w.Header().Add("Cross-Origin-Resource-Policy", string(ResourceSameOrigin))
			} else {
				w.Header().Add("Cross-Origin-Resource-Policy", string(opt.CrossOriginResourcePolicy))
			}

			// Referrer-Policy
			rpCount := len(opt.ReferrerPolicy)
			if rpCount > 0 {
				refP := make([]string, rpCount)
				for i, r := range opt.ReferrerPolicy {
					refP[i] = string(r)
				}
				w.Header().Add("Referrer-Policy", string(NoReferrer))
			} else {
				// default no referer
				w.Header().Add("Referrer-Policy", string(NoReferrer))
			}

			// Origin-Agent-Cluster
			if opt.OriginAgentCluster {
				w.Header().Add("Origin-Agent-Cluster", "?1")
			}

			// Strict-Transport-Security
			if opt.StrictTransportSecurity != nil {
				var sb strings.Builder

				if opt.StrictTransportSecurity.MaxAge == 0 {
					opt.StrictTransportSecurity.MaxAge = YearDuration
				}

				sb.WriteString(fmt.Sprintf("max-age=%d", opt.StrictTransportSecurity.MaxAge))
				if opt.StrictTransportSecurity.IncludeSubDomains {
					sb.WriteString("; includeSubDomains")
				}

				if opt.StrictTransportSecurity.Preload {
					sb.WriteString("; preload")
				}

				w.Header().Add("Strict-Transport-Security", sb.String())
			}

			if !opt.DisableSniffMimeType {
				// MIME types advertised in the Content-Current headers should be followed and not be changed
				w.Header().Add("X-Content-Type-Options", "nosniff")
			}

			if opt.DisableDNSPrefetch {
				w.Header().Add("X-DNS-Prefetch-Control", "off")
			} else {
				w.Header().Add("X-DNS-Prefetch-Control", "on")
			}

			if !opt.DisableXDownload {
				// Instructs Internet Explorer not to open the file directly but to offer it for download first.
				w.Header().Add("X-Download-Options", "noopen")
			}

			// indicate whether a browser should be allowed to render a page in iframe | frame | embed | object
			if opt.XFrameOption == XFrameDefault {
				w.Header().Add("X-Frame-Options", string(XFrameSameOrigin))
			} else {
				w.Header().Add("X-Frame-Options", string(opt.XFrameOption))
			}

			if opt.CrossDomainPolicies == CDPDefault {
				w.Header().Add("X-Permitted-Cross-Domain-Policies", string(CDPNone))
			} else {
				w.Header().Add("X-Permitted-Cross-Domain-Policies", string(opt.CrossDomainPolicies))
			}

			w.Header().Del("X-Powered-By")

			if opt.XssProtection {
				// feature of IE, Chrome and Safari that stops pages from loading when they detect reflected
				// cross-site scripting (XSS) attacks.
				w.Header().Add("X-Xss-Protection", "1; mode=block")
			} else {
				// Following a decision by Google Chrome developers to disable Auditor,
				// developers should be able to disable the auditor for older browsers and set it to 0.
				// The X-XSS-PROTECTION header was found to have a multitude of issues, instead of helping the
				// developers protect their application.
				w.Header().Add("X-Xss-Protection", "0")
			}

			h.ServeHTTP(w, r)
		})
	}
}

func (csp *CSP) value() string {
	var sb strings.Builder

	// should be the first thing
	if csp.UpgradeInsecureRequests {
		sb.WriteString("upgrade-insecure-requests;")
	}

	sb.WriteString(fmt.Sprintf(
		"default-src %s; ",
		cspNormalised(csp.DefaultSrc, []string{"self"}),
	))

	sb.WriteString(fmt.Sprintf(
		"script-src %s; ",
		cspNormalised(csp.ScriptSrc, []string{"self"}),
	))

	sb.WriteString(fmt.Sprintf(
		"script-src-attr %s; ",
		cspNormalised(csp.ScriptSrcAttr, []string{"none"}),
	))

	sb.WriteString(fmt.Sprintf(
		"style-src %s; ",
		cspNormalised(csp.StyleSrc, []string{"self", "https:", "unsafe-inline"}),
	))

	sb.WriteString(fmt.Sprintf(
		"img-src %s; ",
		cspNormalised(csp.ImgSrc, []string{"self", "data:"}),
	))

	sb.WriteString(fmt.Sprintf(
		"object-src %s; ",
		cspNormalised(csp.ObjectSrc, []string{"none"}),
	))

	sb.WriteString(fmt.Sprintf(
		"base-uri %s; ",
		cspNormalised(csp.BaseUri, []string{"self"}),
	))

	sb.WriteString(fmt.Sprintf(
		"font-src %s; ",
		cspNormalised(csp.FontSrc, []string{"self", "https:", "data:"}),
	))

	sb.WriteString(fmt.Sprintf(
		"form-action %s; ",
		cspNormalised(csp.FormAction, []string{"self"}),
	))

	sb.WriteString(fmt.Sprintf(
		"frame-ancestors %s; ",
		cspNormalised(csp.FrameAncestors, []string{"self"}),
	))

	return sb.String()
}

func cspNormalised(v, defaultVal []string) string {
	if len(v) == 0 {
		v = defaultVal
	}

	var sb strings.Builder
	for _, val := range v {
		val = strings.TrimSpace(val)
		if val == "" {
			continue
		}

		sb.WriteString(" " + cspQuoted(val))
	}

	return strings.TrimSpace(sb.String())
}

func cspQuoted(v string) string {
	switch v {
	case "none",
		"self",
		"strict-dynamic",
		"report-sample",
		"inline-speculation-rules",
		"unsafe-inline",
		"unsafe-eval",
		"unsafe-hashes",
		"wasm-unsafe-eval":
		return fmt.Sprintf("'%s'", v)
	default:
		return v
	}
}