// Copyright 2022 ADA Logics Ltd
//
// 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.
//

/*
This file is used to autogenerate a fuzzer for Vitess' API marshaling routines.
To use this file:
git clone https://github.com/vitessio/vitess
cd vitess/go/vt
grep -r ') Unmarshal' .>>/tmp/marshal_targets.txt
cd ../test/fuzzing/autogenerate
go run convert_grep_to_fuzzer.go
# The fuzzer has now been built.
# To instruct OSS-fuzz to build the autogenerated fuzzer:
mv api_marshal_fuzzer.go ../
compile_go_fuzzer vitess.io/vitess/go/test/fuzzing FuzzAPIMarshal api_marshal_fuzzer
*/

package main

import (
	"bufio"
	"fmt"
	"os"
	"strings"
)

// importPathShort is a map to convert import paths to import aliases
var importPathShort = map[string]string{
	"vitess.io/vitess/go/vt/proto/query":             "query",
	"vitess.io/vitess/go/vt/proto/tabletmanagerdata": "tabletmanagerdata",
	"vitess.io/vitess/go/vt/proto/binlogdata":        "binlogdata",
	"vitess.io/vitess/go/vt/proto/vtgate":            "vtgate",
	"vitess.io/vitess/go/vt/proto/tableacl":          "tableacl",
	"vitess.io/vitess/go/vt/proto/logutil":           "logutil",
	"vitess.io/vitess/go/vt/proto/vttest":            "vttest",
	"vitess.io/vitess/go/vt/proto/vttime":            "vttime",
	"vitess.io/vitess/go/vt/proto/workflow":          "workflow",
	"vitess.io/vitess/go/vt/proto/replicationdata":   "replicationdata",
	"vitess.io/vitess/go/vt/proto/vtctldata":         "vtctldata",
	"vitess.io/vitess/go/vt/proto/vtrpc":             "vtrpc",
	"vitess.io/vitess/go/vt/proto/vschema":           "vschema",
	"vitess.io/vitess/go/vt/proto/mysqlctl":          "mysqlctl",
	"vitess.io/vitess/go/vt/proto/vtadmin":           "vtadmin",
	"vitess.io/vitess/go/vt/proto/throttlerdata":     "throttlerdata",
	"vitess.io/vitess/go/vt/proto/topodata":          "topodata",
}

// pathToImportPath is a map to convert filepaths to import paths.
// the filepaths are available in the grep results.
var pathToImportPath = map[string]string{
	"./proto/query/query_vtproto.pb.go":                         "vitess.io/vitess/go/vt/proto/query",
	"./proto/tabletmanagerdata/tabletmanagerdata_vtproto.pb.go": "vitess.io/vitess/go/vt/proto/tabletmanagerdata",
	"./proto/binlogdata/binlogdata_vtproto.pb.go":               "vitess.io/vitess/go/vt/proto/binlogdata",
	"./proto/vtgate/vtgate_vtproto.pb.go":                       "vitess.io/vitess/go/vt/proto/vtgate",
	"./proto/tableacl/tableacl_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/tableacl",
	"./proto/logutil/logutil_vtproto.pb.go":                     "vitess.io/vitess/go/vt/proto/logutil",
	"./proto/vttest/vttest_vtproto.pb.go":                       "vitess.io/vitess/go/vt/proto/vttest",
	"./proto/vttime/vttime_vtproto.pb.go":                       "vitess.io/vitess/go/vt/proto/vttime",
	"./proto/workflow/workflow_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/workflow",
	"./proto/replicationdata/replicationdata_vtproto.pb.go":     "vitess.io/vitess/go/vt/proto/replicationdata",
	"./proto/vtctldata/vtctldata_vtproto.pb.go":                 "vitess.io/vitess/go/vt/proto/vtctldata",
	"./proto/vtrpc/vtrpc_vtproto.pb.go":                         "vitess.io/vitess/go/vt/proto/vtrpc",
	"./proto/vschema/vschema_vtproto.pb.go":                     "vitess.io/vitess/go/vt/proto/vschema",
	"./proto/mysqlctl/mysqlctl_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/mysqlctl",
	"./proto/vtadmin/vtadmin_vtproto.pb.go":                     "vitess.io/vitess/go/vt/proto/vtadmin",
	"./proto/throttlerdata/throttlerdata_vtproto.pb.go":         "vitess.io/vitess/go/vt/proto/throttlerdata",
	"./proto/topodata/topodata_vtproto.pb.go":                   "vitess.io/vitess/go/vt/proto/topodata",
}

// contains checks if a string is present in a string slice
func contains(s []string, e string) bool {
	for _, a := range s {
		if a == e {
			return true
		}
	}
	return false
}

/*
createFunctionCall creates a function call
This is used to create the calls to the internal harnesses:
switch funcOp {
case 0:
        _ = FuzzqueryTarget(data2) <- created by createFunctionCall
case 1:
        _ = FuzzqueryVTGateCallerID(data2) <- created by createFunctionCall
case 2:
        _ = FuzzqueryEventToken(data2) <- created by createFunctionCall
}
*/
func createFunctionCall(shortName, structName string) string {
	return fmt.Sprintf("\t_ = Fuzz%s%s(data2)\n", shortName, structName)
}

