package main

// What the store's reader may read (novox/hq issue 337).
//
// The reader (Reader) is the login every caller's statement runs as: the store seat's `query` verb and
// postgres_query. Until issue 337 it was granted `pg_read_all_data`, so it could read every column of
// every database the store holds — the controller's signing key and its sealed secrets among them — and
// anybody who could call the verb could read them through it.
//
// **Deny by default.** The reader holds no role and no grant but what this file gives it, and it gives
// only what a consuming module declares readable in its contribution to `postgres-database`:
//
//	"contributes": { "postgres-database": { "name": "letta", "readable": ["agents", "steps(id, created_at)"] } }
//
// An entry is `[schema.]table` — every column it has at each pass — or `[schema.]table(col, col)` — those
// columns only, spelled as the catalogue spells them. A database no contribution names (the controller's
// own among them) gives the reader nothing. A list is used rather than an object so the controller's
// reading of `contributes` cannot mistake it for the several-instance shape.
//
// **A column or table whose name says it holds a secret is never granted, declared or not**
// (SecretName): private keys, sealed values, tokens, passwords, credentials, salts. The refusal is said
// in the reader's report (postgres_reader_grants), so a module that meant it learns why.
//
// **Reconciled every pass, on every database, not only made once.** The reader's role memberships are
// revoked (pg_read_all_data included, which is what an existing install holds), and in every database
// every grant the reader holds that is not declared is revoked, and every declared one that is missing
// is granted. A narrowed declaration narrows the reader at the next pass.
//
// **Checked where it is used.** Before a caller's statement runs, the reader's own session is asked
// which columns it can read whose name says secret; if any, the statement is refused and the columns
// named (ExposureProbe). That catches what the reconcile does not own — a grant to PUBLIC, a grant made
// by hand — and fails closed. And a statement that names a known secret table or a secret-sounding name
// is refused before it is sent (refusedStatement): defence in depth, not the guard — the privileges are.

import (
	"context"
	"encoding/json"
	"fmt"
	"os"
	"regexp"
	"sort"
	"strings"
	"sync"
	"time"
)

// ReadableTable is one entry of a contribution's `readable`.
type ReadableTable struct {
	Schema  string
	Table   string
	Columns []string // nil: every column the table has
}

func (r ReadableTable) String() string {
	s := r.Schema + "." + r.Table
	if r.Columns != nil {
		s += "(" + strings.Join(r.Columns, ", ") + ")"
	}
	return s
}

var (
	readableEntry = regexp.MustCompile(`^(?:([A-Za-z_][A-Za-z0-9_$]*)\.)?([A-Za-z_][A-Za-z0-9_$]*)\s*(?:\(([^()]*)\))?$`)
	identifier    = regexp.MustCompile(`^[A-Za-z_][A-Za-z0-9_$]*$`)
)

// ParseReadable reads a contribution's `readable`: absent is nothing readable. Anything malformed refuses
// the whole list, so a typo never grants more than was meant.
func ParseReadable(values map[string]any) ([]ReadableTable, error) {
	raw, ok := values["readable"]
	if !ok || raw == nil {
		return nil, nil
	}
	list, ok := raw.([]any)
	if !ok {
		return nil, fmt.Errorf("readable must be a list of tables, as \"table\" or \"table(column, column)\", not %T", raw)
	}
	var out []ReadableTable
	seen := map[string]bool{}
	for _, v := range list {
		entry, ok := v.(string)
		if !ok {
			return nil, fmt.Errorf("readable: %v is not a table name", v)
		}
		m := readableEntry.FindStringSubmatch(strings.TrimSpace(entry))
		if m == nil {
			return nil, fmt.Errorf("readable: %q is not \"[schema.]table\" or \"[schema.]table(column, …)\"", entry)
		}
		t := ReadableTable{Schema: firstOf(m[1], "public"), Table: m[2]}
		if strings.Contains(entry, "(") {
			t.Columns = []string{}
			for _, col := range strings.Split(m[3], ",") {
				col = strings.TrimSpace(col)
				if !identifier.MatchString(col) {
					return nil, fmt.Errorf("readable: %q names a column %q that is not a plain name", entry, col)
				}
				t.Columns = append(t.Columns, col)
			}
		}
		key := t.Schema + "." + t.Table
		if seen[key] {
			return nil, fmt.Errorf("readable: %s is named twice", key)
		}
		seen[key] = true
		out = append(out, t)
	}
	return out, nil
}

