runbook.go

package runn

import (
	"errors"
	"fmt"
	"io"
	"net/http"
	"os"
	"strconv"
	"strings"

	"github.com/Songmu/axslogparser"
	"github.com/goccy/go-yaml"
	"github.com/goccy/go-yaml/ast"
	"github.com/goccy/go-yaml/lexer"
	"github.com/goccy/go-yaml/parser"
	"github.com/goccy/go-yaml/token"
	"github.com/k1LoW/curlreq"
	"github.com/k1LoW/expand"
	"github.com/k1LoW/grpcurlreq"
)

type position struct {
	Line int
	// Column int
}

type area struct {
	Start *position
	End   *position
}

type areas struct {
	Desc    *area
	Runners *area
	Vars    *area
	Steps   []*area
}

type runbook struct {
	Desc        string            `yaml:"desc"`
	Labels      []string          `yaml:"labels,omitempty"`
	Needs       map[string]string `yaml:"needs,omitempty"`
	Runners     map[string]any    `yaml:"runners,omitempty"`
	Vars        map[string]any    `yaml:"vars,omitempty"`
	Secrets     []string          `yaml:"secrets,omitempty"`
	Steps       []yaml.MapSlice   `yaml:"steps"`
	HostRules   yaml.MapSlice     `yaml:"hostRules,omitempty"`
	Debug       bool              `yaml:"debug,omitempty"`
	Interval    string            `yaml:"interval,omitempty"`
	If          string            `yaml:"if,omitempty"`
	SkipTest    bool              `yaml:"skipTest,omitempty"`
	Loop        any               `yaml:"loop,omitempty"`
	Concurrency any               `yaml:"concurrency,omitempty"`
	Force       bool              `yaml:"force,omitempty"`
	Trace       bool              `yaml:"trace,omitempty"`

	useMap   bool
	stepKeys []string
	wd       string // working directory for curl file uploads
}

type runbookListed runbook

type runbookMapped struct {
	Desc        string            `yaml:"desc,omitempty"`
	Labels      []string          `yaml:"labels,omitempty"`
	Needs       map[string]string `yaml:"needs,omitempty"`
	Runners     map[string]any    `yaml:"runners,omitempty"`
	Vars        map[string]any    `yaml:"vars,omitempty"`
	Secrets     []string          `yaml:"secrets,omitempty"`
	Steps       yaml.MapSlice     `yaml:"steps,omitempty"`
	HostRules   yaml.MapSlice     `yaml:"hostRules,omitempty"`
	Debug       bool              `yaml:"debug,omitempty"`
	Interval    string            `yaml:"interval,omitempty"`
	If          string            `yaml:"if,omitempty"`
	SkipTest    bool              `yaml:"skipTest,omitempty"`
	Loop        any               `yaml:"loop,omitempty"`
	Concurrency any               `yaml:"concurrency,omitempty"`
	Force       bool              `yaml:"force,omitempty"`
	Trace       bool              `yaml:"trace,omitempty"`
}

var decOpts = []yaml.DecodeOption{
	yaml.UseOrderedMap(),
}

var encOpts = []yaml.EncodeOption{
	yaml.UseLiteralStyleIfMultiline(true),
}

// NewRunbook creates a new runbook with the specified description.
// If no description is provided, a default one is used.
func NewRunbook(desc string) *runbook {
	const defaultDesc = "Generated by `runn new`"
	if desc == "" {
		desc = defaultDesc
	}
	r := &runbook{
		Desc:    desc,
		Needs:   map[string]string{},
		Runners: map[string]any{},
		Vars:    map[string]any{},
		Steps:   []yaml.MapSlice{},
		wd:      ".",
	}
	return r
}

// ParseRunbook parses a runbook from the provided reader.
// It reads the YAML content and converts it into a runbook structure.
func ParseRunbook(in io.Reader) (*runbook, error) {
	b, err := io.ReadAll(in)
	if err != nil {
		return nil, err
	}
	return parseRunbook(b)
}

func parseRunbook(b []byte) (*runbook, error) {
	rb := NewRunbook("")

	repFn := expand.InterpolateRepFn(os.LookupEnv)
	rep, err := expand.ReplaceYAML(string(b), repFn)
	if err != nil {
		return nil, err
	}

	if err := yaml.UnmarshalWithOptions([]byte(rep), rb, decOpts...); err != nil {
		if err := parseRunbookAsMapped([]byte(rep), rb); err != nil {
			return nil, err
		}
	}

	if err := rb.validate(); err != nil {
		return nil, err
	}

	return rb, nil
}

