package web

import (
	"bytes"
	"crypto/sha256"
	"encoding/base64"
	"net/http"
	"regexp"
	"slices"
	"strings"
)

// inlineScript is a script element with its code in the page rather than
// in a src: what a policy names by its hash.
var inlineScript = regexp.MustCompile(`(?is)<script(\s[^>]*)?>(.*?)</script>`)

// contentSecurityPolicy is the policy for a rendered page: scripts from this
// server and the page's own inline ones, each by its hash, and nothing else
// — no inline handlers, no eval, no script a reply or a model card manages
// to put in the page. Styles may be inline (the pages and htmx use them),
// and fonts come from Google's for the public site.
//
// Hashed from the page as rendered, not from the templates: html/template
// drops comments inside scripts, so the template's text is not what the
// browser hashes.
func contentSecurityPolicy(page []byte) string {
	var hashes []string
	for _, m := range inlineScript.FindAllSubmatch(page, -1) {
		attrs := strings.ToLower(string(m[1]))
		if strings.Contains(attrs, "src=") || (strings.Contains(attrs, "type=") && !strings.Contains(attrs, "javascript")) {
			continue // fetched, or data such as application/json
		}
		sum := sha256.Sum256(m[2])
		h := "'sha256-" + base64.StdEncoding.EncodeToString(sum[:]) + "'"
		if !slices.Contains(hashes, h) {
			hashes = append(hashes, h)
		}
	}
	return "default-src 'self'; " +
		"script-src 'self'" + prefixSpace(strings.Join(hashes, " ")) + "; " +
		"style-src 'self' 'unsafe-inline' https://fonts.googleapis.com; " +
		"font-src 'self' https://fonts.gstatic.com; " +
		"img-src 'self' data: https:; " + // model cards show their images
		"connect-src 'self'; " +
		"object-src 'none'; base-uri 'self'; form-action 'self'; frame-ancestors 'none'"
}

func prefixSpace(s string) string {
	if s == "" {
		return ""
	}
	return " " + s
}

// securePage sets the headers every page is served with, and writes it.
func securePage(h http.Header, page *bytes.Buffer) {
	h.Set("Content-Security-Policy", contentSecurityPolicy(page.Bytes()))
	h.Set("X-Content-Type-Options", "nosniff")
	if h.Get("Referrer-Policy") == "" {
		h.Set("Referrer-Policy", "same-origin")
	}
}