ENG-1793: Install decrypted private-ingredient wheels at consume time

internal/python/wheelinstall/wheelinstall.go

@@ -0,0 +1,61 @@
+// Package wheelinstall installs a pure-Python wheel into a site-packages
+// directory.
+package wheelinstall
+
+import (
+	"os"
+	"path/filepath"
+	"strings"
+
+	"github.com/ActiveState/cli/internal/errs"
+	"github.com/ActiveState/cli/internal/fileutils"
+	"github.com/ActiveState/cli/internal/unarchiver"
+)
+
+// installerName is the contents of the .dist-info/INSTALLER file (PEP 376).
+const installerName = "state-tool"
+
+// Install extracts the wheel at wheelPath into sitePackagesDir and writes the
+// INSTALLER marker into its .dist-info. Entries that would escape sitePackagesDir,
+// and wheels without a .dist-info, are rejected.
+func Install(wheelPath, sitePackagesDir string) error {
+	if err := fileutils.MkdirUnlessExists(sitePackagesDir); err != nil {
+		return errs.Wrap(err, "could not create site-packages directory")
+	}
+
+	wheel, err := os.Open(wheelPath)
+	if err != nil {
+		return errs.Wrap(err, "could not open wheel")
+	}
+	defer wheel.Close()
+
+	// A wheel is a zip; untrusted-source mode confines entries to the destination.
+	ua := unarchiver.NewZip(unarchiver.WithUntrustedSource())
+	if err := ua.Unarchive(wheel, sitePackagesDir); err != nil {
+		return errs.Wrap(err, "could not extract wheel")
+	}
+
+	if err := writeInstaller(sitePackagesDir); err != nil {
+		return errs.Wrap(err, "could not record installer")
+	}
+	return nil
+}
+
+// writeInstaller writes the INSTALLER marker into the wheel's sole *.dist-info
+// directory, erroring if there is none.
+func writeInstaller(sitePackagesDir string) error {
+	entries, err := os.ReadDir(sitePackagesDir)
+	if err != nil {
+		return errs.Wrap(err, "could not read site-packages directory")
+	}
+	for _, e := range entries {
+		if e.IsDir() && strings.HasSuffix(e.Name(), ".dist-info") {
+			marker := filepath.Join(sitePackagesDir, e.Name(), "INSTALLER")
+			if err := fileutils.WriteFile(marker, []byte(installerName+"\n")); err != nil {
+				return errs.Wrap(err, "could not write INSTALLER")
+			}
+			return nil
+		}
+	}
+	return errs.New("wheel has no .dist-info directory")
+}

internal/python/wheelinstall/wheelinstall_test.go

@@ -0,0 +1,69 @@
+package wheelinstall
+
+import (
+	"archive/zip"
+	"os"
+	"path/filepath"
+	"testing"
+)
+
+// makeWheel writes a zip (wheel) containing the given name->body entries and
+// returns its path.
+func makeWheel(t *testing.T, dir string, entries map[string]string) string {
+	t.Helper()
+	wheelPath := filepath.Join(dir, "greeting-1.0-py3-none-any.whl")
+	f, err := os.Create(wheelPath)
+	if err != nil {
+		t.Fatal(err)
+	}
+	defer f.Close()
+	zw := zip.NewWriter(f)
+	for name, body := range entries {
+		w, err := zw.Create(name)
+		if err != nil {
+			t.Fatal(err)
+		}
+		if _, err := w.Write([]byte(body)); err != nil {
+			t.Fatal(err)
+		}
+	}
+	if err := zw.Close(); err != nil {
+		t.Fatal(err)
+	}
+	return wheelPath
+}
+
+func TestInstall(t *testing.T) {
+	t.Run("extracts the package and records the installer", func(t *testing.T) {
+		dir := t.TempDir()
+		wheel := makeWheel(t, dir, map[string]string{
+			"greeting/__init__.py":            "print('hi')\n",
+			"greeting-1.0.dist-info/METADATA": "Name: greeting\nVersion: 1.0\n",
+			"greeting-1.0.dist-info/RECORD":   "",
+		})
+
+		site := filepath.Join(dir, "site-packages")
+		if err := Install(wheel, site); err != nil {
+			t.Fatalf("Install: %v", err)
+		}
+
+		if got, _ := os.ReadFile(filepath.Join(site, "greeting", "__init__.py")); string(got) != "print('hi')\n" {
+			t.Errorf("package not extracted: got %q", got)
+		}
+		installer, err := os.ReadFile(filepath.Join(site, "greeting-1.0.dist-info", "INSTALLER"))
+		if err != nil {
+			t.Fatalf("INSTALLER not written: %v", err)
+		}
+		if string(installer) != installerName+"\n" {
+			t.Errorf("INSTALLER = %q, want %q", installer, installerName+"\n")
+		}
+	})
+
+	t.Run("a wheel without a .dist-info fails closed", func(t *testing.T) {
+		dir := t.TempDir()
+		wheel := makeWheel(t, dir, map[string]string{"greeting/__init__.py": "x\n"})
+		if err := Install(wheel, filepath.Join(dir, "site-packages")); err == nil {
+			t.Error("expected an error for a wheel without a .dist-info directory")
+		}
+	})
+}

