code.cestus.io/tools/fabricator@v0.4.3/pkg/helpers/exec.go (about)

     1  package helpers
     2  
     3  import (
     4  	"context"
     5  	"encoding/json"
     6  	"fmt"
     7  	"os"
     8  	"os/exec"
     9  	"runtime"
    10  	"strings"
    11  	"time"
    12  
    13  	"code.cestus.io/tools/fabricator/pkg/fabricator"
    14  )
    15  
    16  type Executor struct {
    17  	root string
    18  	io   fabricator.IOStreams
    19  	env  fabricator.Environment
    20  }
    21  
    22  func NewExecutor(root string, io fabricator.IOStreams) *Executor {
    23  	return &Executor{
    24  		root: root,
    25  		io:   io,
    26  		env:  make(fabricator.Environment),
    27  	}
    28  }
    29  
    30  func (e *Executor) WithRoot(root string) *Executor {
    31  	return &Executor{
    32  		root: root,
    33  		io:   e.io,
    34  		env:  e.env,
    35  	}
    36  }
    37  
    38  func (e *Executor) WithEnv(key, value string) *Executor {
    39  	env := fabricator.Environment{}
    40  	env[key] = value
    41  
    42  	return e.WithEnvMap(env)
    43  }
    44  
    45  func (e *Executor) WithEnvMap(env fabricator.Environment) *Executor {
    46  	newEnv := fabricator.Environment{}
    47  
    48  	if e.env != nil {
    49  		for k, v := range e.env {
    50  			newEnv[k] = v
    51  		}
    52  	}
    53  
    54  	for k, v := range env {
    55  		newEnv[k] = v
    56  	}
    57  
    58  	return &Executor{
    59  		root: e.root,
    60  		io:   e.io,
    61  		env:  newEnv,
    62  	}
    63  }
    64  
    65  func (e *Executor) setEnv(cmd *exec.Cmd) {
    66  	if e.env != nil {
    67  		env := append(os.Environ())
    68  
    69  		for key, value := range e.env {
    70  			env = append(env, fmt.Sprintf("%s=%s", key, value))
    71  		}
    72  
    73  		cmd.Env = env
    74  	}
    75  }
    76  
    77  func (e *Executor) Output(ctx context.Context, path string, args ...string) (string, error) {
    78  	cmd := exec.CommandContext(ctx, path, args...)
    79  	cmd.Dir = e.root
    80  	cmd.Stderr = e.io.ErrOut
    81  	e.setEnv(cmd)
    82  
    83  	data, err := cmd.Output()
    84  
    85  	if err != nil {
    86  		return "", err
    87  	}
    88  
    89  	return string(data), nil
    90  }
    91  
    92  func (e *Executor) JSONOutput(ctx context.Context, target interface{}, path string, args ...string) error {
    93  	jsonData, err := e.Output(ctx, path, args...)
    94  
    95  	if err != nil {
    96  		return err
    97  	}
    98  
    99  	if err = json.Unmarshal([]byte(jsonData), target); err != nil {
   100  		return fmt.Errorf("failed to JSON-decode the program's output: %s", err)
   101  	}
   102  
   103  	return nil
   104  }
   105  
   106  func (e *Executor) Run(ctx context.Context, path string, args ...string) error {
   107  	if e.io.Out != nil {
   108  		fmt.Fprintf(e.io.Out, "executing %s %s\n", path, strings.Join(args, " "))
   109  	}
   110  
   111  	cmd := exec.Command(path, args...)
   112  	cmd.Dir = e.root
   113  	cmd.Stdout = e.io.Out
   114  	cmd.Stderr = e.io.ErrOut
   115  	e.setEnv(cmd)
   116  
   117  	go func() {
   118  		<-ctx.Done()
   119  		// Do not send signals on a terminated process
   120  		if cmd.Process == nil {
   121  			return
   122  		}
   123  
   124  		if runtime.GOOS == "windows" {
   125  			_ = cmd.Process.Signal(os.Kill)
   126  			return
   127  		}
   128  
   129  		go func() {
   130  			time.Sleep(5 * time.Second)
   131  			_ = cmd.Process.Signal(os.Kill)
   132  		}()
   133  
   134  		cmd.Process.Signal(os.Interrupt)
   135  	}()
   136  
   137  	return cmd.Run()
   138  }