terraform-playground/pymage/internal/cli/cli.go

// Package cli implements the pymage command line.
package cli

import (
	"context"
	"fmt"
	"os"
	"runtime"
	"strconv"
	"strings"

	"github.com/google/go-containerregistry/pkg/authn"
	"github.com/google/go-containerregistry/pkg/logs"
	"github.com/google/go-containerregistry/pkg/name"
	v1 "github.com/google/go-containerregistry/pkg/v1"
	"github.com/google/go-containerregistry/pkg/v1/empty"
	"github.com/google/go-containerregistry/pkg/v1/layout"
	"github.com/google/go-containerregistry/pkg/v1/mutate"
	"github.com/google/go-containerregistry/pkg/v1/remote"
	"github.com/spf13/cobra"

	"github.com/imjasonh/terraform-playground/pymage/internal/build"
	"github.com/imjasonh/terraform-playground/pymage/internal/cache"
	"github.com/imjasonh/terraform-playground/pymage/internal/lock"
	"github.com/imjasonh/terraform-playground/pymage/internal/sbom"
	"github.com/imjasonh/terraform-playground/pymage/internal/wheel"
	"github.com/imjasonh/terraform-playground/pymage/internal/wheelhouse"
)

// Root returns the root cobra command.
func Root() *cobra.Command {
	root := &cobra.Command{
		Use:           "pymage",
		Short:         "Docker-less, layer-aware Python container image builder",
		SilenceUsage:  true,
		SilenceErrors: true,
	}
	root.AddCommand(buildCmd())
	return root
}

type buildFlags struct {
	base       string
	lockFile   string
	findLinks  []string
	source     string
	repo       string
	tags       []string
	platforms  []string
	extras     []string
	pkg        string
	pythonTag  string
	prefix     string
	workingDir string
	entrypoint []string
	cmd        []string
	env        []string
	user       string
	labels     []string
	ignore     []string
	strategy   string

	maxLayers      int
	maxWheelLayers int
	pushJobs       int

	push        bool
	ociLayout   string
	printDigest bool
	sbomOut     string
	requireHash bool
	insecure    bool
	cacheDir    string
	noCache     bool
}

