teleport/lib/utils/parse/parse.go

/*
Copyright 2017-2020 Gravitational, Inc.

Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at

    http://www.apache.org/licenses/LICENSE-2.0

Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
*/

package parse

import (
	"go/ast"
	"go/parser"
	"net/mail"
	"regexp"
	"strconv"
	"strings"
	"unicode"

	"github.com/gravitational/trace"
)

// Expression is an expression template
// that can interpolate to some variables
type Expression struct {
	// namespace is expression namespace,
	// e.g. internal.traits has a variable traits
	// in internal namespace
	namespace string
	// variable is a variable name, e.g. trait name,
	// e.g. internal.traits has variable name traits
	variable string
	// prefix is a prefix of the string
	prefix string
	// suffix is a suffix
	suffix string
	// transform is an optional transform function to call,
	// currently email.local is the only supported function
	transform func(in string) (string, error)
}

// EmailLocal returns local part of the email
func EmailLocal(in string) (string, error) {
	if in == "" {
		return "", trace.BadParameter("address is empty")
	}
	addr, err := mail.ParseAddress(in)
	if err != nil {
		return "", trace.BadParameter("failed to parse address %q: %q", in, err)
	}
	parts := strings.SplitN(addr.Address, "@", 2)
	if len(parts) != 2 {
		return "", trace.BadParameter("could not find local part in %q", addr.Address)
	}
	return parts[0], nil
}

// Namespace returns a variable namespace, e.g. external or internal
func (p *Expression) Namespace() string {
	return p.namespace
}

// Name returns variable name
func (p *Expression) Name() string {
	return p.variable
}

// Interpolate interpolates the variable adding prefix and suffix if present,
// returns trace.NotFound in case if the trait is not found, nil in case of
// success and BadParameter error otherwise
func (p *Expression) Interpolate(traits map[string][]string) ([]string, error) {
	values, ok := traits[p.variable]
	if !ok {
		return nil, trace.NotFound("variable is not found")
	}
	out := make([]string, len(values))
	for i := range values {
		val := values[i]
		var err error
		if p.transform != nil {
			val, err = p.transform(val)
			if err != nil {
				return nil, trace.Wrap(err)
			}
		}
		out[i] = p.prefix + val + p.suffix
	}
	return out, nil
}

var reVariable = regexp.MustCompile(
	// prefix is anyting that is not { or }
	`^(?P<prefix>[^}{]*)` +
		// variable is antything in brackets {{}} that is not { or }
		`{{(?P<expression>\s*[^}{]*\s*)}}` +
		// prefix is anyting that is not { or }
		`(?P<suffix>[^}{]*)$`,
)

// RoleVariable checks if the passed in string matches the variable pattern
// {{external.foo}} or {{internal.bar}}. If it does, it returns the variable
// prefix and the variable name. In the previous example this would be
// "external" or "internal" for the variable prefix and "foo" or "bar" for the
// variable name. If no variable pattern is found, trace.NotFound is returned.
func RoleVariable(variable string) (*Expression, error) {
	match := reVariable.FindStringSubmatch(variable)
	if len(match) == 0 {
		if strings.Index(variable, "{{") != -1 || strings.Index(variable, "}}") != -1 {
			return nil, trace.BadParameter(
				"%q is using template brackets '{{' or '}}', however expression does not parse, make sure the format is {{variable}}",
				variable)
		}
		return nil, trace.NotFound("no variable found in %q", variable)
	}

	prefix, variable, suffix := match[1], match[2], match[3]

	// parse and get the ast of the expression
	expr, err := parser.ParseExpr(variable)
	if err != nil {
		return nil, trace.NotFound("no variable found in %q: %v", variable, err)
	}

	// walk the ast tree and gather the variable parts
	result, err := walk(expr)
	if err != nil {
		return nil, trace.Wrap(err)
	}

	// the variable must have two parts the prefix and the variable name itself
	if len(result.parts) != 2 {
		return nil, trace.NotFound("no variable found: %v", variable)
	}

	return &Expression{
		prefix:    strings.TrimLeftFunc(prefix, unicode.IsSpace),
		namespace: result.parts[0],
		variable:  result.parts[1],
		suffix:    strings.TrimRightFunc(suffix, unicode.IsSpace),
		transform: result.transform,
	}, nil
}

const (
	// EmailNamespace is a function namespace for email functions
	EmailNamespace = "email"
	// EmailLocalFnName is a name for email.local function
	EmailLocalFnName = "local"
)

// TransformFn is an optional transform function
// that can take in string and replace it with another value
type TransformFn func(in string) (string, error)

type walkResult struct {
	parts     []string
	transform TransformFn
}

// walk will walk the ast tree and gather all the variable parts into a slice and return it.
func walk(node ast.Node) (*walkResult, error) {
	var result walkResult

	switch n := node.(type) {
	case *ast.CallExpr:
		switch call := n.Fun.(type) {
		case *ast.Ident:
			return nil, trace.BadParameter("function %v is not supported", call.Name)
		case *ast.SelectorExpr:
			// Selector expression looks like email.local(parameter)
			namespace, ok := call.X.(*ast.Ident)
			if !ok {
				return nil, trace.BadParameter("expected namespace, e.g. email.local, got %v", call.X)
			}
			// This is the part before the dot
			if namespace.Name != EmailNamespace {
				return nil, trace.BadParameter("unsupported namespace, e.g. email.local, got %v", call.X)
			}
			// This is a function name
			if call.Sel.Name != EmailLocalFnName {
				return nil, trace.BadParameter("unsupported function %v, supported functions are: email.local", call.Sel.Name)
			}
			// Because only one function is supported for now,
			// this makes sure that the function call has exactly one argument
			if len(n.Args) != 1 {
				return nil, trace.BadParameter("expected 1 argument for email.local got %v", len(n.Args))
			}
			result.transform = EmailLocal
			ret, err := walk(n.Args[0])
			if err != nil {
				return nil, trace.Wrap(err)
			}
			result.parts = ret.parts
			return &result, nil
		default:
			return nil, trace.BadParameter("unsupported function %T", n.Fun)
		}
	case *ast.IndexExpr:
		ret, err := walk(n.X)
		if err != nil {
			return nil, err
		}
		result.parts = append(result.parts, ret.parts...)
		ret, err = walk(n.Index)
		if err != nil {
			return nil, err
		}
		result.parts = append(result.parts, ret.parts...)
		return &result, nil
	case *ast.SelectorExpr:
		ret, err := walk(n.X)
		if err != nil {
			return nil, err
		}
		result.parts = append(result.parts, ret.parts...)

		ret, err = walk(n.Sel)
		if err != nil {
			return nil, err
		}
		result.parts = append(result.parts, ret.parts...)
		return &result, nil
	case *ast.Ident:
		return &walkResult{parts: []string{n.Name}}, nil
	case *ast.BasicLit:
		value, err := strconv.Unquote(n.Value)
		if err != nil {
			return nil, err
		}
		return &walkResult{parts: []string{value}}, nil
	default:
		return nil, trace.BadParameter("unknown node type: %T", n)
	}
}