go.mondoo.com/cnquery@v0.0.0-20231005093811-59568235f6ea/providers-sdk/v1/lr/go.go

// Copyright (c) Mondoo, Inc.
// SPDX-License-Identifier: BUSL-1.1

package lr

import (
	"bytes"
	"errors"
	"fmt"
	"path/filepath"
	"regexp"
	"strconv"
	"strings"

	"go.mondoo.com/cnquery/types"
	"go.mondoo.com/cnquery/utils/multierr"
)

// Go produced go code for the LR file
func Go(packageName string, ast *LR, collector *Collector) (string, error) {
	o := goBuilder{
		collector:    collector,
		ast:          ast,
		name:         packageName,
		providerName: filepath.Base(ast.Options["provider"]),
		packsInUse:   map[string]struct{}{},
	}

	// should generally not happen and is primarily used for logging
	if o.providerName == "" {
		o.providerName = "provider"
	}

	o.goCreateResource(ast.Resources)
	o.goGetData(ast.Resources)
	o.goSetData(ast.Resources)

	for i := range ast.Resources {
		o.goResource(ast.Resources[i])
	}

	imports := ""
	for packName := range o.packsInUse {
		importPath, ok := ast.packPaths[packName]
		if !ok {
			return "", errors.New("cannot find import path for pack: " + packName)
		}

		imports += "\n\t" + strconv.Quote(importPath)
	}

	coreImports := ""
	if hasTimeImports(ast, &o) {
		coreImports = "\n\t\"time\""
	}

	header := fmt.Sprintf(goHeader, coreImports, imports)
	return header + o.data, o.errors.Deduplicate()
}

func hasTimeImports(ast *LR, b *goBuilder) bool {
	timeType := primitiveTypes["time"]
	for i := range ast.Resources {
		r := ast.Resources[i]
		for j := range r.Body.Fields {
			field := r.Body.Fields[j]
			if field.BasicField != nil && field.BasicField.Type.goType(b) == timeType {
				return true
			}
		}
	}
	return false
}

var reservedKeywords = map[string]struct{}{}

func init() {
	// Reserved keywords up to go 1.20
	words := `
	break     default      func    interface  select
	case      defer        go      map        struct
	chan      else         goto    package    switch
	const     fallthrough  if      range      type
	continue  for          import  return     var
	`
	re := regexp.MustCompile(`\S+`)
	terms := re.FindAllString(words, -1)
	for _, term := range terms {
		reservedKeywords[term] = struct{}{}
	}
}

func fieldCall(f string) string {
	if _, ok := reservedKeywords[f]; ok {
		return "compute_" + f
	}
	return f
}

type goBuilder struct {
	data         string
	collector    *Collector
	ast          *LR
	errors       multierr.Errors
	name         string
	providerName string
	packsInUse   map[string]struct{}
}

const goHeader = `// Copyright (c) Mondoo, Inc.
// SPDX-License-Identifier: BUSL-1.1

// Code generated by resources. DO NOT EDIT.

package resources

import (
	"errors"%s

	"go.mondoo.com/cnquery/llx"
	"go.mondoo.com/cnquery/providers-sdk/v1/plugin"
	"go.mondoo.com/cnquery/types"%s
)
`