func buildCmd() *cobra.Command {
	f := &buildFlags{}
	cmd := &cobra.Command{
		Use:   "build [source]",
		Short: "Build a Python application image and push it",
		Long: "Build a Python application image and push it.\n\n" +
			"[source] is the uv project directory to build (default: current directory).",
		Args: cobra.MaximumNArgs(1),
		RunE: func(cmd *cobra.Command, args []string) error {
			if len(args) == 1 {
				f.source = args[0]
			}
			return runBuild(cmd, f)
		},
	}
	fs := cmd.Flags()
	fs.StringVar(&f.base, "base", "", "base image reference (default: cgr.dev/chainguard/python:latest)")
	fs.StringVar(&f.lockFile, "lock", "", "lock file (default: uv.lock in the source directory)")
	fs.StringArrayVar(&f.extras, "extra", nil, "enable a project optional-dependency group from uv.lock (repeatable)")
	fs.StringVar(&f.pkg, "package", "", "build a single uv workspace member (by name)")
	fs.StringArrayVar(&f.findLinks, "find-links", nil, "local wheel directory (optional; wheels are downloaded from the lock when omitted)")
	fs.StringVar(&f.repo, "repo", "", "destination repository, pushed by digest, e.g. gcr.io/foo/bar (default: [tool.pymage] repo)")
	fs.StringArrayVarP(&f.tags, "tag", "t", nil, "tag(s) to apply at --repo; tag component only, not a full reference (repeatable; default: [tool.pymage] tags or 'latest')")
	fs.StringSliceVar(&f.platforms, "platform", nil, "target platform(s); repeatable or comma-separated, e.g. linux/amd64,linux/arm64 (multiple builds a multi-arch index)")
	fs.StringVar(&f.pythonTag, "python", "", "interpreter version, e.g. python3.12 (default: auto-detected from the base image; if set, must match the base)")
	fs.StringVar(&f.prefix, "prefix", "/app/.venv", "install prefix (venv-like root)")
	fs.StringVar(&f.workingDir, "workdir", "/app", "image working directory and source destination")
	fs.StringArrayVar(&f.entrypoint, "entrypoint", nil, "entrypoint argv (repeatable)")
	fs.StringArrayVar(&f.cmd, "cmd", nil, "cmd argv (repeatable)")
	fs.StringArrayVar(&f.env, "env", nil, "extra KEY=VALUE env (repeatable)")
	fs.StringVar(&f.user, "user", "", "image user, e.g. 65532")
	fs.StringArrayVar(&f.labels, "label", nil, "image label KEY=VALUE (repeatable)")
	fs.StringArrayVar(&f.ignore, "ignore", nil, "extra source ignore glob (repeatable)")
	fs.StringVar(&f.strategy, "layer-strategy", string(build.Auto), "auto | per-wheel | single-deps-layer")
	fs.IntVar(&f.maxLayers, "max-layers", build.DefaultMaxLayers, "max total image layers (base + deps + app) for the auto strategy")
	fs.IntVar(&f.maxWheelLayers, "max-wheel-layers", 0, "cap the number of dependency layers directly (overrides --max-layers when > 0)")
	fs.IntVar(&f.pushJobs, "push-concurrency", 0, "max concurrent layer uploads when pushing (0 = auto)")
	fs.BoolVar(&f.push, "push", true, "push the image to --repo (by digest)")
	fs.StringVar(&f.ociLayout, "oci-layout", "", "also write the image to this OCI layout directory")
	fs.BoolVar(&f.printDigest, "print-digest", false, "print only the resulting image digest")
	fs.StringVar(&f.sbomOut, "sbom", "", "write a CycloneDX SBOM to this path")
	fs.BoolVar(&f.requireHash, "require-hashes", true, "require every requirement to carry --hash entries")
	fs.BoolVar(&f.insecure, "insecure", false, "use plain HTTP for the registry")
	fs.StringVar(&f.cacheDir, "cache-dir", "", "cache root directory (default: per-user cache dir)")
	fs.BoolVar(&f.noCache, "no-cache", false, "disable all build caches (layers, downloaded wheels, interpreter detection)")
	return cmd
}

func runBuild(cmd *cobra.Command, f *buildFlags) error {
	ctx := cmd.Context()
	// Keep stdout clean (only the final image ref): route go-containerregistry's
	// per-blob progress ("existing blob", "mounted blob", "pushed blob") and
	// warnings to stderr.
	logs.Progress.SetOutput(cmd.ErrOrStderr())
	logs.Warn.SetOutput(cmd.ErrOrStderr())

	if err := applyDefaults(cmd, f); err != nil {
		return err
	}
	if err := validateBuildFlags(f); err != nil {
		return err
	}

	if _, err := keyValues(f.env); err != nil {
		return fmt.Errorf("--env: %w", err)
	}
	labels, err := keyValues(f.labels)
	if err != nil {
		return fmt.Errorf("--label: %w", err)
	}

	lk, err := lock.Load(f.lockFile)
	if err != nil {
		return err
	}

	platforms, err := parsePlatforms(f.platforms)
	if err != nil {
		return err
	}

	var nameOpts []name.Option
	if f.insecure {
		nameOpts = append(nameOpts, name.Insecure)
	}

	// Resolve the base once (single index/manifest fetch) and reuse it for every
	// platform, rather than re-fetching it per platform.
	baseSet, err := build.ResolveBaseSet(ctx, f.base, authn.DefaultKeychain, nameOpts...)
	if err != nil {
		return err
	}

	// With no platform requested, default to whatever the base image supports
	// (a multi-arch base yields a multi-arch build).
	if len(platforms) == 0 {
		platforms = baseSet.Platforms()
	}

	// Caching is on by default (layers, downloaded wheels, and interpreter
	// detection); --no-cache disables it.
	layerCache, metaCache, wheelCache, cleanup, err := setupCaches(f)
	if err != nil {
		return err
	}
	defer cleanup()

	// Build one image per target platform. With more than one platform we
	// assemble them into a multi-arch OCI image index; with zero or one we push
	// a plain image manifest.
	var (
		images  []v1.Image
		adds    []mutate.IndexAddendum
		allDeps []wheelhouse.ResolvedWheel
	)
	targets := platforms
	for i := range targets {
		p := targets[i]
		var pp *v1.Platform
		if p.OS != "" || p.Architecture != "" {
			pp = &targets[i]
		}
		base, err := baseSet.Image(pp)
		if err != nil {
			return err
		}
		img, deps, err := buildOne(ctx, f, lk, pp, base, labels, layerCache, metaCache, wheelCache)
		if err != nil {
			return err
		}
		images = append(images, img)
		allDeps = append(allDeps, deps...)
		desc := v1.Descriptor{}
		if pp != nil {
			desc.Platform = pp
		}
		adds = append(adds, mutate.IndexAddendum{Add: img, Descriptor: desc})
	}

	if f.sbomOut != "" {
		data, err := sbom.Generate(dedupeWheels(allDeps))
		if err != nil {
			return err
		}
		if err := os.WriteFile(f.sbomOut, data, 0o644); err != nil {
			return err
		}
	}

	if len(platforms) > 1 {
		idx := mutate.AppendManifests(empty.Index, adds...)
		return output(cmd, f, nameOpts, nil, idx)
	}
	return output(cmd, f, nameOpts, images[0], nil)
}