func parseRunbookAsMapped(b []byte, rb *runbook) error {
	m := &runbookMapped{}
	if err := yaml.UnmarshalWithOptions(b, m, decOpts...); err != nil {
		return err
	}
	rb.useMap = true
	rb.Desc = m.Desc
	rb.Labels = m.Labels
	rb.Needs = m.Needs
	rb.Runners = m.Runners
	rb.Vars = m.Vars
	rb.Secrets = m.Secrets
	rb.HostRules = m.HostRules
	rb.Debug = m.Debug
	rb.Interval = m.Interval
	rb.If = m.If
	rb.SkipTest = m.SkipTest
	rb.Loop = m.Loop
	rb.Concurrency = m.Concurrency
	rb.Force = m.Force
	rb.Trace = m.Trace

	keys := map[string]struct{}{}
	for _, s := range m.Steps {
		k, ok := s.Key.(string)
		if !ok {
			return fmt.Errorf("failed to parse as mapped steps: %v", s)
		}
		v, ok := s.Value.(yaml.MapSlice)
		if !ok {
			return fmt.Errorf("failed to parse as mapped steps: %v", s)
		}
		rb.stepKeys = append(rb.stepKeys, k)
		rb.Steps = append(rb.Steps, v)
		if _, ok := keys[k]; ok {
			return fmt.Errorf("duplicate step keys: %s", k)
		}
		keys[k] = struct{}{}
	}

	return nil
}

func (rb *runbook) AppendStep(in ...string) error {
	if len(in) == 0 {
		return errors.New("no argument")
	}
	if rb.useMap {
		key := fmt.Sprintf("%s%d", in[0], len(rb.stepKeys))
		rb.stepKeys = append(rb.stepKeys, key)
	}
	switch {
	case strings.HasPrefix(in[0], "curl"):
		return rb.curlToStep(in...)
	case strings.HasPrefix(in[0], "grpcurl"):
		return rb.grpcurlToStep(in...)
	default:
		if len(in) == 1 {
			if err := rb.axsLogToStep(in...); err == nil {
				return nil
			}
		}
		return rb.cmdToStep(in...)
	}
}

func (rb *runbook) MarshalYAML() (any, error) {
	if !rb.useMap {
		return &runbookListed{
			Desc:        rb.Desc,
			Labels:      rb.Labels,
			Needs:       rb.Needs,
			Runners:     rb.Runners,
			Vars:        rb.Vars,
			Secrets:     rb.Secrets,
			Steps:       rb.Steps,
			HostRules:   rb.HostRules,
			Debug:       rb.Debug,
			Interval:    rb.Interval,
			If:          rb.If,
			SkipTest:    rb.SkipTest,
			Loop:        rb.Loop,
			Concurrency: rb.Concurrency,
			Force:       rb.Force,
			Trace:       rb.Trace,

			useMap:   rb.useMap,
			stepKeys: rb.stepKeys,
		}, nil
	}
	if len(rb.stepKeys) != len(rb.Steps) {
		return nil, errors.New("invalid runbook")
	}
	m := &runbookMapped{}
	m.Desc = rb.Desc
	m.Labels = rb.Labels
	m.Needs = rb.Needs
	m.Runners = rb.Runners
	m.Vars = rb.Vars
	m.Secrets = rb.Secrets
	m.HostRules = rb.HostRules
	m.Debug = rb.Debug
	m.Interval = rb.Interval
	m.If = rb.If
	m.SkipTest = rb.SkipTest
	m.Loop = rb.Loop
	m.Concurrency = rb.Concurrency
	m.Force = rb.Force
	m.Trace = rb.Trace
	ms := yaml.MapSlice{}
	for i, k := range rb.stepKeys {
		ms = append(ms, yaml.MapItem{
			Key:   k,
			Value: rb.Steps[i],
		})
	}
	m.Steps = ms
	return m, nil
}