pkg/runtime/setup.go

@@ -19,8 +19,10 @@ import (
 	"github.com/ActiveState/cli/internal/httputil"
 	"github.com/ActiveState/cli/internal/locale"
 	"github.com/ActiveState/cli/internal/logging"
+	"github.com/ActiveState/cli/internal/multilog"
 	"github.com/ActiveState/cli/internal/osutils"
 	"github.com/ActiveState/cli/internal/proxyreader"
+	"github.com/ActiveState/cli/internal/python/wheelinstall"
 	"github.com/ActiveState/cli/internal/sliceutils"
 	"github.com/ActiveState/cli/internal/svcctl"
 	"github.com/ActiveState/cli/internal/unarchiver"
@@ -478,6 +480,18 @@ func (s *setup) unpack(artifact *buildplan.Artifact, b []byte) (rerr error) {
 		if err := s.depot.MarkPrivate(artifact.ArtifactID); err != nil {
 			return errs.Wrap(err, "Could not mark decrypted artifact as private")
 		}
+		switch {
+		case s.isPrivateWheel(unpackPath):
+			if err := s.installPrivateWheel(unpackPath); err != nil {
+				rerr := errs.Wrap(err, "Could not install private wheel")
+				if err2 := os.RemoveAll(unpackPath); err2 != nil {
+					return errs.Pack(rerr, errs.Wrap(err2, "unable to remove artifact directory"))
+				}
+				return rerr
+			}
+		default:
+			multilog.Error("Decrypted private artifact %s (%s) is of an unknown type; cannot install it", artifact.ArtifactID, artifact.Name())
+		}
 	}
 
 	// Camel artifacts do not have runtime.json, so in order to not have multiple paths of logic we generate one based
@@ -638,6 +652,80 @@ func readPayloadHeader(path string) (header artifactcrypto.Header, rerr error) {
 	return artifactcrypto.ParseHeader(f)
 }
 
+// isPrivateWheel reports whether the decrypted payload under dir is a wheel. We
+// control the payload via `state publish --build`, so a .whl extension is a
+// sufficient test.
+func (s *setup) isPrivateWheel(dir string) bool {
+	wheelPath, err := findWheel(dir)
+	return err == nil && wheelPath != ""
+}
+
+// installPrivateWheel installs the decrypted wheel found under artifactDir into a
+// site-packages directory and adds it to PYTHONPATH in the artifact's
+// runtime.json.
+func (s *setup) installPrivateWheel(artifactDir string) error {
+	wheelPath, err := findWheel(artifactDir)
+	if err != nil {
+		return errs.Wrap(err, "could not locate decrypted wheel")
+	}
+	if wheelPath == "" {
+		return errs.New("decrypted private artifact contains no wheel")
+	}
+
+	// site-packages sits in the deploy tree, so the deploy links it into the
+	// runtime where ${INSTALLDIR}/site-packages resolves it.
+	sitePackages := filepath.Join(filepath.Dir(wheelPath), "site-packages")
+	if err := wheelinstall.Install(wheelPath, sitePackages); err != nil {
+		return errs.Wrap(err, "could not install wheel")
+	}
+	if err := os.Remove(wheelPath); err != nil {
+		return errs.Wrap(err, "could not remove installed wheel")
+	}
+
+	return s.exposeSitePackages(artifactDir)
+}
+
+// exposeSitePackages adds the installed site-packages directory to PYTHONPATH in
+// the artifact's runtime.json.
+func (s *setup) exposeSitePackages(artifactDir string) error {
+	rtPath := filepath.Join(artifactDir, envdef.EnvironmentDefinitionFilename)
+	envDef, err := envdef.NewEnvironmentDefinition(rtPath)
+	if err != nil {
+		return errs.Wrap(err, "could not load runtime definition")
+	}
+	envDef.Env = append(envDef.Env, envdef.EnvironmentVariable{
+		Name:      "PYTHONPATH",
+		Values:    []string{"${INSTALLDIR}/site-packages"},
+		Join:      envdef.Prepend,
+		Inherit:   false,
+		Separator: ":", // OS-independent
+	})
+	if err := envDef.Save(artifactDir); err != nil {
+		return errs.Wrap(err, "could not save runtime definition")
+	}
+	return nil
+}
+
+// findWheel returns the path of the single .whl under dir (searched recursively),
+// or "" if none is present.
+func findWheel(dir string) (string, error) {
+	var found string
+	err := filepath.WalkDir(dir, func(path string, d os.DirEntry, err error) error {
+		if err != nil {
+			return err
+		}
+		if !d.IsDir() && strings.HasSuffix(d.Name(), ".whl") {
+			found = path
+			return filepath.SkipAll
+		}
+		return nil
+	})
+	if err != nil {
+		return "", errs.Wrap(err, "could not scan for wheel")
+	}
+	return found, nil
+}
+
 func (s *setup) updateExecutors() error {
 	execPath := ExecutorsPath(s.path)
 	if err := fileutils.MkdirUnlessExists(execPath); err != nil {

test/integration/publish_int_test.go

@@ -1,15 +1,12 @@
 package integration
 
 import (
-	"archive/zip"
 	"encoding/base64"
 	"encoding/json"
 	"fmt"
-	"io"
 	"os"
 	"path/filepath"
 	"regexp"
-	"strings"
 	"testing"
 	"time"
 
@@ -556,85 +553,14 @@ func (suite *PublishIntegrationTestSuite) TestPublishBuildEncrypted() {
 	cp.Expect("All dependencies have been installed and verified", e2e.RuntimeBuildSourcingTimeoutOpt)
 	cp.ExpectExitCode(0)
 
-	// Decryption proof: the decrypted content must be present in the depot and
-	// contain our sentinel. A failed decrypt would skip the artifact, leaving the
-	// sentinel absent.
-	suite.assertDecryptedPayloadContains(ts, sentinel)
-}
-
-// assertDecryptedPayloadContains fails the test unless a decrypted artifact under
-// the depot contains sentinel. It scans every wheel (as a zip) and every small
-// plaintext file, since the exact on-disk path depends on how the artifact is
-// packaged on install.
-func (suite *PublishIntegrationTestSuite) assertDecryptedPayloadContains(ts *e2e.Session, sentinel string) {
-	depot := filepath.Join(ts.Dirs.Cache, "depot")
-
-	var wheels []string
-	fileCount := 0
-	found := false
-	walkErr := filepath.WalkDir(depot, func(path string, d os.DirEntry, err error) error {
-		if err != nil {
-			return err
-		}
-		if found {
-			return filepath.SkipAll // sentinel located; no need to walk the rest of the depot
-		}
-		if d.IsDir() {
-			return nil
-		}
-		fileCount++
-
-		// A decrypted wheel is a zip; scan its entries for the sentinel.
-		if strings.HasSuffix(d.Name(), ".whl") {
-			wheels = append(wheels, path)
-			if suite.wheelContains(path, sentinel) {
-				found = true
-			}
-			return nil
-		}
-
-		// Otherwise scan the raw file, in case the payload was delivered unpacked
-		// rather than as a wheel. Skip large files (the sentinel lives in a tiny
-		// Python source file).
-		if info, err := d.Info(); err != nil || info.Size() > 5<<20 {
-			return nil
-		}
-		content, err := os.ReadFile(path)
-		if err == nil && strings.Contains(string(content), sentinel) {
-			found = true
-		}
-		return nil
-	})
-	suite.Require().NoError(walkErr, "could not walk depot %s", depot)
-	suite.T().Logf("searched %d files under the depot; wheels found: %v", fileCount, wheels)
-	suite.Require().True(found, "sentinel %q not found in the depot; the artifact was likely not decrypted", sentinel)
-}
-
-// wheelContains reports whether any file inside the wheel (a zip) contains
-// sentinel. A wheel that failed to decrypt would not be a readable zip, so an
-// unreadable wheel is logged and treated as not containing the sentinel.
-func (suite *PublishIntegrationTestSuite) wheelContains(wheelPath, sentinel string) bool {
-	zr, err := zip.OpenReader(wheelPath)
-	if err != nil {
-		suite.T().Logf("could not open wheel %s as zip: %v", wheelPath, err)
-		return false
-	}
-	defer zr.Close()
-	for _, f := range zr.File {
-		rc, err := f.Open()
-		if err != nil {
-			continue
-		}
-		content, err := io.ReadAll(rc)
-		rc.Close()
-		if err != nil {
-			continue
-		}
-		if strings.Contains(string(content), sentinel) {
-			return true
-		}
-	}
-	return false
+	// Installation proof: the decrypted wheel is installed into the runtime's
+	// site-packages and is importable. Importing the package runs its
+	// __init__.py, which prints the unique sentinel — a value only the decrypted
+	// plaintext contains, so this confirms decrypt + install + PYTHONPATH wiring
+	// end to end.
+	cp = ts.SpawnWithOpts(e2e.OptArgs("exec", "python3", "--", "-c", "import greeting"))
+	cp.Expect(sentinel)
+	cp.ExpectExitCode(0)
 }
 
 // orgKeyContract builds the org-key contract JSON the key service would serve for