llmgw/internal/cli/keys.go

package cli

import (
	"context"
	"crypto/rand"
	"encoding/base32"
	"errors"
	"fmt"
	"math"
	"slices"
	"strings"
	"text/tabwriter"
	"time"

	"github.com/spf13/cobra"

	"gitea.mynodes.xyz/ilker/llmgw/internal/store"
)

func newTenancyCmd() *cobra.Command {
	var sf storeFlags

	cmd := &cobra.Command{
		Use:   "keys",
		Short: "Manage tenants and API keys",
		Long: "Works directly on the store file rather than through the admin API, so keys\n" +
			"can be issued before the server is running and without exposing an endpoint\n" +
			"that mints credentials.",
	}
	sf.bind(cmd)

	// Takes a context rather than the command, because some of these run
	// before the command tree is entered.
	open := func(ctx context.Context) (*store.DB, error) { return sf.openCtx(ctx) }

	cmd.AddCommand(
		newKeyInitCmd(open),
		newKeyCreateCmd(open),
		newKeyListCmd(open),
		newKeyRevokeCmd(open),
		newKeyRotateCmd(open),
		newKeyTTLCmd(open),
		newLimitsCmd(open),
		newThreadsCmd(open),
		newMemoryCmd(open),
		newMoveStoreCmd(),
	)
	return cmd
}

type opener func(context.Context) (*store.DB, error)

func newKeyInitCmd(open opener) *cobra.Command {
	var org, team, project string
	cmd := &cobra.Command{
		Use:   "init",
		Short: "Create the tenancy tree a key hangs off",
		Long: "Idempotent: running it again renames rather than duplicating, so it can be\n" +
			"kept in whatever provisions the machine.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			ctx := cmd.Context()
			if err := s.PutOrg(ctx, store.Org{ID: slug(org), Name: org}); err != nil {
				return err
			}
			if err := s.PutTeam(ctx, store.Team{ID: slug(team), OrgID: slug(org), Name: team}); err != nil {
				return err
			}
			if err := s.PutProject(ctx, store.Project{
				ID: slug(project), TeamID: slug(team), Name: project,
			}); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "%s / %s / %s\n", slug(org), slug(team), slug(project))
			return nil
		},
	}
	cmd.Flags().StringVar(&org, "org", "default", "organization name")
	cmd.Flags().StringVar(&team, "team", "default", "team name")
	cmd.Flags().StringVar(&project, "project", "default", "project name")
	return cmd
}

