main.go

package main

import (
	"crypto/sha256"
	"encoding/json"
	"fmt"
	"go/format"
	"io/ioutil"
	"log"
	"os"
	"path"
	"sort"
	"strconv"
	"strings"
	"text/template"

	"github.com/getkin/kin-openapi/openapi3"
	"github.com/spf13/cobra"
	trie "github.com/v-pap/trie"
	"gitlab.ops.mist.io/mistio/openapi-cli-generator/shorthand"
	yaml "gopkg.in/yaml.v2"
)

//go:generate go-bindata ./templates/...

// OpenAPI Extensions
const (
	ExtAliases       = "x-cli-aliases"
	ExtDescription   = "x-cli-description"
	ExtIgnore        = "x-cli-ignore"
	ExtHidden        = "x-cli-hidden"
	ExtName          = "x-cli-name"
	ExtWaiters       = "x-cli-waiters"
	ExtTableDefault  = "x-cli-table-default"
	ExtTableExtended = "x-cli-table-extended"
	ExtTableOverride = "x-cli-table-override"
	ExtHelpSchema    = "x-cli-help-schema"
)

// Param describes an OpenAPI parameter (path, query, header, etc)
type Param struct {
	Name        string
	CLIName     string
	GoName      string
	Description string
	In          string
	Required    bool
	Type        string
	TypeNil     string
	Style       string
	Explode     bool
}

// Operation describes an OpenAPI operation (GET/POST/PUT/PATCH/DELETE)
type Operation struct {
	HandlerName     string
	GoName          string
	Use             string
	Aliases         []string
	Short           string
	Long            string
	Method          string
	CanHaveBody     bool
	HasBody         bool
	ReturnType      string
	Path            string
	Tags            []string
	AllParams       []*Param
	RequiredParams  []*Param
	OptionalParams  []*Param
	MediaType       string
	Examples        []string
	Hidden          bool
	NeedsResponse   bool
	Waiters         []*WaiterParams
	TableDefault    []string
	TableExtended   []string
	TableOverride   map[string]string
	OriginalVerb    string
	CommandVerb     string
	UseStripped     string
	IgnoreCommand   bool
	LinkedOperation *Operation
}

// Waiter describes a special command that blocks until a condition has been
// met, after which it exits.
type Waiter struct {
	CLIName     string
	GoName      string
	Use         string
	Aliases     []string
	Short       string
	Long        string
	Delay       int
	Attempts    int
	RetryCodes  []int
	OperationID string `json:"operationId"`
	Operation   *Operation
	Matchers    []*Matcher
	Updates     string
	After       map[string]map[string]string
}

// Matcher describes a condition to match for a waiter.
type Matcher struct {
	Select   string
	Test     string
	Expected json.RawMessage
	State    string
}

// WaiterParams links a waiter with param selector querires to perform wait
// operations after a command has run.
type WaiterParams struct {
	Waiter *Waiter
	Args   []string
	Params map[string]string
}

// Server describes an OpenAPI server endpoint
type Server struct {
	Description string
	URL         string
	// TODO: handle server parameters
}

// Imports describe optional imports based on features in use.
type Imports struct {
	Fmt      bool
	Strings  bool
	Time     bool
	Os       bool
	Terminal bool
	Json     bool
	JmesPath bool
}

// OpenAPI describes an API
type OpenAPI struct {
	Imports      Imports
	Name         string
	GoName       string
	PublicGoName string
	Title        string
	Description  string
	Servers      []*Server
	Operations   []*Operation
	Waiters      []*Waiter
	CommandVerbs map[string]bool
	HandlerNames map[string]*Operation
}