// secretName is what a column or table is called when it holds a secret, or something a secret is
// made from: a private key, a sealed value, a token, a password or its hash, a credential, a salt, a
// one-time-password seed, or a key of any kind. Wide on purpose — a public key refused is a small cost;
// a private one granted is the whole mesh.
var secretName = regexp.MustCompile(`(?i)(private|secret|seal|token|passw|pwd|credential|salt|cookie|nonce|` +
	`api_?key|(^|_)(otp|totp|mfa|key|keys|rands)(_|$))`)

// SecretName says whether a column or table name says it holds a secret.
func SecretName(name string) bool { return secretName.MatchString(name) }

// knownSecretTables are tables whose secret is not in their columns' names — the mesh's own, in the
// controller's databases. A statement naming one is refused before it is sent.
var knownSecretTables = map[string][]string{
	"identity":  {"signing_key", "authority", "node_key"},
	"inventory": {"secret", "module_secret", "operator_key", "bus_membership", "bus_user", "enrolment_token"},
	"licences":  {"licence_holder", "licence_manager"},
}

// statementWord is what the statement guard reads: every word of the statement, literals included, so
// a name in a string handed to a function is read as well.
var statementWord = regexp.MustCompile(`[A-Za-z_][A-Za-z0-9_$]*`)

// guardName is narrower than secretName, because it is read over a caller's whole statement, where
// `key` and `token` are ordinary words (jsonb_object_keys).
var guardName = regexp.MustCompile(`(?i)(private|secret|sealed|passw|credential|api_?key|signing_key|` +
	`sealing_key|overlay_key|operator_key|refresh_token|access_token|authtoken|pg_authid|pg_shadow|pg_user_mapping)`)

// refusedStatement says why a caller's statement is refused before it is sent, or "".
func refusedStatement(database, sql string) string {
	known := map[string]bool{}
	for _, t := range knownSecretTables[database] {
		known[t] = true
	}
	for _, w := range statementWord.FindAllString(sql, -1) {
		lw := strings.ToLower(w)
		if known[lw] {
			return fmt.Sprintf("the statement names %s.%s, which holds the mesh's keys or secrets; the store's reader is "+
				"never given it (novox/hq issue 337)", database, lw)
		}
		if guardName.MatchString(lw) {
			return fmt.Sprintf("the statement names %q, which reads as a key, secret or credential; the store's reader is "+
				"never given one, and statements naming one are refused (novox/hq issue 337)", w)
		}
	}
	return ""
}

// ---- the probe --------------------------------------------------------------------------------

// ExposureProbe lists every column the session's own login may SELECT outside the system schemas — by
// table grant, column grant, PUBLIC or a role it holds. Exposed filters it to the secret ones.
const ExposureProbe = `SELECT n.nspname, c.relname, a.attname FROM pg_catalog.pg_attribute a ` +
	`JOIN pg_catalog.pg_class c ON c.oid = a.attrelid JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace ` +
	`WHERE c.relkind IN ('r','v','m','p','f') AND a.attnum > 0 AND NOT a.attisdropped ` +
	`AND n.nspname NOT IN ('pg_catalog','information_schema') AND n.nspname NOT LIKE 'pg\_toast%' ` +
	`AND pg_catalog.has_column_privilege(c.oid, a.attnum, 'SELECT') ORDER BY 1, 2, 3`

// Exposed is the probe's rows that name a secret: "schema.table.column".
func Exposed(r Result) []string {
	var out []string
	for _, row := range r.Rows {
		if len(row) < 3 || row[0] == nil || row[1] == nil || row[2] == nil {
			continue
		}
		if SecretName(*row[1]) || SecretName(*row[2]) {
			out = append(out, *row[0]+"."+*row[1]+"."+*row[2])
		}
	}
	return out
}

