package frontend

import (
	"crypto/sha256"
	"encoding/base64"
	"regexp"
	"strings"
)

const scriptSrcSelf = "script-src 'self'"

var inlineScriptPattern = regexp.MustCompile(`(?is)<script(\s[^>]*)?>(.*?)</script>`)
var srcAttributePattern = regexp.MustCompile(`(?i)\ssrc\s*=`)

func inlineScriptSources(html []byte) []string {
	var sources []string
	for _, match := range inlineScriptPattern.FindAllSubmatch(html, -1) {
		if srcAttributePattern.Match(match[1]) || strings.TrimSpace(string(match[2])) == "" {
			continue
		}
		digest := sha256.Sum256(match[2])
		sources = append(sources, "'sha256-"+base64.StdEncoding.EncodeToString(digest[:])+"'")
	}
	return sources
}

func AllowInlineScripts(policy string, html []byte) string {
	sources := inlineScriptSources(html)
	if len(sources) == 0 {
		return policy
	}
	return strings.Replace(policy, scriptSrcSelf, scriptSrcSelf+" "+strings.Join(sources, " "), 1)
}