/
envoy.istio
/
talm
GigaCode
beta
ОбзорЦентр заботыВойти

talm

Форк
0
Ветки: 5Коммиты: 82Теги: 17
talm
/tools
/
import_commands.go 
291 строка · 7.9 Кб
1
package main
2


3
import (
4
	"bytes"
5
	"flag"
6
	"fmt"
7
	"go/ast"
8
	"go/format"
9
	"go/parser"
10
	"go/token"
11
	"io"
12
	"io/ioutil"
13
	"log"
14
	"net/http"
15
	"os"
16
	"path/filepath"
17
	"strings"
18
)
19


20
var talosVersion = flag.String("talos-version", "main", "the desired Talos version (branch or tag)")
21


22
func changePackageName(node *ast.File, newPackageName string) {
23
	node.Name = ast.NewIdent(newPackageName)
24
}
25


26
func addFieldToStruct(node *ast.File, varName, fieldType, fieldName string) {
27
	// Variable to track if the variable or type is found
28
	var found bool
29


30
	// Inspect nodes to find and modify the target struct declaration or its type
31
	ast.Inspect(node, func(n ast.Node) bool {
32
		switch decl := n.(type) {
33
		case *ast.GenDecl:
34
			if decl.Tok == token.TYPE {
35
				for _, spec := range decl.Specs {
36
					ts, ok := spec.(*ast.TypeSpec)
37
					if !ok || !strings.HasSuffix(ts.Name.Name, "Type") {
38
						continue
39
					}
40
					if typeName := strings.TrimSuffix(ts.Name.Name, "Type"); typeName != varName {
41
						continue
42
					}
43


44
					st, ok := ts.Type.(*ast.StructType)
45
					if !ok {
46
						continue
47
					}
48


49
					// Add field to the found struct type
50
					field := &ast.Field{
51
						Names: []*ast.Ident{ast.NewIdent(fieldName)},
52
						Type:  ast.NewIdent(fieldType),
53
					}
54
					st.Fields.List = append(st.Fields.List, field)
55
					found = true
56
					return false // stop searching, type found and updated
57
				}
58
			} else if decl.Tok == token.VAR {
59
				for _, spec := range decl.Specs {
60
					vs, ok := spec.(*ast.ValueSpec)
61
					if !ok || len(vs.Names) != 1 || vs.Names[0].Name != varName {
62
						continue
63
					}
64


65
					st, ok := vs.Type.(*ast.StructType)
66
					if !ok {
67
						continue
68
					}
69


70
					// Add field to the found struct variable
71
					field := &ast.Field{
72
						Names: []*ast.Ident{ast.NewIdent(fieldName)},
73
						Type:  ast.NewIdent(fieldType),
74
					}
75
					st.Fields.List = append(st.Fields.List, field)
76
					found = true
77
					return false // stop searching, variable found and updated
78
				}
79
			}
80
		}
81
		return true
82
	})
83


84
	// If the struct or type is not found, create a new variable struct
85
	if !found {
86
		newField := &ast.Field{
87
			Names: []*ast.Ident{ast.NewIdent(fieldName)},
88
			Type:  ast.NewIdent(fieldType),
89
		}
90
		newStruct := &ast.StructType{
91
			Fields: &ast.FieldList{
92
				List: []*ast.Field{newField},
93
			},
94
		}
95
		newSpec := &ast.ValueSpec{
96
			Names: []*ast.Ident{ast.NewIdent(varName)},
97
			Type:  newStruct,
98
		}
99
		newDecl := &ast.GenDecl{
100
			Tok:   token.VAR,
101
			Specs: []ast.Spec{newSpec},
102
		}
103


104
		node.Decls = append(node.Decls, newDecl)
105
		fmt.Println("New struct variable created:", varName)
106
	}
107
}
108