// ---- the plan ---------------------------------------------------------------------------------

// column is one column of one table.
type column struct{ Schema, Table, Name string }

type table struct{ Schema, Table string }

// held is what the reader holds in one database, read from the catalogue as the admin.
type held struct {
	// Tables are those the reader holds any table-level privilege on, with their kind ('S': sequence).
	Tables map[table]string
	// Columns are the column privileges it holds, each column's privilege types.
	Columns map[column][]string
}

// DatabaseReport is what the reader may read in one database, and what was refused or missing.
type DatabaseReport struct {
	Database string   `json:"database"`
	Declared bool     `json:"declared"`
	Readable []string `json:"readable,omitempty"`
	Refused  []string `json:"refused,omitempty"`
	Missing  []string `json:"missing,omitempty"`
	Revoked  []string `json:"revoked,omitempty"`
	Exposed  []string `json:"exposed,omitempty"`
	Error    string   `json:"error,omitempty"`
}

// planDatabase is the DDL that brings the reader's grants in one database to exactly what is declared,
// minus every secret: revoke what is held and not wanted, grant what is wanted and not held.
// existing is every column of every declared table that the database has. Pure; the tests' subject.
func planDatabase(decl []ReadableTable, existing []column, h held, noUsage map[string]bool) ([]string, DatabaseReport) {
	var rep DatabaseReport
	byTable := map[table][]string{}
	for _, c := range existing {
		t := table{c.Schema, c.Table}
		byTable[t] = append(byTable[t], c.Name)
	}
	want := map[column]bool{}
	for _, d := range decl {
		t := table{d.Schema, d.Table}
		cols, ok := byTable[t]
		if !ok {
			rep.Missing = append(rep.Missing, d.String())
			continue
		}
		if SecretName(d.Table) {
			rep.Refused = append(rep.Refused, d.Schema+"."+d.Table+" (its name says it holds secrets)")
			continue
		}
		named := cols
		if d.Columns != nil {
			// A declared column the table lacks is said and left out: granting it would fail the
			// database's whole change.
			named = nil
			has := map[string]bool{}
			for _, c := range cols {
				has[c] = true
			}
			for _, c := range d.Columns {
				if has[c] {
					named = append(named, c)
				} else {
					rep.Missing = append(rep.Missing, d.Schema+"."+d.Table+"."+c)
				}
			}
		}
		for _, c := range named {
			if SecretName(c) {
				rep.Refused = append(rep.Refused, d.Schema+"."+d.Table+"."+c+" (its name says it holds a secret)")
				continue
			}
			want[column{d.Schema, d.Table, c}] = true
		}
	}
	// Every table-level privilege is revoked: the reader is only ever given columns. REVOKE on a table
	// takes its column privileges with it, so those are granted again below if wanted.
	var stmts []string
	cleared := map[table]bool{}
	for _, t := range sortedTables(h.Tables) {
		kind := "TABLE"
		if h.Tables[t] == "S" {
			kind = "SEQUENCE"
		}
		stmts = append(stmts, fmt.Sprintf("REVOKE ALL ON %s %s.%s FROM %s", kind, Ident(t.Schema), Ident(t.Table), Ident(Reader)))
		rep.Revoked = append(rep.Revoked, t.Schema+"."+t.Table+" (every privilege on the table)")
		cleared[t] = true
	}
	have := map[column]bool{}
	unwanted := map[table][]string{}
	for c, privs := range h.Columns {
		if cleared[table{c.Schema, c.Table}] {
			continue
		}
		onlySelect := len(privs) == 1 && privs[0] == "SELECT"
		if want[c] && onlySelect {
			have[c] = true
			continue
		}
		// Not wanted, or held with more than SELECT: revoked, and granted SELECT again below if wanted.
		t := table{c.Schema, c.Table}
		unwanted[t] = append(unwanted[t], c.Name)
	}
	for _, t := range sortedKeys(unwanted) {
		cols := unwanted[t]
		sort.Strings(cols)
		stmts = append(stmts, fmt.Sprintf("REVOKE ALL (%s) ON %s.%s FROM %s", idents(cols), Ident(t.Schema), Ident(t.Table), Ident(Reader)))
		for _, c := range cols {
			rep.Revoked = append(rep.Revoked, t.Schema+"."+t.Table+"."+c)
		}
	}
	grant := map[table][]string{}
	schemas := map[string]bool{}
	for c := range want {
		rep.Readable = append(rep.Readable, c.Schema+"."+c.Table+"."+c.Name)
		if !have[c] {
			t := table{c.Schema, c.Table}
			grant[t] = append(grant[t], c.Name)
			schemas[c.Schema] = true
		}
	}
	sort.Strings(rep.Readable)
	for _, s := range sortedSet(schemas) {
		if noUsage[s] {
			stmts = append(stmts, fmt.Sprintf("GRANT USAGE ON SCHEMA %s TO %s", Ident(s), Ident(Reader)))
		}
	}
	for _, t := range sortedKeys(grant) {
		cols := grant[t]
		sort.Strings(cols)
		stmts = append(stmts, fmt.Sprintf("GRANT SELECT (%s) ON %s.%s TO %s", idents(cols), Ident(t.Schema), Ident(t.Table), Ident(Reader)))
	}
	return stmts, rep
}

