train-osource/pkg/installer/installer.go

package installer

import (
	"archive/tar"
	"compress/gzip"
	"crypto/sha256"
	"encoding/hex"
	"encoding/json"
	"errors"
	"fmt"
	"io"
	"io/ioutil"
	"net/http"
	"os"
	"os/exec"
	"path/filepath"
	"runtime"
	"strings"

	"train/pkg/manifest"
	"train/pkg/paths"
)

var (
	binInstallDir   = filepath.Join(paths.BaseDir(), "installed_bins")
	buildInstallDir = filepath.Join(paths.BaseDir(), "installed_packages")
)

func computeChecksum(filePath string) (string, error) {
	data, err := ioutil.ReadFile(filePath)
	if err != nil {
		return "", err
	}
	hash := sha256.Sum256(data)
	return hex.EncodeToString(hash[:]), nil
}

func downloadFile(url string) (string, error) {
	resp, err := http.Get(url)
	if err != nil {
		return "", err
	}
	defer resp.Body.Close()
	tmpFile, err := ioutil.TempFile("", "package_")
	if err != nil {
		return "", err
	}
	_, err = io.Copy(tmpFile, resp.Body)
	if err != nil {
		tmpFile.Close()
		return "", err
	}
	tmpFile.Close()
	return tmpFile.Name(), nil
}

func unpackPackage(filePath, packageName string) error {
	destDir := filepath.Join(buildInstallDir, packageName)
	if err := os.MkdirAll(destDir, os.ModePerm); err != nil {
		return err
	}
	f, err := os.Open(filePath)
	if err != nil {
		return err
	}
	defer f.Close()
	gzr, err := gzip.NewReader(f)
	if err != nil {
		return err
	}
	defer gzr.Close()
	tr := tar.NewReader(gzr)
	for {
		header, err := tr.Next()
		if err == io.EOF {
			break
		}
		if err != nil {
			return err
		}
		target := filepath.Join(destDir, header.Name)
		switch header.Typeflag {
		case tar.TypeDir:
			if err := os.MkdirAll(target, os.FileMode(header.Mode)); err != nil {
				return err
			}
		case tar.TypeReg:
			outFile, err := os.Create(target)
			if err != nil {
				return err
			}
			if _, err := io.Copy(outFile, tr); err != nil {
				outFile.Close()
				return err
			}
			outFile.Close()
		}
	}
	return nil
}

func executeBuildCommands(commands []string, packageName string) error {
	packagePath := filepath.Join(buildInstallDir, packageName)
	for _, commandStr := range commands {
		cmd := exec.Command("sh", "-c", commandStr)
		cmd.Dir = packagePath
		output, err := cmd.CombinedOutput()
		if err != nil {
			return fmt.Errorf("command '%s' failed: %s", commandStr, string(output))
		}
	}
	return nil
}

func saveInstalledPackage(name string, m *manifest.Manifest) error {
	packagePath := filepath.Join(buildInstallDir, name)
	if err := os.MkdirAll(packagePath, os.ModePerm); err != nil {
		return err
	}
	jsonData, err := json.MarshalIndent(m, "", "  ")
	if err != nil {
		return err
	}
	return ioutil.WriteFile(filepath.Join(packagePath, "manifest.json"), jsonData, 0644)
}

func InstallPackage(name string) error {
	m, err := manifest.FetchManifest(name)
	if err != nil {
		return fmt.Errorf("failed to fetch manifest: %w", err)
	}
	if strings.ToLower(m.Mode) == "bin" {
		platformKey := runtime.GOOS + "_" + runtime.GOARCH
		binaryInfo, ok := m.Binaries[platformKey]
		if !ok {
			return fmt.Errorf("no binary available for platform %s", platformKey)
		}
		packageFile, err := downloadFile(binaryInfo.URL)
		if err != nil {
			return fmt.Errorf("failed to download binary: %w", err)
		}
		defer os.Remove(packageFile)
		computedChecksum, err := computeChecksum(packageFile)
		if err != nil {
			return fmt.Errorf("failed to compute checksum: %w", err)
		}
		if computedChecksum != binaryInfo.Checksum {
			return fmt.Errorf("checksum mismatch: expected %s, got %s", binaryInfo.Checksum, computedChecksum)
		}
		if err := os.MkdirAll(binInstallDir, os.ModePerm); err != nil {
			return err
		}
		destPath := filepath.Join(binInstallDir, m.Name)
		data, err := ioutil.ReadFile(packageFile)
		if err != nil {
			return err
		}
		if err := ioutil.WriteFile(destPath, data, 0755); err != nil {
			return err
		}
		fmt.Printf("Binary package %s installed in %s\n", m.Name, destPath)
		return nil
	} else if strings.ToLower(m.Mode) == "build" {
		packageFile, err := downloadFile(m.Source)
		if err != nil {
			return fmt.Errorf("failed to download package: %w", err)
		}
		defer os.Remove(packageFile)
		computedChecksum, err := computeChecksum(packageFile)
		if err != nil {
			return fmt.Errorf("failed to compute checksum: %w", err)
		}
		if computedChecksum != m.Checksum {
			return fmt.Errorf("checksum mismatch: expected %s, got %s", m.Checksum, computedChecksum)
		}
		if err := unpackPackage(packageFile, name); err != nil {
			return fmt.Errorf("failed to unpack package: %w", err)
		}
		if len(m.Build) > 0 {
			if err := executeBuildCommands(m.Build, name); err != nil {
				return fmt.Errorf("build command error: %w", err)
			}
		}
		if err := saveInstalledPackage(name, m); err != nil {
			return fmt.Errorf("failed to save manifest: %w", err)
		}
		fmt.Printf("Build package %s installed successfully\n", m.Name)
		return nil
	}
	return errors.New("unknown package mode")
}

func RemovePackage(name string) error {
	packagePath := filepath.Join(buildInstallDir, name)
	if _, err := os.Stat(packagePath); os.IsNotExist(err) {
		return errors.New("package not found")
	}
	if err := os.RemoveAll(packagePath); err != nil {
		return fmt.Errorf("failed to remove package: %w", err)
	}
	return nil
}

func UpdatePackage(name string) error {
	if err := RemovePackage(name); err != nil {
		return fmt.Errorf("failed to remove package for update: %w", err)
	}
	if err := InstallPackage(name); err != nil {
		return fmt.Errorf("failed to install package: %w", err)
	}
	return nil
}

func ListInstalledPackages() ([]string, error) {
	var pkgs []string
	entries, err := ioutil.ReadDir(buildInstallDir)
	if err != nil {
		return nil, err
	}
	for _, entry := range entries {
		if entry.IsDir() {
			pkgs = append(pkgs, entry.Name())
		}
	}
	return pkgs, nil
}

func GetPackageInfo(name string) (string, error) {
	infoFile := filepath.Join(buildInstallDir, name, "manifest.json")
	data, err := ioutil.ReadFile(infoFile)
	if err != nil {
		return "", err
	}
	return string(data), nil
}