func (b *goBuilder) goCreateResource(r []*Resource) {
	newCmds := make([]string, len(r))
	for i := range r {
		resource := r[i]
		iName := resource.interfaceName(b)

		var parseArgs string
		if b.collector.HasInit(iName) {
			parseArgs = "Init: init" + iName + ","
		} else {
			parseArgs = "// to override args, implement: init" + iName + "(runtime *plugin.Runtime, args map[string]*llx.RawData) (map[string]*llx.RawData, plugin.Resource, error)"
		}

		newCmds[i] = fmt.Sprintf("\"%s\": {\n\t\t\t%s\n\t\t\tCreate: create%s,\n\t\t},", resource.ID, parseArgs, iName)
	}

	b.data += `
var resourceFactories map[string]plugin.ResourceFactory

func init() {
	resourceFactories = map[string]plugin.ResourceFactory {
		` + strings.Join(newCmds, "\n\t\t") + `
	}
}

// NewResource is used by the runtime of this plugin to create new resources.
// Its arguments may be provided by users. This function is generally not
// used by initializing resources from recordings or from lists.
func NewResource(runtime *plugin.Runtime, name string, args map[string]*llx.RawData) (plugin.Resource, error) {
	f, ok := resourceFactories[name]
	if !ok {
		return nil, errors.New("cannot find resource " + name + " in this provider")
	}

	if f.Init != nil {
		cargs, res, err := f.Init(runtime, args)
		if err != nil {
			return res, err
		}

		if res != nil {
			id := name+"\x00"+res.MqlID()
			if x, ok := runtime.Resources.Get(id); ok {
				return x, nil
			}
			runtime.Resources.Set(id, res)
			return res, nil
		}

		args = cargs
	}

	res, err := f.Create(runtime, args)
	if err != nil {
		return nil, err
	}

	id := name+"\x00"+res.MqlID()
	if x, ok := runtime.Resources.Get(id); ok {
		return x, nil
	}

	runtime.Resources.Set(id, res)
	return res, nil
}

// CreateResource is used by the runtime of this plugin to create resources.
// Its arguments must be complete and pre-processed. This method is used
// for initializing resources from recordings or from lists.
func CreateResource(runtime *plugin.Runtime, name string, args map[string]*llx.RawData) (plugin.Resource, error) {
	f, ok := resourceFactories[name]
	if !ok {
		return nil, errors.New("cannot find resource " + name + " in this provider")
	}

	res, err := f.Create(runtime, args)
	if err != nil {
		return nil, err
	}

	id := name+"\x00"+res.MqlID()
	if x, ok := runtime.Resources.Get(id); ok {
		return x, nil
	}

	runtime.Resources.Set(id, res)
	return res, nil
}
`
}

func (b *goBuilder) goGetData(r []*Resource) {
	fields := []string{}
	for i := range r {
		resource := r[i]
		for j := range resource.Body.Fields {
			field := resource.Body.Fields[j]
			if field.Init != nil {
				continue
			}

			x := fmt.Sprintf(`"%s.%s": func(r plugin.Resource) *plugin.DataRes {
		return (r.(*%s).Get%s()).ToDataRes(%s)
	},`,
				resource.ID, field.BasicField.ID,
				resource.structName(b), field.BasicField.methodname(),
				field.BasicField.Type.mondooType(b),
			)
			fields = append(fields, x)
		}
	}

	b.data += `
var getDataFields = map[string]func(r plugin.Resource) *plugin.DataRes{
	` + strings.Join(fields, "\n\t") + `
}

func GetData(resource plugin.Resource, field string, args map[string]*llx.RawData) *plugin.DataRes {
	f, ok := getDataFields[resource.MqlName()+"."+field]
	if !ok {
		return &plugin.DataRes{Error: "cannot find '" + field + "' in resource '" + resource.MqlName() + "'"}
	}

	return f(resource)
}
`
}

func (b *goBuilder) goSetData(r []*Resource) {
	fields := []string{}
	for i := range r {
		resource := r[i]

		x := fmt.Sprintf(`"%s.%s": func(r plugin.Resource, v *llx.RawData) (ok bool) {
			r.(*%s).__id, ok = v.Value.(string)
			return
		},`,
			resource.ID, "__id",
			resource.structName(b),
		)
		fields = append(fields, x)

		for j := range resource.Body.Fields {
			field := resource.Body.Fields[j]
			if field.Init != nil {
				continue
			}

			x := fmt.Sprintf(`"%s.%s": func(r plugin.Resource, v *llx.RawData) (ok bool) {
		r.(*%s).%s, ok = plugin.RawToTValue[%s](v.Value, v.Error)
		return
	},`,
				resource.ID, field.BasicField.ID,
				resource.structName(b), field.BasicField.methodname(),
				field.BasicField.Type.goType(b),
			)
			fields = append(fields, x)
		}
	}

	b.data += fmt.Sprintf(`
var setDataFields = map[string]func(r plugin.Resource, v *llx.RawData) bool {
	%s
}

func SetData(resource plugin.Resource, field string, val *llx.RawData) error {
	f, ok := setDataFields[resource.MqlName() + "." + field]
	if !ok {
		return errors.New("[%s] cannot set '"+field+"' in resource '"+resource.MqlName()+"', field not found")
	}

	if ok := f(resource, val); !ok {
		return errors.New("[%s] cannot set '"+field+"' in resource '"+resource.MqlName()+"', type does not match")
	}
	return nil
}

func SetAllData(resource plugin.Resource, args map[string]*llx.RawData) error {
	var err error
	for k, v := range args {
		if err = SetData(resource, k, v); err != nil {
			return err
		}
	}
	return nil
}
`,
		strings.Join(fields, "\n\t"),
		b.providerName, b.providerName,
	)
}