// buildOne resolves wheels for a single platform and builds the image from the
// already-resolved base, returning the resolved wheels for SBOM aggregation.
func buildOne(ctx context.Context, f *buildFlags, lk *lock.Lock, platform *v1.Platform, base v1.Image, labels map[string]string, layerCache, metaCache *cache.Cache, wheelCache string) (v1.Image, []wheelhouse.ResolvedWheel, error) {
	// Determine the interpreter: honor --python (validated against the base) or
	// auto-detect it from the base image.
	major, minor, pyTag, err := resolveInterpreter(f, base, metaCache)
	if err != nil {
		return nil, nil, err
	}

	target := platformTarget(platform)
	target.PyMajor, target.PyMinor = major, minor

	// Resolve the lock for this target (uv.lock runtime closure honors
	// --package, --extra, and the target's environment markers).
	reqs, err := lk.Resolve(lock.Options{
		Package: f.pkg, Extras: f.extras,
		OS: target.OS, Arch: target.Arch, PyMajor: target.PyMajor, PyMinor: target.PyMinor,
	})
	if err != nil {
		return nil, nil, err
	}
	if err := lock.Validate(reqs, f.requireHash); err != nil {
		return nil, nil, err
	}

	wheels, err := wheelhouse.ResolveContext(ctx, reqs, f.findLinks, target, wheelCache)
	if err != nil {
		return nil, nil, err
	}
	img, err := build.Build(build.Options{
		Base:           base,
		Wheels:         wheels,
		Layout:         wheel.Layout{Prefix: f.prefix, PythonTag: pyTag},
		Strategy:       build.LayerStrategy(f.strategy),
		MaxLayers:      f.maxLayers,
		MaxWheelLayers: f.maxWheelLayers,
		SourceDir:      f.source,
		Ignore:         f.ignore,
		WorkingDir:     f.workingDir,
		Entrypoint:     f.entrypoint,
		Cmd:            f.cmd,
		Env:            f.env,
		User:           f.user,
		Labels:         labels,
		Cache:          layerCache,
	})
	if err != nil {
		return nil, nil, err
	}
	return img, wheels, nil
}