func (rb *runbook) curlToStep(in ...string) error {
	p, err := curlreq.NewParser(curlreq.WithWorkingDirectory(rb.wd))
	if err != nil {
		return err
	}
	parsed, err := p.Parse(in...)
	if err != nil {
		return err
	}
	req, err := parsed.Request()
	if err != nil {
		return err
	}

	splitted := strings.Split(req.URL.String(), req.URL.Host)
	dsn := fmt.Sprintf("%s%s", splitted[0], req.URL.Host)
	key := rb.setRunner(dsn)
	step, err := CreateHTTPStepMapSlice(key, req)
	if err != nil {
		return err
	}
	rb.Steps = append(rb.Steps, step)
	return nil
}

func (rb *runbook) grpcurlToStep(in ...string) error {
	p, err := grpcurlreq.Parse(in...)
	if err != nil {
		return err
	}
	if p.Addr == "" || p.Method == "" || p.SubCmd != "" {
		return fmt.Errorf("unsupported grpcurl command: %v", in)
	}
	dsn := fmt.Sprintf("grpc://%s", p.Addr)
	key := rb.setRunner(dsn)

	hm := yaml.MapSlice{}
	h := map[string]string{}
	for k, v := range p.Headers {
		h[k] = v[0]
	}
	if len(h) > 0 {
		hm = append(hm, yaml.MapItem{
			Key:   "headers",
			Value: h,
		})
	}

	// messages
	switch {
	case len(p.Messages) == 0:
		hm = append(hm, yaml.MapItem{
			Key:   "message",
			Value: map[string]any{},
		})
	case len(p.Messages) == 1:
		hm = append(hm, yaml.MapItem{
			Key:   "message",
			Value: p.Messages[0],
		})
	case len(p.Messages) > 1:
		hm = append(hm, yaml.MapItem{
			Key:   "messages",
			Value: p.Messages,
		})
	}

	if len(hm) == 0 {
		hm = nil
	}
	step := yaml.MapSlice{
		{Key: key, Value: yaml.MapSlice{
			{Key: p.Method, Value: hm},
		}},
	}
	rb.Steps = append(rb.Steps, step)
	return nil
}

func (rb *runbook) setRunner(dsn string) string {
	const (
		httpRunnerKeyPrefix = "req"
		grpcRunnerKeyPrefix = "greq"
		dbRunnerKeyPrefix   = "db"
	)
	var hc, gc, dc int
	for k, v := range rb.Runners {
		vv, ok := v.(string)
		if !ok {
			continue
		}
		if vv == dsn {
			return k
		}
		switch {
		case strings.HasPrefix(vv, "http"):
			hc += 1
		case strings.HasPrefix(vv, "grpc"):
			gc += 1
		default:
			dc += 1
		}
	}

	var key string
	switch {
	case strings.HasPrefix(dsn, "http"):
		if hc > 0 {
			key = fmt.Sprintf("%s%d", httpRunnerKeyPrefix, hc+1)
		} else {
			key = httpRunnerKeyPrefix
		}
	case strings.HasPrefix(dsn, "grpc"):
		if gc > 0 {
			key = fmt.Sprintf("%s%d", grpcRunnerKeyPrefix, gc+1)
		} else {
			key = grpcRunnerKeyPrefix
		}
	default:
		if dc > 0 {
			key = fmt.Sprintf("%s%d", dbRunnerKeyPrefix, dc+1)
		} else {
			key = dbRunnerKeyPrefix
		}
	}
	rb.Runners[key] = dsn
	return key
}

func (rb *runbook) axsLogToStep(in ...string) error {
	const dummyDSN = "https://dummy.example.com"
	line := strings.Join(in, " ")
	_, l, err := axslogparser.GuessParser(line)
	if err != nil {
		return err
	}
	dsn := dummyDSN
	key := rb.setRunner(dsn)
	req, err := http.NewRequest(l.Method, l.RequestURI, nil)
	if err != nil {
		return err
	}
	if l.UserAgent != "" {
		req.Header.Add("User-Agent", l.UserAgent)
	}
	step, err := CreateHTTPStepMapSlice(key, req)
	if err != nil {
		return err
	}
	rb.Steps = append(rb.Steps, step)
	return nil
}

func (rb *runbook) cmdToStep(in ...string) error {
	step := yaml.MapSlice{
		{Key: execRunnerKey, Value: yaml.MapSlice{
			{Key: "command", Value: joinCommands(in...)},
		}},
	}
	rb.Steps = append(rb.Steps, step)
	return nil
}

var invalidLabelTokens = []string{" ", "\n", "\r", "!", "+", "=", "|", ".", "*", "%", "^", "?", ">", "<"}

