osmedeus/core/module.go

package core

import (
	"context"
	"fmt"
	"path"
	"strings"
	"sync"
	"time"

	"github.com/spf13/cast"

	"github.com/fatih/color"
	"github.com/j3ssie/osmedeus/execution"
	"github.com/j3ssie/osmedeus/libs"
	"github.com/j3ssie/osmedeus/utils"
	"github.com/panjf2000/ants"
)

// RunModule run the module
func (r *Runner) RunModule(module libs.Module) {
	// get reports path
	module = ResolveReports(module, r.Params)

	// check if resume enable or not
	if (r.Opt.Resume || module.Resume) && !module.Forced {
		if CheckResume(module) {
			utils.TSPrintF("The %v module has resume", color.HiGreenString(module.Name))
			return
		}
	}

	r.CurrentModule = module.Name
	timeStart := time.Now()
	utils.TSPrintF("The %v module has begun with the objective %v", color.HiGreenString(module.Name), color.HiCyanString(module.Desc))

	// create report record first because I don't want to wait for them to show up in UI until the module done
	r.DBNewReports(module)

	// pre-run
	if len(module.PreRun) > 0 && r.Opt.NoPreRun == false {
		utils.InforF("Running prepare scripts for module %v", color.CyanString(module.Name))
		r.RunScripts(module.PreRun)
	}

	utils.InforF("Running steps for module %v", color.CyanString(module.Name))
	// main part
	err := r.RunSteps(module.Steps)
	if err != nil {
		utils.BadBlockF(fmt.Sprintf("got an exit call"))
	}

	// post-run
	if len(module.PostRun) > 0 && r.Opt.NoPostRun == false {
		utils.InforF("Running conclude scripts for module %v", color.CyanString(module.Name))
		r.RunScripts(module.PostRun)
	}

	// print the reports file
	printReports(module)

	// estimate time
	elapsedTime := time.Since(timeStart).Seconds()
	utils.TSPrintF("The %v module finished within %v.", color.HiGreenString(module.Name), color.HiMagentaString("%vs", elapsedTime))

	r.RunningTime += cast.ToInt(elapsedTime)

	r.DBUpdateScan()
}

// RunScripts run list of scripts
func (r *Runner) RunScripts(scripts []string) string {
	if r.Opt.Timeout != "" {
		timeout := utils.CalcTimeout(r.Opt.Timeout)
		utils.DebugF("Run scripts with %v seconds timeout", timeout)
		r.RunScriptsWithTimeOut(r.Opt.Timeout, scripts)
		return ""
	}

	for _, script := range scripts {
		outScript := r.RunScript(script)
		if strings.Contains(outScript, "exit") {
			return outScript
		}
	}
	return ""
}

// RunScriptsWithTimeOut run list of scripts with timeout
func (r *Runner) RunScriptsWithTimeOut(timeoutRaw string, scripts []string) string {
	timeout := utils.CalcTimeout(timeoutRaw)
	utils.DebugF("Run scripts with %v seconds timeout", timeout)

	c := context.Background()
	deadline := time.Now().Add(time.Duration(timeout) * time.Second)
	c, cancel := context.WithDeadline(c, deadline)
	defer cancel()

	go func() {
		for _, script := range scripts {
			outScript := r.RunScript(script)
			if strings.Contains(outScript, "exit") {
				return
			}
		}
		cancel()
	}()

	select {
	case <-c.Done():
		utils.DebugF("Scripts done")
		return ""
	case <-time.After(time.Duration(timeout) * time.Second):
		utils.BadBlockF(fmt.Sprintf("Scripts got timeout after %v", color.HiMagentaString(timeoutRaw)))
	}
	return ""
}

// RunScript really run a script
func (r *Runner) RunScript(script string) string {
	return r.ExecScript(script)
}

// RunSteps run list of steps
func (r *Runner) RunSteps(steps []libs.Step) error {
	var stepOut string
	for _, step := range steps {
		r.DoneStep += 1

		if step.Timeout != "" {
			// timeout should be: 30, 30m, 1h
			timeout := utils.CalcTimeout(step.Timeout)
			if timeout != 0 {
				stepOut, _ = r.RunStepWithTimeout(timeout, step)
				if strings.Contains(stepOut, "exit") {
					return fmt.Errorf("got exit call")
				}
				continue
			}
		}

		stepOut, _ = r.RunStep(step)
		if strings.Contains(stepOut, "exit") {
			return fmt.Errorf("got an exit call")
		}
	}
	return nil
}

// RunStepWithTimeout run step with timeout
func (r *Runner) RunStepWithTimeout(timeout int, step libs.Step) (out string, err error) {
	utils.DebugF("Run step with %v seconds timeout", timeout)
	prefix := fmt.Sprintf("timeout -k 1m %vs ", timeout)

	// prepare the os command with prefix timeout first
	var preFixCommands []string
	for _, command := range step.Commands {
		preFixCommand := command
		if !strings.Contains(command, "timeout") {
			preFixCommand = prefix + command
		}
		preFixCommands = append(preFixCommands, preFixCommand)
	}
	step.Commands = preFixCommands

	// override global timeout
	r.Opt.Timeout = step.Timeout
	return r.RunStep(step)
}