// resolveInterpreter decides the target Python version and the site-packages
// directory tag. If --python is set it is authoritative but must match the
// base when the base advertises a version; otherwise the version is detected
// from the base image.
func resolveInterpreter(f *buildFlags, base v1.Image, metaCache *cache.Cache) (major, minor int, pyTag string, err error) {
	baseMaj, baseMin, baseOK := cachedInterpreter(base, metaCache)

	if f.pythonTag != "" {
		maj, min, err := parsePythonTag(f.pythonTag)
		if err != nil {
			return 0, 0, "", err
		}
		if baseOK && (baseMaj != maj || baseMin != min) {
			return 0, 0, "", fmt.Errorf("base %q provides Python %d.%d but --python is python%d.%d; pass --python python%d.%d (or pin a matching base)",
				f.base, baseMaj, baseMin, maj, min, baseMaj, baseMin)
		}
		return maj, min, f.pythonTag, nil
	}

	if !baseOK {
		return 0, 0, "", fmt.Errorf("could not determine the base image's Python version (no PYTHON_VERSION env or python package in /etc/apko.json); pass --python pythonX.Y")
	}
	return baseMaj, baseMin, fmt.Sprintf("python%d.%d", baseMaj, baseMin), nil
}

// output writes/pushes either a single image or a multi-arch index (exactly one
// of img/idx is non-nil).
//
// stdout carries exactly one line — the resulting image reference — so callers
// can run `docker run "$(pymage build)"`. All progress/diagnostic output goes
// to stderr.
func output(cmd *cobra.Command, f *buildFlags, nameOpts []name.Option, img v1.Image, idx v1.ImageIndex) error {
	ctx := cmd.Context()
	stdout := cmd.OutOrStdout()
	stderr := cmd.ErrOrStderr()

	digest, err := artifactDigest(img, idx)
	if err != nil {
		return err
	}

	if f.ociLayout != "" {
		if err := writeLayout(f.ociLayout, img, idx); err != nil {
			return err
		}
		_, _ = fmt.Fprintf(stderr, "wrote OCI layout to %s\n", f.ociLayout)
	}

	if f.printDigest {
		_, _ = fmt.Fprintln(stdout, digest.String())
		return nil
	}

	if !f.push {
		_, _ = fmt.Fprintln(stdout, digest.String())
		return nil
	}

	if f.repo == "" {
		return fmt.Errorf("no repo configured: set repo in [tool.pymage] or pass --repo (or use --push=false)")
	}
	repo, err := name.NewRepository(f.repo, nameOpts...)
	if err != nil {
		return err
	}
	tags := f.tags
	if len(tags) == 0 {
		tags = []string{"latest"}
	}

	var artifact remote.Taggable = img
	if idx != nil {
		artifact = idx
	}

	// A single Pusher is shared across all tags: it uploads layers (and, for an
	// index, child manifests) concurrently, and reuses its per-repo upload state
	// so identical blobs are only checked/uploaded once even across tags.
	pusher, err := remote.NewPusher(
		remote.WithAuthFromKeychain(authn.DefaultKeychain),
		remote.WithJobs(pushConcurrency(f.pushJobs)),
	)
	if err != nil {
		return err
	}
	// The artifact is pushed by content; each tag is just another pointer to
	// the same digest.
	for _, t := range tags {
		ref := repo.Tag(t)
		if err := pusher.Push(ctx, ref, artifact); err != nil {
			return fmt.Errorf("push %s: %w", ref, err)
		}
		_, _ = fmt.Fprintf(stderr, "tagged %s -> %s@%s\n", ref, repo.Name(), digest)
	}

	// The sole stdout line: the immutable by-digest reference.
	_, _ = fmt.Fprintf(stdout, "%s@%s\n", repo.Name(), digest)
	return nil
}

func artifactDigest(img v1.Image, idx v1.ImageIndex) (v1.Hash, error) {
	if idx != nil {
		return idx.Digest()
	}
	return img.Digest()
}

func writeLayout(dir string, img v1.Image, idx v1.ImageIndex) error {
	p, err := layout.Write(dir, empty.Index)
	if err != nil {
		return err
	}
	if idx != nil {
		return p.AppendIndex(idx)
	}
	return p.AppendImage(img)
}

