github.com/vchain-us/vcn@v0.9.11-0.20210921212052-a2484d23c0b3/pkg/bom/python/pip.go

/*
 * Copyright (c) 2021 CodeNotary, Inc. All Rights Reserved.
 * This software is released under GPL3.
 * The full license information can be found under:
 * https://www.gnu.org/licenses/gpl-3.0.en.html
 *
 */

package python

import (
	"bufio"
	"bytes"
	"errors"
	"fmt"
	"io/ioutil"
	"os/exec"
	"path/filepath"
	"strings"

	"github.com/schollz/progressbar/v3"

	"github.com/vchain-us/vcn/pkg/bom/artifact"
)

// pythonArtifactFromPip implements Artifact interface
type pythonArtifactFromPip struct {
	pythonArtifact
}

type module struct {
	needed  bool
	version string
}

type task struct {
	name    string
	version string
}
type result struct {
	name     string
	hash     string
	hashType artifact.HashType
	deps     []string
	license  string
	version  string
	err      error
}

var installArgs = []string{"-m", "pip", "install", "-r"}
var moduleListArgs = []string{"-m", "pip", "list", "-v"}
var moduleDetailsArgs = []string{"-m", "pip", "show"}

// Dependencies returns list of Python dependencies for the artifact
// collect info about all installed modules, find module relations, populate module graph and then recursively
// select only the needed modules, using content of 'requirements.txt' as a starting point
func (a *pythonArtifactFromPip) ResolveDependencies(output artifact.OutputOptions) ([]artifact.Dependency, error) {
	if a.Deps != nil {
		return a.Deps, nil
	}

	// install all required dependencies
	// first try "python", if it fails, try "python3"
	pythonExe := "python"
	reqFile := filepath.Join(a.path, pipFileName)
	cmdArgs := append(installArgs, reqFile)
	err := exec.Command(pythonExe, cmdArgs...).Run()
	if err != nil {
		pythonExe = "python3"
		_, err = exec.Command(pythonExe, cmdArgs...).Output()
		if err != nil {
			exit, ok := err.(*exec.ExitError)
			if ok && len(exit.Stderr) > 0 {
				return nil, fmt.Errorf("cannot install python modules:\n%s", string(exit.Stderr))
			}
			return nil, fmt.Errorf("cannot install python modules: %w", err)
		}
	}

	// collect info about all cached modules
	buf, err := exec.Command(pythonExe, moduleListArgs...).Output()
	if err != nil {
		exit, ok := err.(*exec.ExitError)
		if ok && len(exit.Stderr) > 0 {
			return nil, fmt.Errorf("cannot get python module list:\n%s", string(exit.Stderr))
		}
		return nil, fmt.Errorf("cannot get python module list: %w", err)
	}

	// output has two header lines, and then entries of the format "<package> <version> <location> <installer>"
	scanner := bufio.NewScanner(bytes.NewReader(buf))
	scanner.Split(bufio.ScanLines)
	for i := 0; i < 2; {
		if !scanner.Scan() {
			return nil, errors.New("got unexpected result to pip module list request")
		}
		text := scanner.Text()
		if text[0] == '#' {
			// skip possible Python warnings - lines starting with #
			continue
		}
		i++
	}

	// store all known modules
	moduleGraph := make(map[string]*module)
	for scanner.Scan() {
		text := scanner.Text()
		if text[0] == '#' {
			// skip possible Python warnings - lines starting with #
			continue
		}
		fields := strings.Fields(text)
		moduleGraph[fields[0]] = &module{version: fields[1]}
	}

	// first process root dependencies from requirements.txt
	// then for every dependency process its dependencies
	buf, err = ioutil.ReadFile(reqFile)
	if err != nil {
		return nil, err
	}

	var bar *progressbar.ProgressBar
	if output == artifact.Progress {
		bar = progressbar.Default(int64(len(moduleGraph)))
	}

	// init goroutine throttling - channels, start goroutines.
	// We can be sure that there will be no more in-flight messages in channels than known modules
	tasks := make(chan task, len(moduleGraph))
	results := make(chan result, len(moduleGraph))
	for i := 0; i < artifact.MaxGoroutines; i++ {
		go pipWorker(tasks, results, pythonExe, output, bar)
	}
	defer close(results)
	defer close(tasks) // signal workers to stop

	taskCount := 0

	messages := make([]string, 0)
	// initial tasks - content of requirements.txt
	scanner = bufio.NewScanner(bytes.NewReader(buf))
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		line := scanner.Text()
		fields := strings.SplitN(line, "#", 1)
		line = strings.TrimSpace(fields[0])
		if line == "" {
			continue
		}
		endPos := strings.IndexAny(line, "=><!")
		if endPos > 0 {
			line = line[:endPos]
		}

		mod, ok := moduleGraph[line]
		if !ok {
			name, err := resolveAlias(pythonExe, line)
			if err != nil {
				return nil, err
			}
			if name == "" {
				messages = append(messages, "Unknown module "+line+" - ignoring")
				continue
			}
			mod, ok = moduleGraph[name]
			if !ok {
				messages = append(messages, "Unknown module "+line+" - ignoring")
				continue
			}
			moduleGraph[line] = mod // add alias entry
		}

		if mod.needed {
			continue // already being processed by other name
		}
		mod.needed = true

		tasks <- task{name: line, version: mod.version}
		taskCount++
	}

	// get dependencies, run tasks for dependencies, collect info about all used modules
	res := make([]artifact.Dependency, 0)
	for done := 0; taskCount == 0 || done < taskCount; done++ {
		result := <-results
		if result.err != nil {
			close(tasks) // signal workers to stop
			return nil, err
		}
		res = append(res, artifact.Dependency{
			Name:     result.name,
			Version:  result.version,
			Hash:     result.hash,
			HashType: result.hashType,
			License:  result.license})
		for _, v := range result.deps {
			if v == "" {
				continue
			}
			mod, ok := moduleGraph[v]
			if !ok {
				name, err := resolveAlias(pythonExe, v)
				if err != nil {
					return nil, err
				}
				if name == "" {
					messages = append(messages, "Unknown module "+v+" - ignoring")
					continue
				}
				mod, ok = moduleGraph[name]
				if !ok {
					messages = append(messages, "Unknown module "+v+" - ignoring")
					continue
				}
				moduleGraph[v] = mod // add alias entry
			}

			if mod.needed {
				continue // already being processed
			}
			mod.needed = true
			tasks <- task{name: v, version: mod.version}
			taskCount++
		}
	}
	if bar != nil {
		bar.ChangeMax(taskCount) // to make 100% progress bar
	}
	for _, m := range messages {
		fmt.Println(m)
	}

	a.Deps = res
	return res, nil
}