func (r *Runner) RunStep(step libs.Step) (string, error) {
	var output string
	if step.Label != "" {
		utils.TSPrintF("Initiating step %v", color.HiGreenString(step.Label))
	}

	// checking required file
	err := r.CheckRequired(step.Required)
	if err != nil {
		return output, fmt.Errorf("missing requirements")
	}

	// check conditions and run reverse step
	err = r.CheckCondition(step.Conditions)
	if err != nil {
		if len(step.RCommands) == 0 && len(step.RScripts) == 0 {
			return output, fmt.Errorf("conditions not met")
		}

		// run reverse commands
		utils.InforF("Condition false, run the reverse commands")
		if len(step.RCommands) > 0 {
			r.RunCommands(step.RCommands, step.Std)
		}
		// run reverse scripts
		if len(step.RScripts) > 0 {
			output = r.RunScripts(step.RScripts)
			if strings.Contains(output, "exit") {
				return output, nil
			}
		}
		return output, nil
	}

	// run the step in loop mode
	if step.Source != "" {
		return r.RunStepWithSource(step)
	}

	if len(step.Commands) > 0 {
		r.RunCommands(step.Commands, step.Std)
	}
	if len(step.Scripts) > 0 {
		output = r.RunScripts(step.Scripts)
		if strings.Contains(output, "exit") {
			return output, nil
		}
	}

	// run ose here
	if len(step.Ose) > 0 {
		for _, ose := range step.Ose {
			r.RunOse(ose)
		}
	}

	// post scripts
	if len(step.PConditions) > 0 || len(step.PScripts) > 0 {
		err := r.CheckCondition(step.PConditions)
		if err == nil {
			if len(step.PScripts) > 0 {
				r.RunScripts(step.PScripts)
			}
		}
	}
	return output, nil

}

// RunStepWithSource really run a step
func (r *Runner) RunStepWithSource(step libs.Step) (out string, err error) {
	////// Start to run step but in loop mode
	utils.DebugF("Running the step using the source file: %v", step.Source)
	data := utils.ReadingLines(step.Source)
	if len(data) <= 0 {
		return out, fmt.Errorf("missing source")
	}
	if step.Threads != "" {
		step.Parallel = cast.ToInt(step.Threads)
	}
	if step.Parallel == 0 {
		step.Parallel = 1
	}

	// prepare the data first
	var newGeneratedSteps []libs.Step
	for index, line := range data {
		customParams := make(map[string]string)
		customParams["line"] = line
		customParams["line_id"] = fmt.Sprintf("%v-%v", path.Base(line), index)
		customParams["_id_"] = fmt.Sprintf("%v", index)
		customParams["_line_"] = execution.StripName(line)

		// make completely new Step
		localStep := libs.Step{}

		for _, cmd := range step.Commands {
			localStep.Commands = append(localStep.Commands, AltResolveVariable(cmd, customParams))
		}
		for _, cmd := range step.RCommands {
			localStep.RCommands = append(localStep.RCommands, AltResolveVariable(cmd, customParams))
		}

		if len(step.Ose) > 0 {
			for _, ose := range step.Ose {
				localStep.Ose = append(localStep.Ose, AltResolveVariable(ose, customParams))
			}
		}

		for _, script := range step.RScripts {
			localStep.RScripts = append(localStep.RScripts, AltResolveVariable(script, customParams))
		}

		for _, script := range step.Scripts {
			localStep.Scripts = append(localStep.Scripts, AltResolveVariable(script, customParams))
		}

		for _, script := range step.PConditions {
			localStep.PConditions = append(localStep.PConditions, AltResolveVariable(script, customParams))
		}
		for _, script := range step.PScripts {
			localStep.PScripts = append(localStep.PScripts, AltResolveVariable(script, customParams))
		}

		newGeneratedSteps = append(newGeneratedSteps, localStep)
	}

	// skip concurrency part
	if step.Parallel == 1 {
		for _, newGeneratedStep := range newGeneratedSteps {
			out, err = r.RunStep(newGeneratedStep)
			if err != nil {
				continue
			}
		}
	} else {
		/////////////
		// run multiple steps in concurrency mode

		utils.DebugF("Running the step in parallel: %v", step.Parallel)
		var wg sync.WaitGroup
		p, _ := ants.NewPoolWithFunc(step.Parallel, func(i interface{}) {
			r.startStepJob(i)
			wg.Done()
		}, ants.WithPreAlloc(true))
		defer p.Release()

		for _, newGeneratedStep := range newGeneratedSteps {
			wg.Add(1)
			err = p.Invoke(newGeneratedStep)
			if err != nil {
				utils.ErrorF("Error in parallel: %v", err)
			}
		}

		wg.Wait()
	}
	return out, nil
}

func (r *Runner) startStepJob(j interface{}) {
	localStep := j.(libs.Step)

	err := r.CheckCondition(localStep.Conditions)

	if err != nil {
		// run reverse commands
		if len(localStep.RCommands) > 0 {
			r.RunCommands(localStep.RCommands, localStep.Std)
		}
		if len(localStep.RScripts) > 0 {
			r.RunScripts(localStep.RScripts)
		}
	} else {
		if len(localStep.Commands) > 0 {
			r.RunCommands(localStep.Commands, localStep.Std)
		}
	}

	if len(localStep.Ose) > 0 {
		for _, ose := range localStep.Ose {
			r.RunOse(ose)
		}
	}

	if len(localStep.Scripts) > 0 {
		r.RunScripts(localStep.Scripts)
	}

	// post scripts
	if len(localStep.PConditions) > 0 || len(localStep.PScripts) > 0 {
		err := r.CheckCondition(localStep.PConditions)
		if err == nil {
			if len(localStep.PScripts) > 0 {
				r.RunScripts(localStep.PScripts)
			}
		}
	}
}