func newKeyCreateCmd(open opener) *cobra.Command {
	var project, band, profile, baseURL, forUser string
	var expires, claimTTL time.Duration
	var claim bool

	cmd := &cobra.Command{
		Use:   "create <name>",
		Short: "Issue an API key",
		Long: "The key is printed once and never again: only its hash is stored, so there\n" +
			"is nothing to recover it from. That is the point — a leaked database hands\n" +
			"over no working credential.\n\n" +
			"With --claim it prints a one-time link instead. Send the link rather than\n" +
			"the key: a key pasted into a chat window stays in that chat window forever,\n" +
			"while a link works once and then expires. The key is sealed with the link's\n" +
			"own token, which is not stored, so a pending claim is not a usable\n" +
			"credential sitting in the database either.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			// Three kinds: an operator's key in a project, a person's key in a
			// project they are a member of, and a person's personal key.
			var owner store.User
			if forUser != "" {
				if owner, err = resolveUser(cmd.Context(), s, forUser); err != nil {
					return err
				}
			}
			var p store.ProjectPath
			if project != "" || forUser == "" {
				if p, err = resolveProject(cmd.Context(), s, project); err != nil {
					return err
				}
			}
			if p.ID == "" {
				p.Path = "personal, " + owner.Email
			}
			secret, k, err := store.NewKey(randomID("key"), args[0], p.ID)
			if err != nil {
				return err
			}
			k.CreatedAt = time.Now()
			k.Band, k.RoutingProfile = band, profile
			if expires > 0 {
				k.ExpiresAt = k.CreatedAt.Add(expires)
			}
			if owner.ID != "" {
				err = s.CreateUserKey(cmd.Context(), owner.ID, k)
			} else {
				err = s.CreateKey(cmd.Context(), k)
			}
			if errors.Is(err, store.ErrNotMember) {
				return fmt.Errorf("%s is not a member of %s; `llmgw users join` first", owner.Email, p.Path)
			}
			if err != nil {
				return fmt.Errorf("creating the key in %s: %w", p.Path, err)
			}

			// The one output that matters goes to stdout and everything else to
			// stderr, so `keys create ... > out.txt` captures it and nothing
			// else. Cobra's Print writes to stderr, which would have made the
			// useful half the half that cannot be piped.
			out := cmd.OutOrStdout()

			if claim {
				if baseURL == "" {
					return errors.New("--claim needs --base-url, the address users reach this gateway at")
				}
				token, err := store.NewClaimToken()
				if err != nil {
					return err
				}
				if err := s.CreateClaim(cmd.Context(), k.ID, token, secret, claimTTL); err != nil {
					return err
				}
				link := strings.TrimRight(baseURL, "/") + "/claim/" + token
				fmt.Fprintln(out, link)
				cmd.PrintErrf("\nid       %s\nproject  %s\nexpires  %s\n",
					k.ID, p.Path, time.Now().Add(orDefault(claimTTL, store.ClaimTTL)).Format(time.RFC3339))
				cmd.PrintErrln("\nSend the link. It works once, and the key is shown only to whoever opens it.")
				return nil
			}

			fmt.Fprintln(out, secret)
			cmd.PrintErrf("\nid       %s\nproject  %s\n", k.ID, p.Path)
			if !k.ExpiresAt.IsZero() {
				cmd.PrintErrf("expires  %s\n", k.ExpiresAt.Format(time.RFC3339))
			}
			cmd.PrintErrln("\nThis is the only time the key is shown. Store it now.")
			return nil
		},
	}
	cmd.Flags().StringVar(&project, "project", "",
		"project the key belongs to: id, path (Acme/platform/gateway) or unique name")
	cmd.Flags().StringVar(&forUser, "for", "",
		"a person the key belongs to (id or email); without --project it is their personal key")
	cmd.Flags().StringVar(&band, "band", "", "admission band; empty inherits")
	cmd.Flags().StringVar(&profile, "profile", "", "routing profile; empty inherits")
	cmd.Flags().DurationVar(&expires, "expires-in", 0, "lifetime; 0 never expires")
	cmd.Flags().BoolVar(&claim, "claim", false, "print a one-time link instead of the key")
	cmd.Flags().StringVar(&baseURL, "base-url", "", "public address of the gateway, for --claim")
	cmd.Flags().DurationVar(&claimTTL, "claim-ttl", store.ClaimTTL, "how long the link stays usable")
	return cmd
}

func orDefault(d, fallback time.Duration) time.Duration {
	if d <= 0 {
		return fallback
	}
	return d
}

func newKeyListCmd(open opener) *cobra.Command {
	var project string
	var all bool
	cmd := &cobra.Command{
		Use:   "list",
		Short: "List a project's keys, or with --all every key and whose it is",
		Args:  cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			if all {
				rows, err := s.AllKeys(cmd.Context())
				if err != nil {
					return err
				}
				w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
				fmt.Fprintln(w, "ID\tNAME\tKIND\tKEY\tPROJECT\tOWNER\tSTATE\tEXPIRES\tLAST USED")
				now := time.Now()
				// Service keys — a project's, nobody's own — first: they are
				// the ones nobody leaving the company takes with them.
				slices.SortStableFunc(rows, func(a, b store.KeyRow) int {
					return strings.Compare(keyKind(a.Key), keyKind(b.Key))
				})
				for _, r := range rows {
					where := r.ProjectPath
					if where == "" {
						where = "personal"
					}
					fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\t%s\n", r.ID, r.Name, keyKind(r.Key), r.Prefix,
						where, dash(r.OwnerEmail), keyState(r.Key, now), until(r.ExpiresAt, now), ago(r.LastUsedAt, now))
				}
				return w.Flush()
			}

			p, err := resolveProject(cmd.Context(), s, project)
			if err != nil {
				return err
			}
			keys, err := s.Keys(cmd.Context(), p.ID)
			if err != nil {
				return err
			}
			if len(keys) == 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "no keys in %s\n", p.Path)
				return nil
			}

			w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
			fmt.Fprintln(w, "ID\tNAME\tKEY\tBAND\tSTATE\tLAST USED")
			now := time.Now()
			for _, k := range keys {
				fmt.Fprintf(w, "%s\t%s\t%s\t%s\t%s\t%s\n",
					k.ID, k.Name, k.Prefix, dash(k.Band), keyState(k, now), ago(k.LastUsedAt, now))
			}
			return w.Flush()
		},
	}
	cmd.Flags().StringVar(&project, "project", "", "project to list: id, path or unique name")
	cmd.Flags().BoolVar(&all, "all", false, "every key in every project, with its owner")
	return cmd
}