func idents(names []string) string {
	out := make([]string, len(names))
	for i, n := range names {
		out[i] = Ident(n)
	}
	return strings.Join(out, ", ")
}

func sortedTables(m map[table]string) []table {
	out := make([]table, 0, len(m))
	for t := range m {
		out = append(out, t)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Schema+"."+out[i].Table < out[j].Schema+"."+out[j].Table })
	return out
}

func sortedKeys(m map[table][]string) []table {
	out := make([]table, 0, len(m))
	for t := range m {
		out = append(out, t)
	}
	sort.Slice(out, func(i, j int) bool { return out[i].Schema+"."+out[i].Table < out[j].Schema+"."+out[j].Table })
	return out
}

func sortedSet(m map[string]bool) []string {
	out := make([]string, 0, len(m))
	for s := range m {
		out = append(out, s)
	}
	sort.Strings(out)
	return out
}

// ---- the reconcile ------------------------------------------------------------------------------

// Declarations are what each consumer's database lets the reader read, by database; a database whose
// declaration was malformed is in Malformed and is given nothing.
type Declarations struct {
	Tables    map[string][]ReadableTable
	Malformed map[string]string
}

// DeclarationsFrom reads the contributions file the mesh delivers (MESH_RECEIVES): each consumer's
// database is its login (`as`), and its `readable` is in its values. Unreadable is no declarations:
// the reader is then given nothing, and the reason is answered.
func DeclarationsFrom(path string) (Declarations, error) {
	d := Declarations{Tables: map[string][]ReadableTable{}, Malformed: map[string]string{}}
	if path == "" {
		return d, fmt.Errorf("no contributions file is named (MESH_RECEIVES); the reader is given nothing")
	}
	raw, err := os.ReadFile(path)
	if err != nil {
		return d, fmt.Errorf("contributions file unreadable (%s), the reader is given nothing: %w", path, err)
	}
	var doc struct {
		Given []struct {
			As     string         `json:"as"`
			Values map[string]any `json:"values"`
		} `json:"given"`
	}
	if err := json.Unmarshal(raw, &doc); err != nil {
		return d, fmt.Errorf("contributions file is not JSON (%s), the reader is given nothing: %w", path, err)
	}
	for _, g := range doc.Given {
		if g.As == "" {
			continue
		}
		tables, err := ParseReadable(g.Values)
		if err != nil {
			d.Malformed[g.As] = err.Error()
			continue
		}
		if len(tables) > 0 {
			d.Tables[g.As] = tables
		}
	}
	return d, nil
}