/*
createHarness creates an internal harness.
An example of a generated harness:
func FuzztabletmanagerdataTableDefinition(data []byte) error {
        f := fuzz.NewConsumer(data)
        s := &tabletmanagerdata.TableDefinition{}
        err := f.GenerateStruct(s)
        if err != nil {
                return err
        }
        b, err := s.MarshalVT()
        if err != nil {
                return err
        }
        s2 := &tabletmanagerdata.TableDefinition{}
        err = s2.UnmarshalVT(b)
        if err != nil {
                return err
        }
        newBytes, err := f.GetBytes()
        if err != nil {
                return err
        }
        s3 := &tabletmanagerdata.TableDefinition{}
        err = s3.UnmarshalVT(newBytes)
        return err
}
*/
func createHarness(shortName, structName string) string {
	var harnessString strings.Builder
	harnessString.WriteString(fmt.Sprintf("\n\nfunc Fuzz%s%s(data []byte) error {\n", shortName, structName))
	harnessString.WriteString("\tf := fuzz.NewConsumer(data)\n")
	harnessString.WriteString(fmt.Sprintf("\ts := &%s.%s{}\n", shortName, structName))
	harnessString.WriteString("\terr := f.GenerateStruct(s)\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString("\tb, err := s.MarshalVT()\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString(fmt.Sprintf("\ts2 := &%s.%s{}\n", shortName, structName))
	harnessString.WriteString("\terr = s2.UnmarshalVT(b)\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString("\tnewBytes, err := f.GetBytes()\n")
	harnessString.WriteString("\tif err != nil {\n")
	harnessString.WriteString("\t\treturn err\n")
	harnessString.WriteString("\t}\n")
	harnessString.WriteString(fmt.Sprintf("\ts3 := &%s.%s{}\n", shortName, structName))
	harnessString.WriteString("\terr = s3.UnmarshalVT(newBytes)\n")
	harnessString.WriteString("\treturn err\n")
	harnessString.WriteString("}\n")
	return harnessString.String()
}

// createMainFuzzer creates a file and writes all
// the accumulated sub-elements (imports, main harness,
// internal harnesses) of the fuzzer to it.
func createMainFuzzer(functionList, harnesses []string) {
	var mainFuzzer strings.Builder

	// package and imports
	mainFuzzer.WriteString("package fuzzing\n\n")
	mainFuzzer.WriteString("import (\n")
	for k, v := range importPathShort {
		mainFuzzer.WriteString(fmt.Sprintf("\t%s \"%s\"\n", v, k))
	}
	mainFuzzer.WriteString("\tfuzz \"github.com/AdaLogics/go-fuzz-headers\"\n")
	mainFuzzer.WriteString(")\n\n")

	// main entrypoint of the fuzzer
	mainFuzzer.WriteString("func FuzzAPIMarshal(data []byte) int {\n")
	mainFuzzer.WriteString("\tif len(data)<10{ return 0 }\n")

	// calls to each internal harness
	maxOps := len(functionList)
	mainFuzzer.WriteString(fmt.Sprintf("\tfuncOp := int(data[0])%%%d\n", maxOps))
	mainFuzzer.WriteString("\tdata2 := data[1:]\n")
	mainFuzzer.WriteString("\tswitch funcOp {\n")
	for i := 0; i < len(functionList); i++ {
		mainFuzzer.WriteString(fmt.Sprintf("\tcase %d:\n", i))
		mainFuzzer.WriteString(fmt.Sprintf("\t%s\n", functionList[i]))
	}
	mainFuzzer.WriteString("\t}\n")
	mainFuzzer.WriteString("\treturn 1\n")
	mainFuzzer.WriteString("}")

	// add all the internal harnesses
	for i := 0; i < len(harnesses); i++ {
		mainFuzzer.WriteString(harnesses[i])
	}

	// write the file
	fuzzFile, err := os.Create("api_marshal_fuzzer.go")
	if err != nil {
		panic(err)
	}
	defer fuzzFile.Close()
	fuzzFile.WriteString(mainFuzzer.String())
}

// getGrepData extracts the required data from the grep result
func getGrepData() ([]string, []string) {
	// paths contains all import paths
	paths := make([]string, 0)

	// functionList contains the calls to all the harnesses
	functionList := make([]string, 0)

	// harnesses contains all the harnesses
	harnesses := make([]string, 0)

	// extract data from the grep results:
	f, _ := os.Open("/tmp/marshal_targets.txt")
	defer f.Close()
	scanner := bufio.NewScanner(f)
	for scanner.Scan() {
		line := scanner.Text()
		if strings.Contains(line, "func (") {
			thePath := strings.Split(line, ":")[0]
			if !contains(paths, thePath) {
				paths = append(paths, thePath)
			}
			targetStructTemp := strings.Split(line, "func (m *")
			if len(targetStructTemp) == 1 {
				continue
			}

			// get name of target struct
			structName := strings.Split(targetStructTemp[1], ") ")[0]

			// get import path and short name
			importPath := pathToImportPath[thePath]
			shortName := importPathShort[importPath]
			//fmt.Println(createHarness(shortName, structName))

			// create harness
			harness := createHarness(shortName, structName)
			harnesses = append(harnesses, harness)

			// add harness to list of function calls
			functionList = append(functionList, createFunctionCall(shortName, structName))
		}
	}
	return functionList, harnesses
}

func main() {
	functionList, harnesses := getGrepData()
	createMainFuzzer(functionList, harnesses)
}