func newKeyRevokeCmd(open opener) *cobra.Command {
	return &cobra.Command{
		Use:   "revoke <key-id>",
		Short: "Stop a key working, without deleting it",
		Long: "A soft delete. Usage rows reference the key, and an audit trail that loses\n" +
			"its subject is not an audit trail.",
		Args: cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			if err := s.RevokeKey(cmd.Context(), args[0], time.Now()); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "revoked %s\n", args[0])
			cmd.PrintErrln("A running server keeps accepting it until its cache expires (auth.cache_ttl).")
			return nil
		},
	}
}

func newLimitsCmd(open opener) *cobra.Command {
	var scope, id, period string
	var perMember bool
	var rpm, tpm int
	var budget float64
	var models []string

	cmd := &cobra.Command{
		Use:   "limits",
		Short: "Set limits at one level of the tenancy tree",
		Long: "Limits compose downward by taking the most restrictive value, so a key\n" +
			"cannot raise what its organization set. Model allowlists intersect: a key\n" +
			"cannot add a model its team withheld.",
		Args: cobra.NoArgs,
		RunE: func(cmd *cobra.Command, _ []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()

			scopeID := id
			switch store.Scope(scope) {
			case store.ScopeOrg:
				o, err := resolveOrg(cmd.Context(), s, id)
				if err != nil {
					return err
				}
				scopeID = o.ID
			case store.ScopeTeam:
				t, err := resolveTeam(cmd.Context(), s, id)
				if err != nil {
					return err
				}
				scopeID = t.ID
			case store.ScopeProject:
				p, err := resolveProject(cmd.Context(), s, id)
				if err != nil {
					return err
				}
				scopeID = p.ID
			case store.ScopeUser:
				u, err := resolveUser(cmd.Context(), s, id)
				if err != nil {
					return err
				}
				scopeID = u.ID
			case store.ScopeKey:
				if id == "" {
					return errors.New("--scope key needs --id, the key's id")
				}
			default:
				return fmt.Errorf("--scope is org, team, project, user or key, not %q", scope)
			}
			level := store.Scope(scope)
			if perMember {
				// Written on the scope, applied to each member separately.
				if level = store.MemberScope(level); level == "" {
					return errors.New("--per-member is for --scope org, team or project")
				}
			}
			l := store.Limits{Scope: level, ScopeID: scopeID}
			if cmd.Flags().Changed("rpm") {
				l.RPM = &rpm
			}
			if cmd.Flags().Changed("tpm") {
				l.TPM = &tpm
			}
			if cmd.Flags().Changed("budget") {
				// Whole units, as the quotas page takes them; stored as
				// micro-units so repeated addition does not drift.
				micros := int64(math.Round(budget * 1e6))
				l.Budget = &micros
				l.BudgetPeriod = store.Period(period)
			}
			if cmd.Flags().Changed("models") {
				// An explicitly empty list is an allowlist that permits
				// nothing, which is different from setting no list at all.
				if models == nil {
					models = []string{}
				}
				l.AllowedModels = models
			}
			if err := s.PutLimits(cmd.Context(), l); err != nil {
				return err
			}
			fmt.Fprintf(cmd.OutOrStdout(), "limits set on %s %s\n", level, scopeID)
			return nil
		},
	}
	cmd.Flags().StringVar(&scope, "scope", "project", "org | team | project | user | key")
	cmd.Flags().BoolVar(&perMember, "per-member", false,
		"with org, team or project: a limit for each member there, not one shared by all")
	cmd.Flags().StringVar(&id, "id", "", "which one: id, path or unique name (a key's id for --scope key)")
	cmd.Flags().IntVar(&rpm, "rpm", 0, "requests per minute")
	cmd.Flags().IntVar(&tpm, "tpm", 0, "tokens per minute")
	cmd.Flags().Float64Var(&budget, "budget", 0, "budget per period, in whole units as the quotas page shows them")
	cmd.Flags().StringVar(&period, "period", "month", "day | month")
	cmd.Flags().StringSliceVar(&models, "models", nil, "allowed models; empty allows none")
	return cmd
}