// ReaderReport is the last reconcile's answer, for postgres_reader_grants.
type ReaderReport struct {
	At        time.Time        `json:"at"`
	Problem   string           `json:"problem,omitempty"`
	Databases []DatabaseReport `json:"databases"`
}

type readerState struct {
	mu   sync.Mutex
	last ReaderReport
}

// ReconcileReader brings the reader to exactly what is declared, in every database the server holds.
// Each database is its own: one that fails is said in its report and the rest go on.
func (c *Client) ReconcileReader(ctx context.Context, d Declarations, problem string) ReaderReport {
	rep := ReaderReport{At: time.Now().UTC(), Problem: problem}
	if err := c.ensureReaderRole(ctx); err != nil {
		rep.Problem = strings.TrimSpace(rep.Problem + "; the reader's role: " + err.Error())
		c.keep(rep)
		return rep
	}
	r, err := c.Query(ctx, "", "SELECT datname FROM pg_database WHERE datallowconn AND NOT datistemplate ORDER BY datname")
	if err != nil {
		rep.Problem = strings.TrimSpace(rep.Problem + "; listing the databases: " + err.Error())
		c.keep(rep)
		return rep
	}
	present := map[string]bool{}
	for _, row := range r.Rows {
		if len(row) == 0 || row[0] == nil {
			continue
		}
		db := *row[0]
		present[db] = true
		dr := c.reconcileDatabase(ctx, db, d.Tables[db])
		if why, ok := d.Malformed[db]; ok {
			dr.Error = strings.TrimSpace(dr.Error + " its readable declaration is refused, nothing granted: " + why)
		}
		rep.Databases = append(rep.Databases, dr)
	}
	for db := range d.Tables {
		if !present[db] {
			rep.Databases = append(rep.Databases, DatabaseReport{Database: db, Declared: true, Error: "declared readable, but the server holds no such database"})
		}
	}
	sort.Slice(rep.Databases, func(i, j int) bool { return rep.Databases[i].Database < rep.Databases[j].Database })
	c.keep(rep)
	return rep
}

func (c *Client) keep(rep ReaderReport) {
	c.reader.mu.Lock()
	c.reader.last = rep
	c.reader.mu.Unlock()
}

// LastReaderReport is what the last reconcile found.
func (c *Client) LastReaderReport() ReaderReport {
	c.reader.mu.Lock()
	defer c.reader.mu.Unlock()
	return c.reader.last
}