func (b *goBuilder) goResource(r *Resource) {
	b.goStruct(r)
	b.goFactory(r)
	b.goFields(r)
}

func (b *goBuilder) goStruct(r *Resource) {
	internalStruct := r.structName(b) + "Internal"
	if !b.collector.HasStruct(internalStruct) {
		internalStruct = "// optional: if you define " + internalStruct + " it will be used here"
	}

	fields := []string{}
	for i := range r.Body.Fields {
		field := r.Body.Fields[i]
		if field.Init != nil {
			continue
		}
		fields = append(fields, field.BasicField.goName()+" plugin.TValue["+field.BasicField.Type.goType(b)+"]")
	}

	sFields := strings.Join(fields, "\n\t")
	if len(fields) != 0 {
		sFields = "\n\t" + sFields
	}

	b.data += fmt.Sprintf(`
// %s for the %s resource
type %s struct {
	MqlRuntime *plugin.Runtime
	__id string
	%s%s
}
`,
		r.structName(b), r.ID, r.structName(b),
		internalStruct,
		sFields,
	)
}

func (b *goBuilder) goFactory(r *Resource) {
	createName := "create" + r.interfaceName(b)
	structName := r.structName(b)

	var idCode string
	if b.collector.HasID(structName) {
		idCode = `if res.__id == "" {
	res.__id, err = res.id()
		if err != nil {
			return nil, err
		}
	}`
	} else {
		idCode = `// to override __id implement: id() (string, error)`
	}

	b.data += fmt.Sprintf(`
// %s creates a new instance of this resource
func %s(runtime *plugin.Runtime, args map[string]*llx.RawData) (plugin.Resource, error) {
	res := &%s{
		MqlRuntime: runtime,
	}

	err := SetAllData(res, args)
	if err != nil {
		return res, err
	}

	%s

	if runtime.HasRecording {
		args, err = runtime.ResourceFromRecording(%s, res.__id)
		if err != nil || args == nil {
			return res, err
		}
		return res, SetAllData(res, args)
	}

	return res, nil
}

func (c *%s) MqlName() string {
	return "%s"
}

func (c *%s) MqlID() string {
	return c.__id
}
`,
		createName, createName,
		structName,
		idCode,
		strconv.Quote(r.ID),
		structName, r.ID,
		structName,
	)
}

func (b *goBuilder) goFields(r *Resource) {
	for i := range r.Body.Fields {
		field := r.Body.Fields[i]
		if field.Init != nil {
			continue
		}

		b.goField(r, field)
	}
}

func (b *goBuilder) goStaticField(r *Resource, field *Field) {
	goName := field.BasicField.goName()
	b.data += fmt.Sprintf(`
func (c *%s) Get%s() *plugin.TValue[%s] {
	return &c.%s
}
`,
		r.structName(b), goName, field.BasicField.Type.goType(b),
		goName,
	)
}

func (b *goBuilder) goField(r *Resource, field *Field) {
	if field.BasicField.isStatic() {
		b.goStaticField(r, field)
		return
	}

	if field.BasicField.ID == "id" {
		b.errors.Add(errors.New("cannot create a dynamically computed `id` field, please turn it into a static field (ie: `id(..)` => `id`)"))
		return
	}

	goName := field.BasicField.goName()
	goType := field.BasicField.Type.goType(b)
	goZero := field.BasicField.Type.goZeroValue()

	argDefs := []string{}
	argCall := []string{}
	if field.BasicField.Args != nil {
		args := field.BasicField.Args.List
		for i := range args {
			arg := args[i]
			name := resource2goname(arg.Type, b)
			argDefs = append(argDefs, fmt.Sprintf(`varg%s := c.Get%s()
		if varg%s.Error != nil {
			return %s, varg%s.Error
		}

		`, name, name, name, goZero, name))
			argCall = append(argCall, "varg"+name+".Data")
		}
	}

	// resource types may be loaded from recordings
	var fromRecording string
	if field.BasicField.Type.containsResource(b) {
		fromRecording = fmt.Sprintf(`if c.MqlRuntime.HasRecording {
			d, err := c.MqlRuntime.FieldResourceFromRecording(%s, c.__id, %s)
			if err != nil {
				return %s, err
			}
			if d != nil {
				return d.Value.(%s), nil
			}
		}

		`,
			strconv.Quote(r.ID), strconv.Quote(field.BasicField.ID),
			goZero,
			goType,
		)
	}

	b.data += fmt.Sprintf(`
func (c *%s) Get%s() *plugin.TValue[%s] {
	return plugin.GetOrCompute[%s](&c.%s, func() (%s, error) {
		%s%sreturn c.%s(%s)
	})
}
`,
		r.structName(b), goName, goType,
		goType, goName, goType,
		fromRecording, strings.Join(argDefs, ""),
		fieldCall(field.BasicField.ID), strings.Join(argCall, ", "),
	)
}