// keyKind says whose a key is: a project's (service), a person's in a
// project (member), or a person's own (personal).
func keyKind(k store.Key) string {
	switch {
	case k.UserID == "":
		return "service"
	case k.ProjectID == "":
		return "personal"
	}
	return "member"
}

// until is how long a key has left, or "never".
func until(t, now time.Time) string {
	if t.IsZero() {
		return "never"
	}
	if !t.After(now) {
		return "-"
	}
	d := t.Sub(now)
	if d > 48*time.Hour {
		return fmt.Sprintf("in %dd", int(d.Hours()/24))
	}
	return "in " + d.Round(time.Minute).String()
}

func keyState(k store.Key, now time.Time) string {
	switch {
	case !k.RevokedAt.IsZero():
		return "revoked"
	case !k.ExpiresAt.IsZero() && !now.Before(k.ExpiresAt):
		return "expired"
	default:
		return "active"
	}
}

func ago(t, now time.Time) string {
	if t.IsZero() {
		return "never"
	}
	return now.Sub(t).Round(time.Second).String() + " ago"
}

// slug turns a human name into an id.
func slug(name string) string {
	var b strings.Builder
	for _, r := range strings.ToLower(strings.TrimSpace(name)) {
		switch {
		case r >= 'a' && r <= 'z', r >= '0' && r <= '9':
			b.WriteRune(r)
		case r == ' ' || r == '-' || r == '_':
			b.WriteByte('-')
		}
	}
	if b.Len() == 0 {
		return "default"
	}
	return b.String()
}

func randomID(prefix string) string {
	var buf [10]byte
	_, _ = rand.Read(buf[:])
	return prefix + "-" + strings.ToLower(base32.StdEncoding.WithPadding(base32.NoPadding).EncodeToString(buf[:]))
}

func newKeyRotateCmd(open opener) *cobra.Command {
	var overlap time.Duration
	cmd := &cobra.Command{
		Use:   "rotate <key-id>",
		Short: "Replace a key with a new one, the old one working on for a while",
		Long: "The new key belongs where the old one did — the same project, owner, band\n" +
			"and routing profile, and as long a life — and is printed once. The old\n" +
			"one stops working after --overlap, so whatever uses it can be moved over\n" +
			"without an outage.",
		Example: "  llmgw keys rotate key-7f3a --overlap 48h",
		Args:    cobra.ExactArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()
			old, err := s.KeyByID(cmd.Context(), args[0])
			if err != nil {
				return fmt.Errorf("no key %s: %w", args[0], err)
			}
			if !old.RevokedAt.IsZero() {
				return fmt.Errorf("%s is revoked; make a new key instead", old.ID)
			}
			secret, k, err := store.NewKey(randomID("key"), old.Name, old.ProjectID)
			if err != nil {
				return err
			}
			k.CreatedAt = time.Now()
			k.Band, k.RoutingProfile, k.UserID = old.Band, old.RoutingProfile, old.UserID
			if !old.ExpiresAt.IsZero() {
				k.ExpiresAt = k.CreatedAt.Add(old.ExpiresAt.Sub(old.CreatedAt))
			}
			if err := s.CreateKey(cmd.Context(), k); err != nil {
				return err
			}
			retire := time.Now().Add(overlap)
			if err := s.ExpireKeyBy(cmd.Context(), old.ID, retire); err != nil {
				return err
			}
			fmt.Fprintln(cmd.OutOrStdout(), secret)
			cmd.PrintErrf("\nid       %s (replaces %s)\nold key  stops working at %s\n", k.ID, old.ID,
				retire.UTC().Format(time.RFC3339))
			return nil
		},
	}
	cmd.Flags().DurationVar(&overlap, "overlap", 24*time.Hour, "how long the old key keeps working")
	return cmd
}