// cachedInterpreter returns the base image's Python version, consulting the
// metadata cache (keyed by base digest) so the apko.json layer isn't
// re-downloaded on repeat builds.
func cachedInterpreter(base v1.Image, metaCache *cache.Cache) (major, minor int, ok bool) {
	if metaCache != nil {
		if d, err := base.Digest(); err == nil {
			key := "interp:v1:" + d.String()
			if v, hit := metaCache.GetText(key); hit {
				if maj, min, parsed := splitMajorMinor(v); parsed {
					return maj, min, true
				}
			}
			maj, min, found := build.InterpreterVersion(base)
			if found {
				_ = metaCache.PutText(key, fmt.Sprintf("%d.%d", maj, min))
			}
			return maj, min, found
		}
	}
	return build.InterpreterVersion(base)
}

func splitMajorMinor(s string) (int, int, bool) {
	a, b, ok := strings.Cut(strings.TrimSpace(s), ".")
	if !ok {
		return 0, 0, false
	}
	maj, e1 := strconv.Atoi(a)
	min, e2 := strconv.Atoi(b)
	if e1 != nil || e2 != nil {
		return 0, 0, false
	}
	return maj, min, true
}

// pushConcurrency resolves the number of concurrent layer uploads. 0 means
// auto: scale with CPU count but never below go-containerregistry's default of
// 4 (layer uploads are network-bound, so some parallelism always helps).
func pushConcurrency(requested int) int {
	if requested > 0 {
		return requested
	}
	n := runtime.NumCPU()
	if n < 4 {
		n = 4
	}
	return n
}

// parsePlatforms parses the (comma-split, repeatable) --platform values.
func parsePlatforms(specs []string) ([]v1.Platform, error) {
	var out []v1.Platform
	for _, s := range specs {
		s = strings.TrimSpace(s)
		if s == "" {
			continue
		}
		p, err := v1.ParsePlatform(s)
		if err != nil {
			return nil, err
		}
		out = append(out, *p)
	}
	return out, nil
}

// dedupeWheels removes duplicate wheels (same name+version+sha) so a multi-arch
// SBOM doesn't list shared pure-python wheels multiple times.
func dedupeWheels(in []wheelhouse.ResolvedWheel) []wheelhouse.ResolvedWheel {
	seen := map[string]bool{}
	var out []wheelhouse.ResolvedWheel
	for _, w := range in {
		k := w.Name + "\x00" + w.Version + "\x00" + w.SHA256
		if seen[k] {
			continue
		}
		seen[k] = true
		out = append(out, w)
	}
	return out
}

// platformTarget derives the OS/arch part of the wheel-selection target from
// the --platform value, defaulting to linux/amd64. The interpreter fields are
// filled in by the caller.
func platformTarget(platform *v1.Platform) wheel.Target {
	t := wheel.Target{OS: "linux", Arch: "amd64"}
	if platform != nil {
		if platform.OS != "" {
			t.OS = platform.OS
		}
		if platform.Architecture != "" {
			t.Arch = platform.Architecture
		}
	}
	return t
}

// parsePythonTag parses "python3.12" into (3, 12). The tag must start with
// "python" because it is also used verbatim as the site-packages lib directory
// (".../lib/python3.12/site-packages").
func parsePythonTag(tag string) (int, int, error) {
	s, ok := strings.CutPrefix(tag, "python")
	if !ok {
		return 0, 0, fmt.Errorf("--python %q must look like pythonX.Y", tag)
	}
	majStr, minStr, ok := strings.Cut(s, ".")
	if !ok {
		return 0, 0, fmt.Errorf("--python %q must look like pythonX.Y", tag)
	}
	maj, err := strconv.Atoi(majStr)
	if err != nil {
		return 0, 0, fmt.Errorf("--python %q: bad major version", tag)
	}
	min, err := strconv.Atoi(minStr)
	if err != nil {
		return 0, 0, fmt.Errorf("--python %q: bad minor version", tag)
	}
	return maj, min, nil
}

func keyValues(in []string) (map[string]string, error) {
	if len(in) == 0 {
		return nil, nil
	}
	out := map[string]string{}
	for _, kv := range in {
		k, v, ok := strings.Cut(kv, "=")
		if !ok || k == "" {
			return nil, fmt.Errorf("expected KEY=VALUE with a non-empty key, got %q", kv)
		}
		out[k] = v
	}
	return out, nil
}