// ProcessAPI returns the API description to be used with the commands template
// for a loaded and dereferenced OpenAPI 3 document.
func ProcessAPI(shortName string, api *openapi3.Swagger) *OpenAPI {
	apiName := shortName
	if api.Info.Extensions[ExtName] != nil {
		apiName = extStr(api.Info.Extensions[ExtName])
	}

	apiDescription := api.Info.Description
	if api.Info.Extensions[ExtDescription] != nil {
		apiDescription = extStr(api.Info.Extensions[ExtDescription])
	}

	result := &OpenAPI{
		Name:         apiName,
		GoName:       toGoName(shortName, false),
		PublicGoName: toGoName(shortName, true),
		Title:        api.Info.Title,
		Description:  escapeString(apiDescription),
		CommandVerbs: make(map[string]bool),
		HandlerNames: make(map[string]*Operation),
	}

	for _, s := range api.Servers {
		result.Servers = append(result.Servers, &Server{
			Description: s.Description,
			URL:         s.URL,
		})
	}

	// Convenience map for operation ID -> operation
	operationMap := make(map[string]*Operation)

	verbSubcommandMap := make(map[string]trie.Trie)

	var keys []string
	for path := range api.Paths {
		keys = append(keys, path)
	}
	sort.Strings(keys)

	for _, path := range keys {
		item := api.Paths[path]

		if item.Extensions[ExtIgnore] != nil {
			// Ignore this path.
			continue
		}

		pathHidden := false
		if item.Extensions[ExtHidden] != nil {
			json.Unmarshal(item.Extensions[ExtHidden].(json.RawMessage), &pathHidden)
		}

		operationsMap := item.Operations()
		operationsMapKeys := make([]string, 0, len(operationsMap))
		for k := range operationsMap {
			operationsMapKeys = append(operationsMapKeys, k)
		}
		sort.Strings(operationsMapKeys)
		for _, method := range operationsMapKeys {
			operation := operationsMap[method]
			if operation.Extensions[ExtIgnore] != nil {
				// Ignore this operation.
				continue
			}

			name := operation.OperationID
			if operation.Extensions[ExtName] != nil {
				name = extStr(operation.Extensions[ExtName])
			}

			var aliases []string
			if operation.Extensions[ExtAliases] != nil {
				// We need to decode the raw extension value into our string slice.
				json.Unmarshal(operation.Extensions[ExtAliases].(json.RawMessage), &aliases)
			}

			var tableDefault []string
			if operation.Extensions[ExtTableDefault] != nil {
				// We need to decode the raw extension value into our string slice.
				json.Unmarshal(operation.Extensions[ExtTableDefault].(json.RawMessage), &tableDefault)
			}

			var tableExtended []string
			if operation.Extensions[ExtTableExtended] != nil {
				// We need to decode the raw extension value into our string slice.
				json.Unmarshal(operation.Extensions[ExtTableExtended].(json.RawMessage), &tableExtended)
			}

			params := getParams(item, method)

			paramsTypeMap := make(map[string]string)
			for _, param := range params {
				paramsTypeMap[param.Name] = param.Type
			}

			var tableOverrideList []string
			tableOverrideMap := make(map[string]string)
			if operation.Extensions[ExtTableOverride] != nil {
				// We need to decode the raw extension value into our string slice.
				json.Unmarshal(operation.Extensions[ExtTableOverride].(json.RawMessage), &tableOverrideList)
			}

			for _, tableOverrideParam := range tableOverrideList {
				tableOverrideMap[tableOverrideParam] = paramsTypeMap[tableOverrideParam]
			}

			requiredParams := getRequiredParams(params)
			optionalParams := getOptionalParams(params)
			short := operation.Summary
			if short == "" {
				short = name
			}

			use := usage(name, requiredParams)

			description := operation.Description
			if operation.Extensions[ExtDescription] != nil {
				description = extStr(operation.Extensions[ExtDescription])
			}

			hasBody := false
			if operation.RequestBody != nil {
				hasBody = true
			}

			reqMt, reqSchema, reqExamples := getRequestInfo(operation)

			var examples []string
			if len(reqExamples) > 0 {
				wroteHeader := false
				for _, ex := range reqExamples {
					if _, ok := ex.(string); !ok {
						// Not a string, so it's structured data. Let's marshal it to the
						// shorthand syntax if we can.
						if m, ok := ex.(map[string]interface{}); ok {
							ex = shorthand.Get(m)
							examples = append(examples, ex.(string))
							continue
						}

						b, _ := json.Marshal(ex)

						if !wroteHeader {
							description += "\n## Input Example\n\n"
							wroteHeader = true
						}

						description += "\n" + string(b) + "\n"
						continue
					}

					if !wroteHeader {
						description += "\n## Input Example\n\n"
						wroteHeader = true
					}

					description += "\n" + ex.(string) + "\n"
				}
			}

			showHelpSchema := false
			if operation.Extensions[ExtHelpSchema] != nil {
				json.Unmarshal(operation.Extensions[ExtHelpSchema].(json.RawMessage), &showHelpSchema)
			}

			if showHelpSchema && reqSchema != "" {
				description += "\n## Request Schema (" + reqMt + ")\n\n" + reqSchema
			}

			method := strings.Title(strings.ToLower(method))

			hidden := pathHidden
			if operation.Extensions[ExtHidden] != nil {
				json.Unmarshal(operation.Extensions[ExtHidden].(json.RawMessage), &hidden)
			}

			returnType := "interface{}"
		returnTypeLoop:
			for code, ref := range operation.Responses {
				if num, err := strconv.Atoi(code); err != nil || num < 200 || num >= 300 {
					// Skip invalid responses
					continue
				}

				if ref.Value != nil {
					for _, content := range ref.Value.Content {
						if _, ok := content.Example.(map[string]interface{}); ok {
							returnType = "map[string]interface{}"
							break returnTypeLoop
						}

						if content.Schema != nil && content.Schema.Value != nil {
							if content.Schema.Value.Type == "object" || len(content.Schema.Value.Properties) != 0 {
								returnType = "map[string]interface{}"
								break returnTypeLoop
							}
						}
					}
				}
			}

			sort.Strings(examples)

			o := &Operation{
				HandlerName:    slug(name),
				GoName:         toGoName(name, true),
				Use:            use,
				Aliases:        aliases,
				Short:          short,
				Long:           escapeString(description),
				Method:         method,
				CanHaveBody:    method == "Post" || method == "Put" || method == "Patch",
				HasBody:        hasBody,
				ReturnType:     returnType,
				Path:           path,
				AllParams:      params,
				Tags:           operation.Tags,
				RequiredParams: requiredParams,
				OptionalParams: optionalParams,
				MediaType:      reqMt,
				Examples:       examples,
				Hidden:         hidden,
				TableDefault:   tableDefault,
				TableExtended:  tableExtended,
				TableOverride:  tableOverrideMap,
				OriginalVerb:   extractVerb(slug(name)),
				CommandVerb:    convertVerb(extractVerb(slug(name))),
				UseStripped:    stripVerb(use),
			}

			if _, ok := verbSubcommandMap[o.CommandVerb]; !ok {
				verbSubcommandMap[o.CommandVerb] = trie.New()
			}
			verbSubcommandMap[o.CommandVerb].Insert(stripVerb(o.HandlerName))

			result.CommandVerbs[o.CommandVerb] = true

			operationMap[operation.OperationID] = o

			result.Operations = append(result.Operations, o)

			result.HandlerNames[o.HandlerName] = o

			for _, p := range params {
				if p.In == "path" {
					result.Imports.Strings = true
				}
			}

			for _, p := range optionalParams {
				if p.In == "query" || p.In == "header" {
					result.Imports.Fmt = true
				}
			}
		}
	}

	for handler, op := range result.HandlerNames {
		suffix, ok := verbSubcommandMap[op.CommandVerb].FindLongestUniqueSuffix(stripVerb(op.HandlerName))
		if ok {
			op.Aliases = calculateAliases(stripVerb(op.HandlerName), suffix)
		}
		if op.OriginalVerb != "get" {
			continue
		}
		getOp := result.HandlerNames["list-"+stripVerb(handler)+"s"]
		if getOp != nil {
			result.Imports.Json = true
			result.Imports.JmesPath = true
			op.IgnoreCommand = true
			getOp.LinkedOperation = op
			splitUse := strings.Split(op.UseStripped, " ")
			getOp.UseStripped += " [" + splitUse[len(splitUse)-1] + "]"
			duplicateParams := findDuplicateParams(getOp.OptionalParams, op.OptionalParams)
			for i := range getOp.OptionalParams {
				if !duplicateParams[convertToSHA256(*getOp.OptionalParams[i])] {
					if strings.TrimSpace(getOp.OptionalParams[i].Description) != "" {
						getOp.OptionalParams[i].Description += " "
					}
					getOp.OptionalParams[i].Description += "(Only for listings)"
				}
			}
			for i := range op.OptionalParams {
				if !duplicateParams[convertToSHA256(*op.OptionalParams[i])] {
					if strings.TrimSpace(op.OptionalParams[i].Description) != "" {
						op.OptionalParams[i].Description += " "
					}
					op.OptionalParams[i].Description += "(Only for single resource)"
					getOp.OptionalParams = append(getOp.OptionalParams, op.OptionalParams[i])
				}
			}
		}
	}

	if api.Extensions[ExtWaiters] != nil {
		var waiters map[string]*Waiter

		if err := json.Unmarshal(api.Extensions[ExtWaiters].(json.RawMessage), &waiters); err != nil {
			panic(err)
		}

		for name, waiter := range waiters {
			waiter.CLIName = slug(name)
			waiter.GoName = toGoName(name+"-waiter", true)
			waiter.Operation = operationMap[waiter.OperationID]
			waiter.Use = usage(name, waiter.Operation.RequiredParams)

			for _, matcher := range waiter.Matchers {
				if matcher.Test == "" {
					matcher.Test = "equal"
				}
				if matcher.State == "failure" {
					result.Imports.Os = true
					result.Imports.Terminal = true
				}
			}

			for operationID, waitOpParams := range waiter.After {
				operationAfter := operationMap[operationID]
				if operationAfter == nil {
					panic(fmt.Errorf("Unknown after operation %s", operationID))
				}

				operationWaiter := operationMap[waiter.Operation.HandlerName]
				if operationWaiter == nil {
					panic(fmt.Errorf("Unknown waiter operation %s", waiter.Operation.HandlerName))
				}

				var args []string
				for _, p := range operationWaiter.RequiredParams {
					selector := waitOpParams[p.Name]
					if selector == "" {
						panic(fmt.Errorf("Missing required parameter %s", p.Name))
					}
					delete(waitOpParams, p.Name)

					args = append(args, selector)

					result.Imports.Fmt = true
					operationAfter.NeedsResponse = true
				}

				// Transform from OpenAPI param names to CLI names
				wParams := make(map[string]string)
				for p, s := range waitOpParams {
					found := false
					for _, optional := range operationWaiter.OptionalParams {
						if optional.Name == p {
							wParams[optional.CLIName] = s
							found = true
							break
						}
					}
					if !found {
						panic(fmt.Errorf("Unknown parameter %s for waiter %s", p, name))
					}
				}

				operationAfter.Waiters = append(operationAfter.Waiters, &WaiterParams{
					Waiter: waiter,
					Args:   args,
					Params: wParams,
				})
			}

			result.Waiters = append(result.Waiters, waiter)
		}

		if len(waiters) > 0 {
			result.Imports.Time = true
		}
	}

	return result
}