func newKeyTTLCmd(open opener) *cobra.Command {
	return &cobra.Command{
		Use:   "ttl [<org|team|project>:<name> <duration|off>]",
		Short: "How long a key may live at most, in a part of the tree",
		Long: "A key is made within the shortest limit that applies to it — its\n" +
			"project's, its team's, its organization's, and auth.max_key_ttl in\n" +
			"server.yaml for every key. A key asked for with no expiry is given the\n" +
			"longest allowed. Existing keys are not changed. Without arguments, lists\n" +
			"the limits.",
		Example: "  llmgw keys ttl team:Acme/ml 720h\n  llmgw keys ttl team:Acme/ml off",
		Args:    cobra.RangeArgs(0, 2),
		RunE: func(cmd *cobra.Command, args []string) error {
			s, err := open(cmd.Context())
			if err != nil {
				return err
			}
			defer s.Close()
			if len(args) == 0 {
				limits, err := s.KeyTTLs(cmd.Context())
				if err != nil {
					return err
				}
				if len(limits) == 0 {
					fmt.Fprintln(cmd.OutOrStdout(), "No limit: keys may live for ever.")
					return nil
				}
				names := make([]string, 0, len(limits))
				for k := range limits {
					names = append(names, k)
				}
				slices.Sort(names)
				w := tabwriter.NewWriter(cmd.OutOrStdout(), 0, 0, 2, ' ', 0)
				fmt.Fprintln(w, "SCOPE\tLONGEST")
				for _, n := range names {
					fmt.Fprintf(w, "%s\t%s\n", strings.TrimSuffix(n, ":"), limits[n])
				}
				return w.Flush()
			}
			if len(args) != 2 {
				return errors.New("give a scope and a duration, or off")
			}
			kind, name, ok := strings.Cut(args[0], ":")
			if !ok {
				return fmt.Errorf("%q is not org:NAME, team:NAME or project:NAME", args[0])
			}
			var id string
			switch store.Scope(kind) {
			case store.ScopeOrg:
				o, err := resolveOrg(cmd.Context(), s, name)
				if err != nil {
					return err
				}
				id = o.ID
			case store.ScopeTeam:
				t, err := resolveTeam(cmd.Context(), s, name)
				if err != nil {
					return err
				}
				id = t.ID
			case store.ScopeProject:
				p, err := resolveProject(cmd.Context(), s, name)
				if err != nil {
					return err
				}
				id = p.ID
			default:
				return fmt.Errorf("%q: org, team or project — the gateway's is auth.max_key_ttl in server.yaml", kind)
			}
			var d time.Duration
			if args[1] != "off" {
				if d, err = time.ParseDuration(args[1]); err != nil || d <= 0 {
					return fmt.Errorf("%q is not a duration like 720h, or off", args[1])
				}
			}
			if err := s.SetMaxKeyTTL(cmd.Context(), store.Scope(kind), id, d); err != nil {
				return err
			}
			if d == 0 {
				fmt.Fprintf(cmd.OutOrStdout(), "%s: no limit of its own\n", args[0])
			} else {
				fmt.Fprintf(cmd.OutOrStdout(), "%s: keys live %s at most\n", args[0], d)
			}
			return nil
		},
	}
}