sandbox/pythonRuntime.go

package sandbox

import (
	"errors"
	"fmt"
	"net/http"
	"net/url"
	"os"
	"path/filepath"
	"regexp"
	"runtime"
	"time"

	"github.com/ssgo/httpclient"
	"github.com/ssgo/log"
	"github.com/ssgo/u"
)

func init() {
	RegisterRuntime("python", &Runtime{
		Check: func(runtimePath, venvPath, projectPath string, uid, gid int, cfg *RuntimeConfig) (*RuntimeSandboxConfig, error) {
			if err := checkPythonRuntime(runtimePath, cfg); err != nil {
				return nil, err
			}
			if err := checkPythonVenv(runtimePath, venvPath, uid, gid); err != nil {
				return nil, err
			}
			if err := checkPythonProject(venvPath, projectPath, uid, gid, cfg); err != nil {
				return nil, err
			}

			venvExe := getPythonExe(venvPath)
			return &RuntimeSandboxConfig{
				StartCmd: venvExe,
				Envs: map[string]string{
					"PATH": filepath.Dir(venvExe) + string(os.PathListSeparator) + "$PATH",
				},
			}, nil
		},
	})
}

func checkPythonRuntime(runtimePath string, cfg *RuntimeConfig) error {
	runtimeVersion := filepath.Base(runtimePath)
	lock := getRuntimeLock("python", "runtime", runtimePath)
	lock.Lock()
	defer lock.Unlock()
	if !u.FileExists(runtimePath) {
		osPart := u.StringIf(runtime.GOOS == "linux", "unknown-linux-gnu", u.StringIf(runtime.GOOS == "windows", "pc-windows-msvc-shared", "apple-darwin"))
		archPart := u.StringIf(runtime.GOARCH == "arm64", "aarch64", "x86_64")
		instFile := filepath.Join(filepath.Dir(runtimePath), ".installs", fmt.Sprintf("%s-%s-%s.tar.gz", runtimeVersion, osPart, archPart))
		if !u.FileExists(instFile) {
			downloadClient := httpclient.GetClient(3600 * time.Second)
			downloadClient.EnableRedirect()
			if cfg.HttpProxy != "" {
				if proxyURL, err := url.Parse(cfg.HttpProxy); err == nil {
					downloadClient.GetRawClient().Transport = &http.Transport{Proxy: http.ProxyURL(proxyURL)}
				}
			}
			log.DefaultLogger.Info("[Sandbox] fetching python runtime release list", "from", "https://github.com/astral-sh/python-build-standalone/releases", "httpProxy", cfg.HttpProxy)
			r1 := downloadClient.Get("https://github.com/astral-sh/python-build-standalone/releases")
			if r1.Error != nil {
				return r1.Error
			}
			if m1 := regexp.MustCompile(`data-deferred-src="([^"]+)"`).FindStringSubmatch(r1.String()); len(m1) >= 2 {
				log.DefaultLogger.Info("[Sandbox] fetching python runtime download url", "from", m1[1], "httpProxy", cfg.HttpProxy)
				r2 := downloadClient.Get(m1[1])
				if r2.Error != nil {
					return r2.Error
				}
				pattern := fmt.Sprintf(`href="([^"]+cpython-%s(?:\.\d+)*\+[^"]+-%s-%s-install_only\.tar\.gz)"`, regexp.QuoteMeta(runtimeVersion), archPart, osPart)
				if m2 := regexp.MustCompile(pattern).FindStringSubmatch(r2.String()); len(m2) >= 2 {
					log.DefaultLogger.Info("[Sandbox] downloading python runtime", "from", "https://github.com"+m2[1], "to", instFile)
					if _, err := downloadClient.Download(instFile, "https://github.com"+m2[1], nil); err != nil {
						return err
					} else {
						log.DefaultLogger.Info("[Sandbox] downloaded python runtime", "version", runtimeVersion, "os", osPart, "arch", archPart, "url", "https://github.com"+m2[1], "to", instFile, "httpProxy", cfg.HttpProxy)
					}
				} else {
					return fmt.Errorf("no binary for version %s on %s-%s on %s", runtimeVersion, osPart, archPart, m1[1])
				}
			} else {
				return errors.New("can't read release list on https://github.com/astral-sh/python-build-standalone/releases")
			}
		}

		if err := u.Extract(instFile, runtimePath, true); err != nil {
			log.DefaultLogger.Error("[Sandbox] unpack python runtime failed", "version", runtimeVersion, "os", osPart, "arch", archPart, "from", instFile, "err", err.Error())
			return err
		}
	}
	return nil
}

func checkPythonVenv(runtimePath, venvPath string, uid, gid int) error {
	lock := getRuntimeLock("python", "venv", venvPath)
	lock.Lock()
	defer lock.Unlock()
	if !u.FileExists(venvPath) {
		pyExe := getPythonExe(runtimePath)
		os.MkdirAll(venvPath, 0755)
		if uid != 0 && runtime.GOOS == "linux" {
			os.Chown(venvPath, uid, gid)
		}
		if out, err := RunCmd(venvPath, uid, gid, pyExe, "-m", "venv", venvPath); err != nil {
			log.DefaultLogger.Error("[Sandbox] venv create failed", "runtime", runtimePath, "venv", venvPath, "err", err.Error(), "output", out)
			return err
		} else {
			log.DefaultLogger.Info("[Sandbox] venv create success", "runtime", runtimePath, "venv", venvPath)
		}
	}
	return nil
}

func checkPythonProject(venvPath, projectPath string, uid, gid int, cfg *RuntimeConfig) error {
	lock := getRuntimeLock("python", "project", projectPath)
	lock.Lock()
	defer lock.Unlock()
	if uid != 0 && runtime.GOOS == "linux" {
		os.Chown(projectPath, uid, gid)
	}
	reqFile := filepath.Join(projectPath, "requirements.txt")
	if u.FileExists(reqFile) {
		shaFile := filepath.Join(projectPath, ".requirements.sha1")
		currentSha := u.Sha1String(u.ReadFileN(reqFile))
		oldSha := u.ReadFileN(shaFile)
		if oldSha != currentSha {
			pipExe := getPipExe(venvPath)
			pipArgs := []string{"install", "-r", reqFile}
			if cfg.Mirror != "" {
				pipArgs = append(pipArgs, "-i", cfg.Mirror)
			}

			if out, err := RunCmd(projectPath, uid, gid, pipExe, pipArgs...); err != nil {
				log.DefaultLogger.Error("[Sandbox] pip install failed", "venv", venvPath, "project", projectPath, "err", err.Error(), "output", out)
				return err
			} else {
				log.DefaultLogger.Info("[Sandbox] pip install success", "venv", venvPath, "project", projectPath, "output", out)
				u.WriteFile(shaFile, currentSha)
			}
		}
	}
	return nil
}

func getPythonExe(base string) string {
	if runtime.GOOS == "windows" {
		return filepath.Join(base, "python.exe")
	}
	exe := filepath.Join(base, "bin", "python")
	if !u.FileExists(exe) {
		exe = filepath.Join(base, "bin", "python3")
	}
	return exe
}

func getPipExe(base string) string {
	if runtime.GOOS == "windows" {
		return filepath.Join(base, "pip.exe")
	}
	exe := filepath.Join(base, "bin", "pip")
	if !u.FileExists(exe) {
		exe = filepath.Join(base, "bin", "pip3")
	}
	return exe
}