func (rb *runbook) validate() error {
	// labels:
	if len(rb.Labels) == 0 {
		return nil
	}
	for _, l := range rb.Labels {
		for _, t := range invalidLabelTokens {
			if strings.Contains(l, t) {
				return fmt.Errorf("invalid label: %q", l)
			}
		}
	}
	return nil
}

func (rb *runbook) toBook() (*book, error) {
	var (
		ok  bool
		err error
	)
	bk := newBook()
	bk.desc = rb.Desc
	bk.labels = rb.Labels
	bk.needs = rb.Needs
	bk.runners, ok = normalize(rb.Runners).(map[string]any)
	if !ok {
		return nil, fmt.Errorf("failed to normalize runners: %v", rb.Runners)
	}
	bk.vars, ok = normalize(rb.Vars).(map[string]any)
	if !ok {
		return nil, fmt.Errorf("failed to normalize vars: %v", rb.Vars)
	}
	bk.secrets = rb.Secrets
	for _, s := range rb.Steps {
		v, ok := normalize(s).(map[string]any)
		if !ok {
			return nil, fmt.Errorf("failed to normalize step values: %v", s)
		}
		preserveInlineStepOrder(s, v)
		bk.rawSteps = append(bk.rawSteps, v)
	}
	for _, r := range rb.HostRules {
		host, ok := r.Key.(string)
		if !ok {
			return nil, fmt.Errorf("invalid host rule: %v", r)
		}
		rule, ok := r.Value.(string)
		if !ok {
			return nil, fmt.Errorf("invalid host rule: %v", r)
		}
		bk.hostRules = append(bk.hostRules, hostRule{host, rule})
	}
	bk.debug = rb.Debug
	bk.intervalStr = rb.Interval
	bk.ifCond = rb.If
	bk.skipTest = rb.SkipTest
	bk.force = rb.Force
	bk.trace = rb.Trace
	if rb.Loop != nil {
		bk.loop, err = newLoop(rb.Loop)
		if err != nil {
			return nil, err
		}
	}
	if rb.Concurrency != nil {
		bk.concurrency, err = newConcurrency(rb.Concurrency)
		if err != nil {
			return nil, err
		}
	}
	bk.useMap = rb.useMap
	bk.stepKeys = rb.stepKeys

	return bk, nil
}

func newConcurrency(v any) ([]string, error) {
	switch vv := v.(type) {
	case string:
		return []string{vv}, nil
	case []string:
		return vv, nil
	case []any:
		var ss []string
		for _, vvv := range vv {
			s, ok := vvv.(string)
			if !ok {
				return nil, fmt.Errorf("invalid concurrency: %v", v)
			}
			ss = append(ss, s)
		}
		return ss, nil
	default:
		return nil, fmt.Errorf("invalid concurrency: %v", v)
	}
}

func joinCommands(in ...string) string {
	var cmd []string
	for _, i := range in {
		i = strings.TrimSuffix(i, "\n")
		if strings.Contains(i, " ") {
			cmd = append(cmd, fmt.Sprintf("%#v", i))
		} else {
			cmd = append(cmd, i)
		}
	}
	return strings.Join(cmd, " ") + "\n"
}

// preserveInlineStepOrder replaces the normalized map[string]any for inline
// include map-style steps with the original yaml.MapSlice so that
// parseInlineSteps can iterate in YAML definition order.
func preserveInlineStepOrder(original yaml.MapSlice, normalized map[string]any) {
	for _, item := range original {
		if fmt.Sprintf("%v", item.Key) != includeRunnerKey {
			continue
		}
		incSlice, ok := item.Value.(yaml.MapSlice)
		if !ok {
			return
		}
		for _, incItem := range incSlice {
			if fmt.Sprintf("%v", incItem.Key) != "steps" {
				continue
			}
			if _, ok := incItem.Value.(yaml.MapSlice); !ok {
				return
			}
			if incNorm, ok := normalized[includeRunnerKey].(map[string]any); ok {
				incNorm["steps"] = incItem.Value
			}
			return
		}
		return
	}
}