func convertToSHA256(o interface{}) string {
	h := sha256.New()
	h.Write([]byte(fmt.Sprintf("%v", o)))
	return fmt.Sprintf("%x", h.Sum(nil))
}

func findDuplicateParams(params1 []*Param, params2 []*Param) map[string]bool {
	duplicateParamsSHA := make(map[string]bool)
	params1SHAs := make(map[string]bool)
	for _, p1 := range params1 {
		params1SHAs[convertToSHA256(*p1)] = true
	}
	for _, p2 := range params2 {
		if params1SHAs[convertToSHA256(*p2)] {
			duplicateParamsSHA[convertToSHA256(*p2)] = true
		}
	}
	return duplicateParamsSHA
}

func convertVerb(verb string) string {
	verbsMap := map[string]string{"list": "get"}
	if val, ok := verbsMap[verb]; ok {
		verb = val
	}
	return verb
}

func extractVerb(handlerName string) string {
	return strings.Split(handlerName, "-")[0]
}

func stripVerb(use string) string {
	subStrings := strings.SplitN(use, "-", 2)
	if len(subStrings) == 2 {
		if subStrings[0] == "list" {
			return strings.TrimSuffix(subStrings[1], "s")
		}
		return subStrings[1]
	}
	return use
}

func calculateAliases(command, suffix string) []string {
	if len(command) == 0 {
		return []string{}
	}
	prefix := strings.TrimSuffix(command, suffix)
	subString := ""
	aliases := []string{}
	for i := len(prefix); i < len(command)-1; i++ {
		subString += string(command[i])
		aliases = append(aliases, prefix+subString)
	}
	if len(aliases) > 0 && string(command[len(command)-1]) != "s" {
		aliases = append(aliases, command+"s")
	}
	return aliases
}