109
func prependStmtToInit(node *ast.File, cmdName string) {
110
	ast.Inspect(node, func(n ast.Node) bool {
111
		fn, ok := n.(*ast.FuncDecl)
112
		if ok && fn.Name.Name == "init" {
113
			stmt := &ast.ExprStmt{
114
				X: &ast.CallExpr{
115
					Fun: &ast.SelectorExpr{
116
						X: &ast.SelectorExpr{
117
							X:   ast.NewIdent(cmdName + "Cmd"),
118
							Sel: ast.NewIdent("Flags()"),
119
						},
120
						Sel: ast.NewIdent("StringSliceVarP"),
121
					},
122
					Args: []ast.Expr{
123
						&ast.UnaryExpr{
124
							Op: token.AND,
125
							X:  ast.NewIdent(cmdName + "CmdFlags.configFiles"),
126
						},
127
						ast.NewIdent(`"file"`),
128
						ast.NewIdent(`"f"`),
129
						ast.NewIdent("nil"),
130
						ast.NewIdent(`"specify config files or patches in a YAML file (can specify multiple)"`),
131
					},
132
				},
133
			}
134
			fn.Body.List = append([]ast.Stmt{stmt}, fn.Body.List...)
135
			return false
136
		}
137
		return true
138
	})
139
}
140


141
func insertInitCode(node *ast.File, cmdName, initCode string) {
142
	anonFuncCode := fmt.Sprintf(`func() { %s }`, initCode)
143


144
	initCodeExpr, err := parser.ParseExpr(anonFuncCode)
145
	if err != nil {
146
		log.Fatalf("Failed to parse init code: %v", err)
147
	}
148


149
	ast.Inspect(node, func(n ast.Node) bool {
150
		switch x := n.(type) {
151
		case *ast.FuncDecl:
152
			if x.Name.Name == "init" {
153
				if x.Body != nil {
154
					initFunc, ok := initCodeExpr.(*ast.FuncLit)
155
					if !ok {
156
						log.Fatalf("Failed to extract function body from init code expression")
157
					}
158


159
					x.Body.List = append(initFunc.Body.List, x.Body.List...)
160
				}
161
			}
162
		}
163
		return true
164
	})
165
}
166