func pipWorker(tasks <-chan task, results chan<- result, pythonExe string, output artifact.OutputOptions, bar *progressbar.ProgressBar) {
	for task := range tasks {
		lic, hashType, hash, err := QueryPkgDetails(task.name, task.version)
		if err != nil {
			results <- result{err: err}
			continue
		}
		deps, err := preRequisites(pythonExe, task.name)
		if err != nil {
			results <- result{err: err}
			continue
		}

		results <- result{name: task.name, version: task.version, hash: hash, deps: deps, hashType: hashType, license: lic, err: nil}
		switch output {
		case artifact.Progress:
			bar.Add(1)
		case artifact.Debug:
			fmt.Printf("%s@%s (%s)\n", task.name, task.version, hash)
		}
	}
}

func preRequisites(pythonExe string, module string) ([]string, error) {
	output, err := exec.Command(pythonExe, append(moduleDetailsArgs, module)...).Output()
	if err != nil {
		exit, ok := err.(*exec.ExitError)
		if ok && len(exit.Stderr) > 0 {
			return nil, fmt.Errorf("cannot get python module details:\n%s", string(exit.Stderr))
		}
		return nil, fmt.Errorf("cannot get python module details: %w", err)
	}

	scanner := bufio.NewScanner(bytes.NewReader(output))
	scanner.Split(bufio.ScanLines)
	var deps []string
	for scanner.Scan() {
		fields := strings.SplitN(scanner.Text(), ": ", 2)
		if fields[0] == "Requires" {
			deps = strings.Split(fields[1], ", ")
			break
		}
	}
	return deps, nil
}

// sometimes module is known by several names, in this case query module
// by its alias and add en entry for hashmap, pointing to the module
func resolveAlias(pythonExe, module string) (string, error) {
	output, err := exec.Command(pythonExe, append(moduleDetailsArgs, module)...).Output()
	if err != nil {
		exit, ok := err.(*exec.ExitError)
		if ok && len(exit.Stderr) > 0 {
			return "", fmt.Errorf("cannot get python module by alias:\n%s", string(exit.Stderr))
		}
		return "", fmt.Errorf("cannot get python module by alias: %w", err)
	}

	scanner := bufio.NewScanner(bytes.NewReader(output))
	scanner.Split(bufio.ScanLines)
	for scanner.Scan() {
		fields := strings.SplitN(scanner.Text(), ": ", 2)
		if fields[0] == "Name" {
			return fields[1], nil
		}
	}
	return "", nil
}