// GO METHODS FOR AST

func (b *goBuilder) importName(symbol string) (string, bool) {
	parts := strings.SplitN(symbol, ".", 2)
	if len(parts) >= 2 {
		if ref, ok := b.ast.imports[parts[0]]; ok {
			if _, ok := ref[parts[1]]; ok {
				return parts[1], true
			}
		}
	}
	return "", false
}

func indent(s string, depth int) string {
	space := ""
	for i := 0; i < depth; i++ {
		space += "\t"
	}
	return space + strings.Replace(s, "\n", "\n"+space, -1)
}

func (r *Resource) structName(b *goBuilder) string {
	return "mql" + r.interfaceName(b)
}

var reMethodName = regexp.MustCompile("\\.[a-z]")

func capitalizeDot(in []byte) []byte {
	return bytes.ToUpper([]byte{in[1]})
}

func (r *Resource) interfaceName(b *goBuilder) string {
	return resource2goname(r.ID, b)
}

func resource2goname(s string, b *goBuilder) string {
	pack := strings.SplitN(s, ".", 2)
	var name string
	if pack[0] != s {
		resources, ok := b.ast.imports[pack[0]]
		if ok {
			if _, ok := resources[pack[1]]; !ok {
				b.errors.Add(errors.New("cannot find resource " + pack[1] + " in imported resource pack " + pack[0]))
			}
			name = pack[0] + "." + strings.Title(string(
				reMethodName.ReplaceAllFunc([]byte(pack[1]), capitalizeDot),
			))
			b.packsInUse[pack[0]] = struct{}{}
		}
	}
	if name == "" {
		name = strings.Title(string(
			reMethodName.ReplaceAllFunc([]byte(s), capitalizeDot),
		))
	}

	return name
}

func (b *ResourceDef) staticFields() []*BasicField {
	res := []*BasicField{}
	for _, f := range b.Fields {
		if f.BasicField != nil {
			if f.BasicField.isStatic() {
				res = append(res, f.BasicField)
			}
		}
	}
	return res
}

func (f *BasicField) goName() string {
	return strings.Title(f.ID)
}

func (f *BasicField) isStatic() bool {
	return f.Args == nil
}

func (f *BasicField) methodname() string {
	return strings.Title(f.ID)
}

// Retrieve the raw mondoo equivalent type, which can be looked up
// as a resource.
func (t *Type) Type(ast *LR) types.Type {
	if t.SimpleType != nil {
		return t.SimpleType.typeItems(ast)
	}
	if t.ListType != nil {
		return t.ListType.typeItems(ast)
	}
	if t.MapType != nil {
		return t.MapType.typeItems(ast)
	}
	return types.Any
}

func (t *MapType) typeItems(ast *LR) types.Type {
	return types.Map(t.Key.typeItems(ast), t.Value.Type(ast))
}

func (t *ListType) typeItems(ast *LR) types.Type {
	return types.Array(t.Type.Type(ast))
}

func (t *SimpleType) typeItems(ast *LR) types.Type {
	switch t.Type {
	case "bool":
		return types.Bool
	case "int":
		return types.Int
	case "float":
		return types.Float
	case "string":
		return types.String
	case "regex":
		return types.Regex
	case "time":
		return types.Time
	case "dict":
		return types.Dict
	default:
		return resourceType(t.Type, ast)
	}
}