167
func processFile(filename, cmdName string) {
168
	content, err := ioutil.ReadFile(filename)
169
	if err != nil {
170
		log.Fatalf("Failed to read the file: %v", err)
171
	}
172
	src := string(content)
173


174
	src = strings.ReplaceAll(src, "\"f\"", "\"F\"")
175
	src = strings.ReplaceAll(src, "github.com/siderolabs/talos/internal", "github.com/aenix-io/talm/internal")
176


177
	// Create a new set of tokens and parse the source code
178
	fset := token.NewFileSet()
179
	node, err := parser.ParseFile(fset, filename, src, parser.ParseComments)
180
	if err != nil {
181
		log.Fatalf("Failed to parse file: %v", err)
182
	}
183


184
	changePackageName(node, "commands")
185
	addFieldToStruct(node, cmdName+"CmdFlags", "[]string", "configFiles")
186


187
	initCode := fmt.Sprintf(`%sCmd.Flags().StringSliceVarP(&%sCmdFlags.configFiles, "file", "f", nil, "specify config files or patches in a YAML file (can specify multiple)")
188
	%sCmd.PreRunE = func(cmd *cobra.Command, args []string) error {
189
		nodesFromArgs := len(GlobalArgs.Nodes) > 0
190
		endpointsFromArgs := len(GlobalArgs.Endpoints) > 0
191
		for _, configFile := range %sCmdFlags.configFiles {
192
			if err :=	processModelineAndUpdateGlobals(configFile,	nodesFromArgs, endpointsFromArgs, false); err != nil {
193
				return err
194
			}
195
		}
196
		return nil
197
	}
198
	`, cmdName, cmdName, cmdName, cmdName)
199


200
	if cmdName == "etcd" {
201
		for _, subCmdName := range []string{"etcdAlarmCmd", "etcdDefragCmd", "etcdForfeitLeadershipCmd", "etcdLeaveCmd", "etcdMemberListCmd", "etcdMemberRemoveCmd", "etcdSnapshotCmd", "etcdStatusCmd"} {
202
			initCode = fmt.Sprintf("%s\n%s", initCode, fmt.Sprintf(`
203
	%s.Flags().StringSliceVarP(&etcdCmdFlags.configFiles, "file", "f", nil, "specify config files or patches in a YAML file (can specify multiple)")
204
	%s.PreRunE = etcdCmd.PreRunE
205
		`, subCmdName, subCmdName))
206
		}
207
	}
208
	if cmdName == "image" {
209
		for _, subCmdName := range []string{"imageDefaultCmd", "imageListCmd", "imagePullCmd"} {
210
			initCode = fmt.Sprintf("%s\n%s", initCode, fmt.Sprintf(`
211
	%s.Flags().StringSliceVarP(&etcdCmdFlags.configFiles, "file", "f", nil, "specify config files or patches in a YAML file (can specify multiple)")
212
	%s.PreRunE = etcdCmd.PreRunE
213
		`, subCmdName, subCmdName))
214
		}
215
	}
216


217
	if cmdName == "etcd" {
218
		for _, subCmdName := range []string{"etcdAlarmListCmd", "etcdAlarmDisarmCmd"} {
219
			initCode = fmt.Sprintf("%s\n%s", initCode, fmt.Sprintf(`
220
	%s.Flags().StringSliceVarP(&etcdCmdFlags.configFiles, "file", "f", nil, "specify config files or patches in a YAML file (can specify multiple)")
221
	%s.PreRunE = etcdCmd.PreRunE
222
		`, subCmdName, subCmdName))
223
		}
224
	}
225


226
	insertInitCode(node, cmdName, initCode)
227


228
	var buf bytes.Buffer
229
	comment := fmt.Sprintf("// Code generated by go run tools/import_commands.go --talos-version %s %s\n// DO NOT EDIT.\n\n", *talosVersion, cmdName)
230
	buf.WriteString(comment)
231


232
	if err := format.Node(&buf, fset, node); err != nil {
233
		log.Fatalf("Failed to format the AST: %v", err)
234
	}
235


236
	if err := ioutil.WriteFile(filename, buf.Bytes(), 0644); err != nil {
237
		log.Fatalf("Failed to write the modified file: %v", err)
238
	}
239


240
	log.Printf("File %s updated successfully.", filename)
241
}
242


243
func main() {
244
	flag.Parse()
245
	url := fmt.Sprintf("https://github.com/siderolabs/talos/raw/%s/cmd/talosctl/cmd/talos/", *talosVersion)
246


247
	args := flag.Args()
248
	if len(args) == 0 {
249
		fmt.Println("Please provide commands to import")
250
		return
251
	}
252


253
	for _, cmd := range args {
254
		srcName := cmd + ".go"
255
		dstName := "pkg/commands/imported_" + srcName
256


257
		err := downloadFile(srcName, dstName, url)
258
		if err != nil {
259
			log.Fatalf("Error downloading file: %v", err)
260
		}
261


262
		log.Printf("File %s succefully downloaded to %s", srcName, dstName)
263


264
		cmdName := strings.TrimSuffix(filepath.Base(srcName), ".go")
265
		if cmdName == "list" {
266
			cmdName = "ls"
267
		}
268
		processFile(dstName, cmdName)
269
	}
270
}
271


272
func downloadFile(srcName, dstName string, url string) error {
273
	resp, err := http.Get(url + "/" + srcName)
274
	if err != nil {
275
		return err
276
	}
277
	defer resp.Body.Close()
278


279
	file, err := os.Create(dstName)
280
	if err != nil {
281
		return err
282
	}
283
	defer file.Close()
284


285
	_, err = io.Copy(file, resp.Body)
286
	if err != nil {
287
		return err
288
	}
289


290
	return nil
291
}
292

Использование cookies

Мы используем файлы cookie в соответствии с Политикой конфиденциальности и Политикой использования cookies.

Нажимая кнопку «Принимаю», Вы даете АО «СберТех» согласие на обработку Ваших персональных данных в целях совершенствования нашего веб-сайта и Сервиса GitVerse, а также повышения удобства их использования.

Запретить использование cookies Вы можете самостоятельно в настройках Вашего браузера.