// extStr returns the string value of an OpenAPI extension stored as a JSON
// raw message.
func extStr(i interface{}) (decoded string) {
	if err := json.Unmarshal(i.(json.RawMessage), &decoded); err != nil {
		panic(err)
	}

	return
}

func toGoName(input string, public bool) string {
	transformed := strings.Replace(input, "-", " ", -1)
	transformed = strings.Replace(transformed, "_", " ", -1)
	transformed = strings.Title(transformed)
	transformed = strings.Replace(transformed, " ", "", -1)

	if !public {
		transformed = strings.ToLower(string(transformed[0])) + transformed[1:]
	}

	return transformed
}

func escapeString(value string) string {
	transformed := strings.Replace(value, "\n", "\\n", -1)
	transformed = strings.Replace(transformed, "\"", "\\\"", -1)
	return transformed
}

func slug(operationID string) string {
	transformed := strings.ToLower(operationID)
	transformed = strings.Replace(transformed, "_", "-", -1)
	transformed = strings.Replace(transformed, " ", "-", -1)
	return transformed
}

func usage(name string, requiredParams []*Param) string {
	usage := slug(name)

	for _, p := range requiredParams {
		usage += " " + strings.ToUpper(slug(p.Name))
	}

	return usage
}

func getParams(path *openapi3.PathItem, httpMethod string) []*Param {
	operation := path.Operations()[httpMethod]
	allParams := make([]*Param, 0, len(path.Parameters))

	var total openapi3.Parameters
	total = append(total, path.Parameters...)
	total = append(total, operation.Parameters...)

	for _, p := range total {
		if p.Value != nil && p.Value.Extensions["x-cli-ignore"] == nil {
			t := "string"
			tn := "\"\""
			if p.Value.Schema != nil && p.Value.Schema.Value != nil && p.Value.Schema.Value.Type != "" {
				switch p.Value.Schema.Value.Type {
				case "boolean":
					t = "bool"
					tn = "false"
				case "integer":
					t = "int64"
					tn = "0"
				case "number":
					t = "float64"
					tn = "0.0"
				}
			}

			cliName := slug(p.Value.Name)
			if p.Value.Extensions[ExtName] != nil {
				cliName = extStr(p.Value.Extensions[ExtName])
			}

			description := p.Value.Description
			if p.Value.Extensions[ExtDescription] != nil {
				description = extStr(p.Value.Extensions[ExtDescription])
			}

			allParams = append(allParams, &Param{
				Name:        p.Value.Name,
				CLIName:     cliName,
				GoName:      toGoName("param "+cliName, false),
				Description: description,
				In:          p.Value.In,
				Required:    p.Value.Required,
				Type:        t,
				TypeNil:     tn,
			})
		}
	}

	return allParams
}