// Try to build an MQL resource from the given name. It may or may not exist in
// a pack. If it doesn't exist at all
func resourceType(name string, ast *LR) types.Type {
	pack := strings.SplitN(name, ".", 2)
	if pack[0] != name {
		resources, ok := ast.imports[pack[0]]
		if ok {
			if _, ok := resources[pack[1]]; ok {
				return types.Resource(pack[1])
			}
		}
	}

	// TODO: look up resources in the current registry and notify if they are not found

	return types.Resource(name)
}

// Retrieve the mondoo equivalent of the type. This is a stringified type
// i.e. it can be compiled with the MQL imports
func (t *Type) mondooType(b *goBuilder) string {
	i := t.mondooTypeItems(b)
	if i == "" {
		return "NO_TYPE_DETECTED"
	}
	return i
}

func (t *Type) mondooTypeItems(b *goBuilder) string {
	if t.SimpleType != nil {
		return t.SimpleType.mondooTypeItems(b)
	}
	if t.ListType != nil {
		return t.ListType.mondooTypeItems(b)
	}
	if t.MapType != nil {
		return t.MapType.mondooTypeItems(b)
	}
	return ""
}

func (t *MapType) mondooTypeItems(b *goBuilder) string {
	return "types.Map(" + t.Key.mondooTypeItems(b) + ", " + t.Value.mondooTypeItems(b) + ")"
}

func (t *ListType) mondooTypeItems(b *goBuilder) string {
	return "types.Array(" + t.Type.mondooTypeItems(b) + ")"
}

func (t *SimpleType) mondooTypeItems(b *goBuilder) string {
	switch t.Type {
	case "bool":
		return "types.Bool"
	case "int":
		return "types.Int"
	case "float":
		return "types.Float"
	case "string":
		return "types.String"
	case "regex":
		return "types.Regex"
	case "time":
		return "types.Time"
	case "dict":
		return "types.Dict"
	default:
		if name, ok := b.importName(t.Type); ok {
			return "types.Resource(\"" + name + "\")"
		}
		return "types.Resource(\"" + t.Type + "\")"
	}

	// TODO: check that this type if a proper resource
	// panic("Cannot convert type '" + t.Type + "' to mondoo type")
}

func (t *Type) containsResource(b *goBuilder) bool {
	if t.ListType != nil {
		return t.ListType.Type.containsResource(b)
	}
	if t.MapType != nil {
		return t.MapType.Value.containsResource(b)
	}
	if t.SimpleType != nil {
		if _, ok := primitiveTypes[t.SimpleType.Type]; !ok {
			return true
		}
	}
	return false
}

// The go type is the golang-equivalent code type, i.e. the type of the
// actual objects that are being moved around.
func (t *Type) goType(b *goBuilder) string {
	if t.SimpleType != nil {
		return t.SimpleType.goType(b)
	}
	if t.ListType != nil {
		return t.ListType.goType()
	}
	if t.MapType != nil {
		return t.MapType.goType(b)
	}
	return "NO_TYPE_DETECTED"
}

func (t *MapType) goType(b *goBuilder) string {
	// limited to interface{} because we cannot cast as universally
	// between types yet
	return "map[" + t.Key.goType(b) + "]interface{}"
}

func (t *ListType) goType() string {
	// limited to []interface{} because we cannot cast as universally
	// between types yet
	return "[]interface{}"
}

var primitiveTypes = map[string]string{
	"string": "string",
	"bool":   "bool",
	"int":    "int64",
	"float":  "float64",
	"time":   "*time.Time",
	"regex":  "string",
	"dict":   "interface{}",
	"any":    "interface{}",
}

func (t *SimpleType) goType(b *goBuilder) string {
	pt, ok := primitiveTypes[t.Type]
	if ok {
		return pt
	}

	if _, ok := b.importName(t.Type); ok {
		return "plugin.Resource"
	}

	return "*mql" + resource2goname(t.Type, b)
}

func (t *Type) goZeroValue() string {
	if t.SimpleType != nil {
		return t.SimpleType.goZeroValue()
	}
	return "nil"
}

var primitiveZeros = map[string]string{
	"string": "\"\"",
	"bool":   "false",
	"int":    "0",
	"float":  "0.0",
	"time":   "nil",
	"dict":   "nil",
	"any":    "nil",
}

func (t *SimpleType) goZeroValue() string {
	pt, ok := primitiveZeros[t.Type]
	if ok {
		return pt
	}

	// TODO: check if the resource exists
	return "nil"
}