func normalize(v any) any {
	switch v := v.(type) {
	case []any:
		res := make([]any, len(v))
		for i, vv := range v {
			res[i] = normalize(vv)
		}
		return res
	case map[any]any:
		res := make(map[string]any, len(v))
		for k, vv := range v {
			res[fmt.Sprintf("%v", k)] = normalize(vv)
		}
		return res
	case map[string]any:
		res := make(map[string]any, len(v))
		for k, vv := range v {
			res[k] = normalize(vv)
		}
		return res
	case []map[string]any:
		res := make([]map[string]any, len(v))
		for i, vv := range v {
			var ok bool
			res[i], ok = normalize(vv).(map[string]any)
			if !ok {
				return fmt.Errorf("failed to normalize: %v", vv)
			}
		}
		return res
	case yaml.MapSlice:
		res := make(map[string]any, len(v))
		for _, i := range v {
			res[fmt.Sprintf("%v", i.Key)] = normalize(i.Value)
		}
		return res
	case int:
		if v < 0 {
			return int64(v)
		}
		return uint64(v)
	default:
		return v
	}
}

func detectRunbookAreas(in string) *areas {
	a := &areas{}
	tokens := lexer.Tokenize(in)
	parsed, err := parser.Parse(tokens, 0)
	if err != nil {
		return a
	}
	if len(parsed.Docs) == 0 {
		return a
	}
	m, ok := parsed.Docs[0].Body.(*ast.MappingNode)
	if !ok {
		return a
	}
	sections := m.Values
	for _, s := range sections {
		key, ok := s.Key.(*ast.StringNode)
		if !ok {
			return a
		}
		switch key.Value {
		case "desc":
			a.Desc = detectAreaFromNode(s)
		case "vars":
			a.Vars = detectAreaFromNode(s)
		case "runners":
			a.Runners = detectAreaFromNode(s)
		case "steps":
			switch steps := s.Value.(type) {
			case *ast.MappingValueNode:
				a.Steps = append(a.Steps, detectAreaFromNode(steps.Value))
			case *ast.MappingNode:
				for _, v := range steps.Values {
					a.Steps = append(a.Steps, detectAreaFromNode(v))
				}
			case *ast.SequenceNode:
				for _, v := range steps.Values {
					aa := detectAreaFromNode(v)
					// Get `-` token
					t := v.GetToken()
					for {
						if t.Value == "-" {
							aa.Start = &position{
								Line: t.Position.Line,
							}
							break
						}
						t = t.Prev
					}
					a.Steps = append(a.Steps, aa)
				}
			}
		}
	}

	return a
}

type areaDetector struct {
	start *token.Token
	end   *token.Token
}

// Visit implements ast.Visitor interface.
// It detects the start and end token of the area.
func (d *areaDetector) Visit(node ast.Node) ast.Visitor {
	if d.start == nil {
		d.start = node.GetToken()
	}
	if d.end == nil {
		d.end = node.GetToken()
	}
	if d.start.Position.Line > node.GetToken().Position.Line ||
		(d.start.Position.Line == node.GetToken().Position.Line && d.start.Position.Column > node.GetToken().Position.Column) {
		d.start = node.GetToken()
	}
	if d.end.Position.Line < node.GetToken().Position.Line ||
		(d.end.Position.Line == node.GetToken().Position.Line && d.end.Position.Column < node.GetToken().Position.Column) {
		d.end = node.GetToken()
	}
	return d
}

// detectAreaFromNode detects the start and end position of the area from the node.
func detectAreaFromNode(node ast.Node) *area {
	d := &areaDetector{}
	ast.Walk(d, node)
	a := &area{
		Start: &position{
			Line: d.start.Position.Line,
		},
		End: &position{
			Line: d.end.Position.Line,
		},
	}
	if (strings.Count(d.end.Value, "\n") - 1) > 0 {
		if d.end.Next != nil {
			a.End.Line += strings.Count(d.end.Value, "\n") - 1
		}
	}
	return a
}

func pickStepYAML(in string, idx int) (string, error) {
	a := detectRunbookAreas(in)
	if len(a.Steps)-1 < idx {
		return "", fmt.Errorf("step not found: %d", idx)
	}
	step := a.Steps[idx]
	start := step.Start.Line
	end := step.End.Line
	lines := strings.Split(in, "\n")
	if len(lines) < end {
		end = len(lines)
	}
	w := len(strconv.Itoa(end))
	var picked []string
	for i := start; i <= end; i++ {
		picked = append(picked, yellow(fmt.Sprintf("%s ", fmt.Sprintf(fmt.Sprintf("%%%dd", w), i)))+lines[i-1])
	}
	return strings.Join(picked, "\n"), nil
}