func getRequiredParams(allParams []*Param) []*Param {
	required := make([]*Param, 0)

	for _, param := range allParams {
		if param.Required || param.In == "path" {
			required = append(required, param)
		}
	}

	return required
}

func getOptionalParams(allParams []*Param) []*Param {
	optional := make([]*Param, 0)

	for _, param := range allParams {
		if !param.Required && param.In != "path" {
			optional = append(optional, param)
		}
	}

	return optional
}

func getRequestInfo(op *openapi3.Operation) (string, string, []interface{}) {
	mts := make(map[string][]interface{})

	if op.RequestBody != nil && op.RequestBody.Value != nil {
		for mt, item := range op.RequestBody.Value.Content {
			var schema string
			var examples []interface{}

			if item.Schema != nil && item.Schema.Value != nil {
				// Let's make this a bit more concise. Since it has special JSON
				// marshalling functions, we do a dance to get it into plain JSON before
				// converting to YAML.
				data, err := json.Marshal(item.Schema.Value)
				if err != nil {
					continue
				}

				var unmarshalled interface{}
				json.Unmarshal(data, &unmarshalled)

				data, err = yaml.Marshal(unmarshalled)
				if err == nil {
					schema = string(data)
				}
			}

			if item.Example != nil {
				examples = append(examples, item.Example)
			} else if len(item.Examples) > 0 {
				for _, ex := range item.Examples {
					if ex.Value != nil {
						examples = append(examples, ex.Value.Value)
					}
				}
			} else if item.Schema.Value.Example != nil {
				examples = append(examples, item.Schema.Value.Example)
			}

			mts[mt] = []interface{}{schema, examples}
		}
	}

	// Prefer JSON.
	for mt, item := range mts {
		if strings.Contains(mt, "json") {
			return mt, item[0].(string), item[1].([]interface{})
		}
	}

	// Fall back to YAML next.
	for mt, item := range mts {
		if strings.Contains(mt, "yaml") {
			return mt, item[0].(string), item[1].([]interface{})
		}
	}

	// Last resort: return the first we find!
	for mt, item := range mts {
		return mt, item[0].(string), item[1].([]interface{})
	}

	return "", "", nil
}

