runner.go

package core

import (
	"errors"
	"fmt"
	"github.com/panjf2000/ants"
	"leviathan/libs"
	"leviathan/utils"
	"os"
	"path"
	"sync"
)

type Runner struct {
	Target    string
	Workspace string
	Workflow  libs.Workflow
	Options   *libs.Options
	Reports   []string
	DSL       map[string]interface{}
}

func InitRunner(target string, options *libs.Options) (*Runner, error) {
	var runner Runner
	runner.Target = target
	runner.Options = options

	runner.Workspace = path.Join(options.Environment.Workspaces, utils.CleanPath(target))
	if _, err := os.Stat(runner.Workspace); os.IsNotExist(err) {
		utils.Logger.Debug().Str("workspace", runner.Workspace).Msg("Create workspace folder")
		if err = os.MkdirAll(runner.Workspace, 0700); err != nil {
			return &runner, err
		}
	}

	if err := runner.readWorkflow(); err != nil {
		return &runner, err
	}

	runner.initDSL()
	return &runner, nil
}

func (r *Runner) readWorkflow() error {
	// check if workflow exists
	utils.Logger.Debug().Str("workflow", r.Options.Scan.Flow).Msg("Check if workflow exists")
	if !WorkflowExists(r.Options.Scan.Flow, r.Options.Environment.Workflows) {
		return errors.New(fmt.Sprintf("%s workflow was not found", r.Options.Scan.Flow))
	}

	// parse workflow
	utils.Logger.Debug().Str("workflow", r.Options.Scan.Flow).Msg("Try to parse workflow")
	workflow, err := WorkflowParse(r)
	if err != nil {
		return err
	}
	r.Workflow = workflow

	// retrieve all modules used by the workflow
	for id, routine := range workflow.Routines {
		var modules []libs.Module
		for _, moduleName := range routine.Modules {
			utils.Logger.Debug().Str("module", moduleName).Msg("Try to parse module")
			module, err := ModuleParse(moduleName, r)
			if err != nil {
				return err
			}

			modules = append(modules, module)
		}
		workflow.Routines[id].ModulesParsed = modules
	}

	return nil
}

func (r *Runner) Start() {
	utils.Logger.Info().
		Str("runner", r.Target).
		Str("workflow", r.Workflow.Name).
		Msg("Start runner")

	for _, routine := range r.Workflow.Routines {
		var wg sync.WaitGroup
		p, _ := ants.NewPoolWithFunc(r.Options.Scan.Threads*10, func(m interface{}) {
			module := m.(libs.Module)
			r.RunModule(module)
			wg.Done()
		}, ants.WithPreAlloc(true))
		defer p.Release()

		for _, module := range routine.ModulesParsed {
			p.Invoke(module)
			wg.Add(1)
		}

		wg.Wait()
	}

	if !r.Options.Scan.NoClean {
		utils.Logger.Info().Str("workspace", r.Workspace).Msg("Clean workspace")
		CleanWorkspace(r.Workspace, r.Reports)
	}
}

func (r *Runner) RunModule(m libs.Module) {
	utils.Logger.Info().Str("module", m.Name).Msg("Run module")

	// Save the reports generated by the module
	for _, report := range m.Reports {
		utils.Logger.Debug().Str("module", m.Name).Str("report", report).Msg("Add report to the list")
		r.Reports = append(r.Reports, utils.NormalizePath(report))
	}

	// Check if module has to be skipped
	if r.Options.Scan.Resume {
		skipModule := true
		for _, report := range m.Reports {
			if !utils.FileExists(report) {
				skipModule = false
			}
		}
		if skipModule {
			utils.Logger.Info().Str("module", m.Name).Msg("skip module (--resume flag used)")
			return
		}
	}

	// Execute pre_run scripts
	if m.PreRun != nil {
		utils.Logger.Debug().Str("module", m.Name).Msg("Execute Pre_run scripts")
		r.runScripts(m.PreRun)
	}

	utils.Logger.Debug().Str("module", m.Name).Msg("Execute step")
	for _, step := range m.Steps {
		// Check requirements, skip step if requirements not met
		if !r.CheckRequirements(step.Requirements) {
			continue
		}

		// If conditions are validated
		if r.CheckConditions(step.Conditions) {
			// Execute commands
			r.runCommands(step.Commands)

			// Execute scripts
			r.runScripts(step.Scripts)
		} else {
			// Execute commands
			r.runCommands(step.RCommands)

			// Execute scripts
			r.runScripts(step.RScripts)
		}

	}

	if m.PostRun != nil {
		utils.Logger.Debug().Str("module", m.Name).Msg("Execute post_run scripts")
		r.runScripts(m.PostRun)
	}

}