func (c *Client) reconcileDatabase(ctx context.Context, db string, decl []ReadableTable) DatabaseReport {
	fail := func(err error) DatabaseReport {
		return DatabaseReport{Database: db, Declared: decl != nil, Error: err.Error()}
	}
	h := held{Tables: map[table]string{}, Columns: map[column][]string{}}
	readerOID := "(SELECT oid FROM pg_catalog.pg_roles WHERE rolname = " + Literal(Reader) + ")"
	r, err := c.Query(ctx, db, "SELECT n.nspname, c.relname, c.relkind::text FROM pg_catalog.pg_class c "+
		"JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace CROSS JOIN LATERAL pg_catalog.aclexplode(c.relacl) x "+
		"WHERE x.grantee = "+readerOID+" GROUP BY 1, 2, 3")
	if err != nil {
		return fail(err)
	}
	for _, row := range r.Rows {
		h.Tables[table{cell(row, 0), cell(row, 1)}] = cell(row, 2)
	}
	r, err = c.Query(ctx, db, "SELECT n.nspname, c.relname, a.attname, x.privilege_type FROM pg_catalog.pg_attribute a "+
		"JOIN pg_catalog.pg_class c ON c.oid = a.attrelid JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace "+
		"CROSS JOIN LATERAL pg_catalog.aclexplode(a.attacl) x WHERE x.grantee = "+readerOID)
	if err != nil {
		return fail(err)
	}
	for _, row := range r.Rows {
		k := column{cell(row, 0), cell(row, 1), cell(row, 2)}
		h.Columns[k] = append(h.Columns[k], cell(row, 3))
	}
	var existing []column
	noUsage := map[string]bool{}
	if len(decl) > 0 {
		var pairs []string
		for _, d := range decl {
			pairs = append(pairs, "("+Literal(d.Schema)+", "+Literal(d.Table)+")")
		}
		r, err = c.Query(ctx, db, "SELECT n.nspname, c.relname, a.attname FROM pg_catalog.pg_attribute a "+
			"JOIN pg_catalog.pg_class c ON c.oid = a.attrelid JOIN pg_catalog.pg_namespace n ON n.oid = c.relnamespace "+
			"WHERE c.relkind IN ('r','v','m','p','f') AND a.attnum > 0 AND NOT a.attisdropped "+
			"AND (n.nspname, c.relname) IN ("+strings.Join(pairs, ", ")+") ORDER BY a.attnum")
		if err != nil {
			return fail(err)
		}
		for _, row := range r.Rows {
			existing = append(existing, column{cell(row, 0), cell(row, 1), cell(row, 2)})
		}
		r, err = c.Query(ctx, db, "SELECT nspname FROM pg_catalog.pg_namespace WHERE NOT pg_catalog.has_schema_privilege("+
			Literal(Reader)+", oid, 'USAGE')")
		if err != nil {
			return fail(err)
		}
		for _, row := range r.Rows {
			noUsage[cell(row, 0)] = true
		}
	}
	stmts, rep := planDatabase(decl, existing, h, noUsage)
	rep.Database, rep.Declared = db, decl != nil
	if len(stmts) > 0 {
		// One string: the simple protocol runs it as one implicit transaction, all or nothing.
		if _, err := c.Query(ctx, db, strings.Join(stmts, "; ")); err != nil {
			rep.Error = "changing the reader's grants: " + err.Error()
			return rep
		}
	}
	// Asked as the admin on the reader's behalf: whatever else gives the reader a secret (PUBLIC, a
	// grant made by hand) is found here, said, and makes the reader's statements against it refused.
	probe := strings.Replace(ExposureProbe, "pg_catalog.has_column_privilege(c.oid", "pg_catalog.has_column_privilege("+Literal(Reader)+", c.oid", 1)
	if r, err = c.Query(ctx, db, probe); err != nil {
		rep.Error = "probing what the reader can read: " + err.Error()
		return rep
	}
	rep.Exposed = Exposed(r)
	return rep
}

// RunReader reconciles the reader at once and then every `every`, until ctx ends, reading the
// declarations from the contributions file each time. A change, a refusal and an exposure are said.
func (c *Client) RunReader(ctx context.Context, receives string, every time.Duration, say func(string, ...any)) {
	lastSaid := ""
	for {
		d, err := DeclarationsFrom(receives)
		problem := ""
		if err != nil {
			problem = err.Error()
		}
		rep := c.ReconcileReader(ctx, d, problem)
		var lines []string
		if rep.Problem != "" {
			lines = append(lines, "reader: "+rep.Problem)
		}
		for _, db := range rep.Databases {
			if len(db.Exposed) > 0 {
				lines = append(lines, fmt.Sprintf("reader: EXPOSED in %s — the reader can read %s; its statements there are refused until it cannot",
					db.Database, strings.Join(db.Exposed, ", ")))
			}
			if db.Error != "" {
				lines = append(lines, fmt.Sprintf("reader: %s: %s", db.Database, db.Error))
			}
			if len(db.Refused) > 0 {
				lines = append(lines, fmt.Sprintf("reader: %s: not granted, %s", db.Database, strings.Join(db.Refused, "; ")))
			}
			if len(db.Revoked) > 0 {
				say("reader: %s: revoked %s", db.Database, strings.Join(db.Revoked, ", "))
			}
		}
		// The same standing is said once, not every pass.
		if now := strings.Join(lines, "\n"); now != lastSaid {
			for _, l := range lines {
				say("%s", l)
			}
			lastSaid = now
		}
		select {
		case <-ctx.Done():
			return
		case <-time.After(every):
		}
	}
}