func writeFormattedFile(filename string, data []byte) {
	formatted, errFormat := format.Source(data)
	if errFormat != nil {
		formatted = data
	}

	err := ioutil.WriteFile(filename, formatted, 0600)
	if errFormat != nil {
		panic(errFormat)
	} else if err != nil {
		panic(err)
	}
}

func initCmd(cmd *cobra.Command, args []string) {
	if _, err := os.Stat("main.go"); err == nil {
		fmt.Println("Refusing to overwrite existing main.go")
		return
	}

	data, _ := Asset("templates/main.tmpl")
	tmpl, err := template.New("cli").Parse(string(data))
	if err != nil {
		panic(err)
	}

	templateData := map[string]string{
		"Name":    args[0],
		"NameEnv": strings.Replace(strings.ToUpper(args[0]), "-", "_", -1),
	}

	var sb strings.Builder
	err = tmpl.Execute(&sb, templateData)
	if err != nil {
		panic(err)
	}

	writeFormattedFile("main.go", []byte(sb.String()))
}

func generate(cmd *cobra.Command, args []string) {
	data, err := ioutil.ReadFile(args[0])
	if err != nil {
		log.Fatal(err)
	}

	// Load the OpenAPI document.
	loader := openapi3.NewSwaggerLoader()
	var swagger *openapi3.Swagger
	swagger, err = loader.LoadSwaggerFromData(data)
	if err != nil {
		log.Fatal(err)
	}

	funcs := template.FuncMap{
		"escapeStr": escapeString,
		"slug":      slug,
		"title":     strings.Title,
	}

	data, _ = Asset("templates/commands.tmpl")
	tmpl, err := template.New("cli").Funcs(funcs).Parse(string(data))
	if err != nil {
		panic(err)
	}

	shortName := strings.TrimSuffix(path.Base(args[0]), ".yaml")

	templateData := ProcessAPI(shortName, swagger)

	var sb strings.Builder
	err = tmpl.Execute(&sb, templateData)
	if err != nil {
		panic(err)
	}

	writeFormattedFile(shortName+".go", []byte(sb.String()))
}

func main() {
	root := &cobra.Command{}

	root.AddCommand(&cobra.Command{
		Use:   "init <app-name>",
		Short: "Initialize and generate a `main.go` file for your project",
		Args:  cobra.ExactArgs(1),
		Run:   initCmd,
	})

	root.AddCommand(&cobra.Command{
		Use:   "generate <api-spec>",
		Short: "Generate a `commands.go` file from an OpenAPI spec",
		Args:  cobra.ExactArgs(1),
		Run:   generate,
	})

	root.Execute()
}