yakcmds

package
v1.3.4-alpha0723 Latest Latest
Warning

This package is not in the latest version of its module.

Go to latest
Published: Jul 23, 2024 License: AGPL-3.0 Imports: 91 Imported by: 0

Documentation

Overview

Package yakcmds @Author bcy2007 2024/3/12 16:06

Index

Constants

This section is empty.

Variables

View Source
var AICommands = []*cli.Command{
	{
		Name: "ai",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "type",
				Value: "chatglm",
			},
		},
		Action: func(c *cli.Context) error {
			var t string
			switch strings.ToLower(c.String("type")) {
			case "openai":
				t = "openai"
			case "chatglm":
			default:
				return utils.Error("unsupported type: " + c.String("type"))
			}
			_ = t
			return nil
		},
	},
}
View Source
var CVEUtilCommands = []*cli.Command{
	{
		Name:    "translating",
		Aliases: []string{"ai-desc", "desc"},
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "apikey",
				Usage: "API Key for AI",
			},
			cli.BoolFlag{
				Name: "no-critical",
			},
			cli.IntFlag{
				Name:  "concurrent",
				Value: 10,
			},
			cli.StringFlag{
				Name: "cve-database",
			},
			cli.BoolFlag{
				Name: "cwe",
			},
			cli.BoolFlag{
				Name: "chaosmaker-rules,chaosmaker",
			},
			cli.StringFlag{Name: "proxy", Usage: "Network Proxy", EnvVar: "http_proxy"},
			cli.StringFlag{Name: "ai", Usage: "Which AI Gateway? (openai/chatglm)", Value: "openai"},
			cli.Float64Flag{Name: "timeout", Usage: "timeout for seconds", Value: 60},
			cli.Float64Flag{Name: "total-timeout", Usage: "total timeout (useful in CI)"},
		},
		Usage:  "Translate CVE Models to Chinese, Supported in OPENAI",
		Hidden: true,
		Action: func(c *cli.Context) error {
			totalTimeout := c.Float64("total-timeout")
			if totalTimeout <= 0 {
			}

			if c.Bool("chaosmaker-rules") {
				rule.DecorateRules(c.String("ai"), c.Int("concurrent"), c.String("proxy"))
				return nil
			}

			if c.Bool("cwe") {
				return cve.TranslatingCWE(c.String("keyfile"), c.Int("concurrent"), c.String("cve-database"))
			}
			_ = consts.GetGormCVEDatabase()
			_ = consts.GetGormCVEDescriptionDatabase()
			return cve.Translating(
				c.String("ai"),
				c.Bool("no-critical"),
				c.Int("concurrent"),
				c.String("cve-database"),
				aispec.WithAPIKey(c.String("apikey")),
				aispec.WithProxy(c.String("proxy")),
				aispec.WithTimeout(c.Float64("timeout")),
			)
		},
	},
	{
		Name:  "build-cve-database",
		Usage: "Build CVE Database in SQLite",
		Flags: []cli.Flag{
			cli.BoolFlag{Name: "cwe"},
			cli.BoolFlag{Name: "cache"},
			cli.StringFlag{Name: "output,o"},
			cli.StringFlag{Name: "description-db"},
			cli.IntFlag{Name: "year"},
			cli.BoolFlag{Name: "no-gzip"},
		},
		Action: func(c *cli.Context) error {
			cvePath := filepath.Join(consts.GetDefaultYakitBaseTempDir(), "cve")
			os.MkdirAll(cvePath, 0o755)

			outputFile := c.String("output")
			if outputFile == "" {
				outputFile = consts.GetCVEDatabasePath()
			}
			outputDB, err := consts.CreateCVEDatabase(outputFile)
			if err != nil {
				return err
			}

			gzipHandler := func() error {
				if c.Bool("no-gzip") {
					return nil
				}
				log.Infof("start to zip... %v", outputFile)
				zipFile := outputFile + ".gzip"
				fp, err := os.OpenFile(zipFile, os.O_CREATE|os.O_RDWR, 0o644)
				if err != nil {
					return err
				}
				defer fp.Close()

				w := gzip.NewWriter(fp)
				srcFp, err := os.Open(outputFile)
				if err != nil {
					return err
				}
				io.Copy(w, srcFp)
				defer srcFp.Close()
				w.Flush()
				w.Close()
				return nil
			}

			descDBPath := c.String("description-db")
			log.Infof("description-db: %v", descDBPath)
			if descDBPath == "" {
				_, _ = consts.InitializeCVEDescriptionDatabase()
				descDBPath = consts.GetCVEDescriptionDatabasePath()
			}
			descDB, err := gorm.Open("sqlite3", descDBPath)
			if err != nil {
				log.Warnf("cannot found sqlite3 cve description: %v", err)
			}

			if c.Bool("cwe") {
				cveDB := outputDB

				if descDB != nil && descDB.HasTable("cwes") && cveDB != nil {
					log.Info("cve-description database is detected, merge cve db")
					if cveDB.HasTable("cwes") {
						if db := cveDB.DropTable("cwes"); db.Error != nil {
							log.Errorf("drop cwe table failed: %s", db.Error)
						}
					}
					log.Infof("start to migrate cwe for cvedb")
					cveDB.AutoMigrate(&cveresources.CWE{})
					for cwe := range cveresources.YieldCWEs(descDB.Model(&cveresources.CVE{}), context.Background()) {
						cveresources.CreateOrUpdateCWE(cveDB, cwe.IdStr, cwe)
					}
					return gzipHandler()
				}

				log.Info("start to download cwe")
				fp, err := cvequeryops.DownloadCWE()
				if err != nil {
					return err
				}
				log.Info("start to load cwes")
				cwes, err := cvequeryops.LoadCWE(fp)
				if err != nil {
					return err
				}
				log.Infof("total cwes: %v", len(cwes))
				db := cveDB
				db.AutoMigrate(&cveresources.CWE{})
				cvequeryops.SaveCWE(db, cwes)
				return gzipHandler()
			}

			wg := new(sync.WaitGroup)
			wg.Add(2)
			var downloadFailed bool
			go func() {
				defer wg.Done()
				log.Infof("start to save cve data from database: %v", cvePath)
				if !c.Bool("cache") {
					err := cvequeryops.DownLoad(cvePath)
					if err != nil {
						log.Error("download failed: %s, err")
						downloadFailed = true
						return
					}
				}
			}()
			go func() {
				defer wg.Done()

				log.Infof("using description database: %s", descDBPath)
				db, err := gorm.Open("sqlite3", descDBPath)
				if err != nil {
					log.Error("sqlite3 failed: %s", err)
					return
				}
				log.Info("start to handling cve description db")
				v := make(map[string]cveresources.CVEDesc)
				var count int
				for i := range cve.YieldCVEDescriptions(db, context.Background()) {
					count++

					v[i.CVE] = cveresources.CVEDesc{
						TitleZh:           i.ChineseTitle,
						Solution:          i.OpenAISolution,
						DescriptionMainZh: i.ChineseDescription,
					}
				}
				cveresources.RegisterDesc(v)
				log.Infof("register description finished! total: %v", count)
			}()

			wg.Wait()
			if downloadFailed {
				return utils.Error("download failed")
			}

			var years []int
			if ret := c.Int("year"); ret > 0 {
				years = append(years, ret)
			}
			cvequeryops.LoadCVE(cvePath, outputFile, years...)
			return gzipHandler()
		},
	},
	{
		Name: "cve-merge",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "desc-db", Value: consts.GetCVEDescriptionDatabasePath()},
			cli.StringFlag{Name: "db", Value: consts.GetCVEDatabasePath()},
		},
		Action: func(c *cli.Context) error {
			log.Info("start to cve description and origin database")
			desc, err := gorm.Open("sqlite3", c.String("desc-db"))
			if err != nil {
				return err
			}
			cvedb, err := gorm.Open("sqlite3", c.String("db"))
			if err != nil {
				return err
			}

			cvedb = cvedb.Where("title_zh is '' or title_zh is null")
			count := 0
			updateCount := 0
			log.Infof("start to merge cve info from %s", c.String("desc-db"))
			for ins := range cveresources.YieldCVEs(cvedb, context.Background()) {
				count++
				var descIns cve.CVEDescription
				if err := desc.Where("cve = ?", ins.CVE).First(&descIns).Error; err != nil {
					continue
				}
				if descIns.CVE == "" {
					continue
				}
				if descIns.ChineseTitle != "" {

					ins.TitleZh = descIns.ChineseTitle
					ins.DescriptionMainZh = descIns.ChineseDescription
					ins.Solution = descIns.OpenAISolution
					cvedb.Save(ins)
					log.Infof("update cve: %v %v", ins.CVE, ins.TitleZh)
					updateCount++
				}
			}
			_ = cvedb
			log.Info("count: ", count, "updated: ", updateCount)
			desc.Close()
			cvedb.Close()

			existedGzipCVE := c.String("db") + ".gzip"
			log.Infof("start gzip origin db: %v", existedGzipCVE)

			if utils.GetFirstExistedPath(existedGzipCVE) != "" {
				backup := existedGzipCVE + ".tmp.bak"
				os.RemoveAll(backup)
				err := os.Rename(existedGzipCVE, backup)
				if err != nil {
					return err
				}
			}

			cvefp, err := os.OpenFile(existedGzipCVE, os.O_CREATE|os.O_RDWR, 0o666)
			if err != nil {
				return err
			}
			w := gzip.NewWriter(cvefp)
			cveOrigin, err := os.Open(c.String("db"))
			if err != nil {
				return err
			}
			io.Copy(w, cveOrigin)
			w.Flush()
			w.Close()
			cvefp.Close()
			cveOrigin.Close()
			log.Infof("gzip cve finished: %v", existedGzipCVE)

			existedGzipCVE = c.String("desc-db") + ".gzip"
			log.Infof("start gzip description cve db: %v", existedGzipCVE)

			if utils.GetFirstExistedPath(existedGzipCVE) != "" {
				backup := existedGzipCVE + ".tmp.bak"
				os.RemoveAll(backup)
				err := os.Rename(existedGzipCVE, backup)
				if err != nil {
					return err
				}
			}
			descfp, err := os.OpenFile(existedGzipCVE, os.O_CREATE|os.O_RDWR, 0o666)
			if err != nil {
				return err
			}
			descGzipW := gzip.NewWriter(descfp)
			descorigin, err := os.Open(c.String("desc-db"))
			if err != nil {
				return err
			}
			io.Copy(descGzipW, descorigin)
			descGzipW.Flush()
			descGzipW.Close()
			descfp.Close()
			descorigin.Close()
			log.Infof("gzip cve desc finished: %v", existedGzipCVE)

			return nil
		},
	},
	{
		Name:  "cve-upload",
		Usage: "upload local cve to aliyun oss (gzip)",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "ak",
				Usage: "oss aliyun access key",
			},
			cli.StringFlag{
				Name: "sk", Usage: "oss aliyun secret key",
			},
			cli.StringFlag{
				Name: "endpoint", Usage: "endpoint for aliyun oss",
				Value: `oss-accelerate.aliyuncs.com`,
			},
			cli.StringFlag{
				Name:  "bucket",
				Usage: `aliyunoss bucket name`,
				Value: "cve-db",
			},
		},
		Action: func(c *cli.Context) error {
			client, err := oss.New(c.String("endpoint"), c.String("ak"), c.String("sk"))
			if err != nil {
				log.Errorf("oss new client failed: %s", err)
				return nil
			}
			bucket, err := client.Bucket("cve-db")
			if err != nil {
				log.Errorf("fetch bucket failed: %s", err)
				return nil
			}

			cvePath := consts.GetCVEDatabaseGzipPath()
			log.Infof("start to upload cve database: %v", cvePath)
			if cvePath == "" {
				return utils.Errorf("no path found for cve: %s", cvePath)
			}
			if utils.GetFirstExistedPath(cvePath) == "" {
				return utils.Errorf("no cve database found: %s", cvePath)
			}

			if fi, err := os.Stat(cvePath); err != nil {
				log.Errorf("stat cve failed: %s", err)
				return err
			} else {
				if fi.Size() < 10*1024*1024 {
					log.Errorf("cve file size is too small: %v", fi.Size())
					return nil
				}
			}
			if err := bucket.PutObjectFromFile("default-cve.db.gzip", cvePath); err != nil {
				log.Errorf("upload cve failed: %s", err)
				return err
			}

			cveDescPath := consts.GetCVEDescriptionDatabaseGzipPath()
			log.Infof("start to upload cve(translating description database: %s)", cveDescPath)
			if cveDescPath == "" {
				log.Errorf("cannot found cve database gzip path")
				return nil
			}
			if utils.GetFirstExistedPath(cveDescPath) == "" {
				return utils.Errorf("no cve database found: %s", cveDescPath)
			}

			if fi, err := os.Stat(cveDescPath); err != nil {
				log.Errorf("stat cve desc failed: %s", err)
				return err
			} else {
				if fi.Size() < 10*1024*1024 {
					log.Errorf("cve desc file size is too small: %v", fi.Size())
					return nil
				}
			}

			if err := bucket.PutObjectFromFile("default-cve-description.db.gzip", cveDescPath); err != nil {
				log.Errorf("upload cve desc failed: %s", err)
				return nil
			}
			return nil
		},
	},
	{
		Name: "cve-download",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "ak",
				Usage: "oss aliyun access key",
			},
			cli.StringFlag{
				Name: "sk", Usage: "oss aliyun secret key",
			},
			cli.StringFlag{
				Name: "endpoint", Usage: "endpoint for aliyun oss",
				Value: `oss-accelerate.aliyuncs.com`,
			},
			cli.StringFlag{
				Name:  "bucket",
				Usage: `aliyunoss bucket name`,
				Value: "cve-db",
			},
		},
		Action: func(c *cli.Context) error {
			client, err := oss.New(c.String("endpoint"), c.String("ak"), c.String("sk"))
			if err != nil {
				log.Errorf("oss new client failed: %s", err)
				return nil
			}
			bucket, err := client.Bucket("cve-db")
			if err != nil {
				log.Errorf("fetch bucket failed: %s", err)
				return nil
			}

			cvePath := consts.GetCVEDatabaseGzipPath()
			log.Infof("start to download cve database: %v", cvePath)
			if cvePath == "" {
				return utils.Errorf("no path found for cve: %s", cvePath)
			}
			if utils.GetFirstExistedPath(cvePath) != "" {
				bak := cvePath + ".bak"
				if err := os.RemoveAll(bak); err != nil {
					return err
				}
				err := os.Rename(cvePath, cvePath+".bak")
				if err != nil {
					return utils.Errorf("%v' s backup failed: %s", cvePath, err)
				}
			}
			if err := bucket.DownloadFile("default-cve.db.gzip", cvePath, 20*1024*1024); err != nil {
				log.Errorf("download cve failed: %s", err)
				return err
			}

			log.Infof("start to extract db from gzip: %v", cvePath)

			cvePathDB := consts.GetCVEDatabasePath()
			os.RemoveAll(cvePathDB)
			cveFile, err := os.OpenFile(cvePathDB, os.O_RDWR|os.O_CREATE, 0o666)
			if err != nil {
				return utils.Errorf("open file failed: %s", err)
			}
			defer cveFile.Close()
			gzipFile, err := os.Open(cvePath)
			if err != nil {
				return err
			}
			defer gzipFile.Close()
			r, err := gzip.NewReader(gzipFile)
			if err != nil {
				return utils.Errorf("gzip new reader failed: %s", err)
			}
			_, err = io.Copy(cveFile, r)
			if err != nil {
				return utils.Errorf("cve(db) copy failed: %s", err)
			}
			log.Infof("download gzip database finished: %v", cvePathDB)

			cveDescPath := consts.GetCVEDescriptionDatabaseGzipPath()
			log.Infof("start to handle cve(translating description database: %s)", cveDescPath)
			if cveDescPath == "" {
				log.Errorf("cannot found cve database gzip path")
				return nil
			}
			var newDescDB bool
			if utils.GetFirstExistedPath(cveDescPath) == "" {
				newDescDB = true
			}
			if !newDescDB {
				err := os.Rename(cveDescPath, cveDescPath+".bak")
				if err != nil {
					return utils.Errorf("%v' s backup failed: %s", cveDescPath, err)
				}
			}

			log.Infof("start to download bucket: %s", "default-cve-description.db.gzip")
			err = bucket.DownloadFile("default-cve-description.db.gzip", cveDescPath, 20*1024*1024)
			if err != nil {
				log.Errorf("download cve desc failed: %s", err)
				return nil
			}

			log.Infof("start to un-gzip: %v", cveDescPath)
			cveDescPathDB := consts.GetCVEDescriptionDatabasePath()
			os.RemoveAll(cveDescPathDB)
			cveDescFile, err := os.OpenFile(cveDescPathDB, os.O_RDWR|os.O_CREATE, 0o666)
			if err != nil {
				return utils.Errorf("open file failed: %s", err)
			}
			defer cveDescFile.Close()
			gzipDescFile, err := os.Open(cveDescPath)
			if err != nil {
				return err
			}
			defer gzipDescFile.Close()

			r, err = gzip.NewReader(gzipDescFile)
			if err != nil {
				return utils.Errorf("gzip new reader failed: %s", err)
			}
			_, err = io.Copy(cveDescFile, r)
			if err != nil {
				return utils.Errorf("cve(desc) copy failed: %s", err)
			}
			log.Infof("download gzip database finished: %v", cveDescPathDB)
			return nil
		},
	},
}
View Source
var ChaosMakerAIHelperCommand = cli.Command{}
View Source
var DistributionCommands = []*cli.Command{
	&scannode.DistYakCommand,
	{
		Name:   "mq",
		Usage:  "distributed by private amqp application protocol, execute yak via rabbitmq",
		Before: nil,
		After:  nil,
		Action: func(c *cli.Context) error {
			config := spec.LoadAMQPConfigFromCliContext(c)
			node, err := scannode.NewScanNode(c.String("id"), c.String("server-port"), config)
			if err != nil {
				return err
			}
			node.Run()
			return nil
		},
		Flags: spec.GetCliBasicConfig("scannode"),
	},
	{
		Name:  "tunnel",
		Usage: "Create Tunnel For CyberTunnel Service",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "server", Value: "cybertunnel.run:64333"},
			cli.IntFlag{Name: "local-port", Value: 53},
			cli.StringFlag{Name: "local-host", Value: "127.0.0.1"},
			cli.IntFlag{Name: "remote-port", Value: 53},
			cli.StringFlag{Name: "secret", Value: ""},
			cli.StringFlag{Name: "network,proto", Value: "tcp"},
		},
		Action: func(c *cli.Context) error {
			return cybertunnel.MirrorLocalPortToRemoteEx(
				c.String("network"),
				c.String("local-host"),
				c.Int("local-port"),
				c.Int("remote-port"),
				"test-cli",
				c.String("server"),
				c.String("secret"),
				context.Background(),
			)
		},
	},
	{
		Name:    "inspect-tuns",
		Usage:   "Inspect Registered Tunnels",
		Aliases: []string{"lst"},
		Flags: []cli.Flag{
			cli.StringFlag{Name: "server", Usage: "远程 Yak Bridge X 服务器", Value: "127.0.0.1:64333"},
			cli.StringFlag{Name: "secret", Usage: "远程 Yak Bridge X 服务器密码"},
			cli.StringFlag{Name: "secondary-password,x", Usage: "远程 Yak Bridge X 服务器的二级密码,避免别人查看注册管道"},
			cli.StringFlag{Name: "id", Usage: "指定 ID 查看 Tunnel 信息与认证"},
		},
		Action: func(c *cli.Context) error {
			ctx, client, _, err := cybertunnel.GetClient(context.Background(), c.String("server"), c.String("secret"))
			if err != nil {
				return err
			}

			showTunnel := func(tun *tpb.RegisterTunnelMeta) {
				withAuth, _ := client.GetRegisteredTunnelDescriptionByID(ctx, &tpb.GetRegisteredTunnelDescriptionByIDRequest{
					Id:                tun.GetId(),
					SecondaryPassword: c.String("secondary-password"),
				})
				fmt.Printf(`Tunnel: %v
	addr: %v
	note:
%v
	auth: 
%v
-----------------

`, tun.GetId(), utils.HostPort(tun.GetConnectHost(), tun.GetConnectPort()), tun.GetVerbose(), string(withAuth.GetAuth()))
			}

			id := c.String("id")
			if id != "" {
				rsp, err := client.GetRegisteredTunnelDescriptionByID(ctx, &tpb.GetRegisteredTunnelDescriptionByIDRequest{
					Id:                id,
					SecondaryPassword: c.String("secondary-password"),
				})
				if err != nil {
					return err
				}

				if len(rsp.GetAuth()) <= 0 {
					return utils.Errorf("cannot generate auth bytes for tun: %s", id)
				}

				showTunnel(rsp.GetInfo())
				println(string(rsp.GetAuth()))
				return nil
			}

			resp, err := client.GetAllRegisteredTunnel(ctx, &tpb.GetAllRegisteredTunnelRequest{
				SecondaryPassword: c.String("secondary-password"),
			})
			if err != nil {
				return err
			}
			for i := 0; i < len(resp.GetTunnels()); i++ {
				showTunnel(resp.Tunnels[i])
			}

			return nil
		},
	},
}
View Source
var DocCommands = []*cli.Command{
	{
		Name:  "doc",
		Usage: "Show Help Information for coding, document in YakLang",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "lib,extlib,l,t",
				Usage: "展示特定第三方扩展包的定义和帮助信息",
			},
			cli.StringFlag{
				Name:  "func,f",
				Usage: "展示特定第三方扩展包函数的定义",
			},
			cli.BoolFlag{
				Name:  "all-lib,all-libs,libs",
				Usage: "展示所有第三方包的帮助信息",
			},
		},
		Action: func(c *cli.Context) error {
			helper := doc.DefaultDocumentHelper

			if c.Bool("all-lib") {
				for _, libName := range helper.GetAllLibs() {
					helper.ShowLibHelpInfo(libName)
				}
				return nil
			}

			extLib := c.String("extlib")
			function := c.String("func")
			if extLib == "" && function != "" {
				extLib = "__GLOBAL__"
			}

			if extLib == "" {
				helper.ShowHelpInfo()
				return nil
			}

			if function != "" {
				if info := helper.LibFuncHelpInfo(extLib, function); info == "" {
					log.Errorf("palm script engine no such function in %s: %v", extLib, function)
					return nil
				} else {
					helper.ShowLibFuncHelpInfo(extLib, function)
				}
			} else {
				if info := helper.LibHelpInfo(extLib); info == "" {
					log.Errorf("palm script engine no such extlib: %v", extLib)
					return nil
				} else {
					helper.ShowLibHelpInfo(extLib)
				}
			}

			return nil
		},
	},

	{
		Name:  "gendoc",
		Usage: "Generate Basic Yaml Structure for YakLang",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "dir",
				Usage: "生成的文档路径",
				Value: "docs",
			},
		},
		Action: func(c *cli.Context) error {
			libs := yak.EngineToLibDocuments(yaklang.New())
			baseDir := filepath.Join(".", c.String("dir"))

			_ = os.MkdirAll(baseDir, 0o777)
			for _, lib := range libs {
				targetFile := filepath.Join(baseDir, fmt.Sprintf("%v.yakdoc.yaml", lib.Name))
				existed := yakdocument.LibDoc{}
				if utils.GetFirstExistedPath(targetFile) != "" {
					raw, _ := ioutil.ReadFile(targetFile)
					_ = yaml.Unmarshal(raw, &existed)
				}

				lib.Merge(&existed)
				raw, _ := yaml.Marshal(lib)
				_ = ioutil.WriteFile(targetFile, raw, os.ModePerm)
			}

			for _, s := range yakdocument.LibsToRelativeStructs(libs...) {
				targetFile := filepath.Join(baseDir, "structs", fmt.Sprintf("%v.struct.yakdoc.yaml", s.StructName))
				dir, _ := filepath.Split(targetFile)
				_ = os.MkdirAll(dir, 0o777)
				existed := yakdocument.StructDocForYamlMarshal{}
				if utils.GetFirstExistedPath(targetFile) != "" {
					raw, err := ioutil.ReadFile(targetFile)
					if err != nil {
						log.Errorf("cannot find file[%s]: %s", targetFile, err)
						continue
					}
					err = yaml.Unmarshal(raw, &existed)
					if err != nil {
						log.Errorf("unmarshal[%s] failed: %s", targetFile, err)
					}
				}

				if existed.StructName != "" {
					s.Merge(&existed)
				}
				raw, _ := yaml.Marshal(s)
				_ = ioutil.WriteFile(targetFile, raw, os.ModePerm)
			}
			return nil
		},
	},
	{
		Name:  "builddoc",
		Usage: "Build Markdown Documents for YakLang(From Structured Yaml Text)",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "from",
				Usage: "生成的文档源文件路径",
				Value: "docs",
			},
			cli.StringFlag{
				Name:  "to",
				Usage: "生成 Markdown 内容",
				Value: "build/yakapis/",
			},
			cli.StringFlag{
				Name:  "to-vscode-data,tovd",
				Value: "build/yaklang-completion.json",
			},
		},
		Action: func(c *cli.Context) error {
			libs := yak.EngineToLibDocuments(yaklang.New())
			baseDir := filepath.Join(".", c.String("from"))

			outputDir := filepath.Join(".", c.String("to"))
			_ = os.MkdirAll(outputDir, os.ModePerm)

			_ = os.MkdirAll(baseDir, os.ModePerm)
			for _, lib := range libs {
				targetFile := filepath.Join(baseDir, fmt.Sprintf("%v.yakdoc.yaml", lib.Name))
				existed := yakdocument.LibDoc{}
				if utils.GetFirstExistedPath(targetFile) != "" {
					raw, _ := ioutil.ReadFile(targetFile)
					_ = yaml.Unmarshal(raw, &existed)
				}

				lib.Merge(&existed)

				outputFileName := filepath.Join(outputDir, fmt.Sprintf("%v.md", strings.ReplaceAll(lib.Name, ".", "_")))
				_ = outputFileName

				results := lib.ToMarkdown()
				if results == "" {
					return utils.Errorf("markdown empty... for %v", lib.Name)
				}
				err := ioutil.WriteFile(outputFileName, []byte(results), os.ModePerm)
				if err != nil {
					return err
				}
			}

			completionJsonRaw, err := yakdocument.LibDocsToCompletionJson(libs...)
			if err != nil {
				return err
			}
			err = ioutil.WriteFile(c.String("to-vscode-data"), completionJsonRaw, os.ModePerm)
			if err != nil {
				return utils.Errorf("write vscode auto-completions json failed: %s", err)
			}
			return nil
		},
	},
}
View Source
var GitCommands = []*cli.Command{
	{
		Name: "git-ai-commit",
		Action: func(c *cli.Context) error {
			dir, err := os.Getwd()
			if err != nil {
				return err
			}
			repo, err := git.PlainOpen(dir)
			if err != nil {
				return err
			}

			worktree, err := repo.Worktree()
			if err != nil {
				return utils.Wrap(err, "plain worktree failed")
			}

			headRef, err := repo.Head()
			if err != nil {
				return err
			}
			headCommit, err := repo.CommitObject(headRef.Hash())
			if err != nil {
				return utils.Wrap(err, `wtBase, err := repo.CommitObject(ref.Hash())`)
			}

			stats, _ := worktree.Status()
			if len(stats) > 0 {

				_, err = worktree.Add(".")
				if err != nil {
					return utils.Wrap(err, "worktree.Add failed")
				}
				commit, err := worktree.Commit("Auto Commit (Wait for amend)", &git.CommitOptions{
					All:       true,
					Author:    &headCommit.Author,
					Committer: &headCommit.Committer,
					Parents:   []plumbing.Hash{headCommit.ParentHashes[0]},
				})
				if err != nil {
					return utils.Wrap(err, "worktree.Commit failed")
				}
				headCommit, err = repo.CommitObject(commit)
				if err != nil {
					return utils.Wrap(err, `auto-commit changes failed`)
				}
				headRef, err = repo.Head()
				log.Infof("use head ref after commit")
			}

			headTree, err := headCommit.Tree()
			if err != nil {
				return utils.Wrap(err, `baseTree, err := wtBase.Tree()`)
			}

			parentBase, err := headCommit.Parent(0)
			if err != nil {
				return utils.Wrap(err, `wtBase.Parent(0) failed`)
			}
			parentTree, err := parentBase.Tree()
			if err != nil {
				return utils.Wrap(err, `fetch parent commit tree`)
			}

			changes, err := parentTree.Diff(headTree)
			if err != nil {
				return err
			}

			result := bytes.NewBuffer(nil)
			result.WriteString(`以下是 Git Diff` + "\r\n\r\n")
			for _, change := range changes {
				patch, err := change.Patch()
				if err != nil {
					continue
				}
				var buf bytes.Buffer
				buf.WriteString(change.String())
				buf.WriteByte('\n')
				raw := patch.String()
				if len(raw) > 409600 {
					raw = utils.ShrinkString(raw, 409600)
				}
				buf.WriteString(raw)
				fmt.Println(buf.String())
				io.Copy(result, &buf)
			}

			var finalMessage string
		GENERATE:
			for {
				log.Info("AI start to generate commit message(current)...")
				results, err := ai.FunctionCall(result.String(), map[string]any{
					"message":    "英文:适用于上述 git diff 的 commit message浓缩成一句话,标注 feature/bugfix/doc 之类的",
					"message_zh": `中文:要求含义同 message`,
				}, aispec.WithDebugStream(), aispec.WithType("openai"))
				if err != nil {
					log.Warnf("AI failed: %v", err)
					_, err := fmt.Scanf("start to retry? (enter to continue)")
					if err != nil {
						log.Warnf("retry failed: %v", err)
					}
					continue
				}

				commitMessage, ok := results["message"]
				if !ok {
					log.Infof("cannot generate ai commit message: ")
					_, err := fmt.Scanf("start to retry? (enter to continue)")
					if err != nil {
						log.Warnf("retry failed: %v", err)
					}
					continue
				}

				log.Infof("AI commit message: %v", commitMessage)
				commitMessageZh, _ := results["message_zh"]
				fmt.Println(commitMessageZh)
				// yes or retry?
				var result string
				fmt.Println()
				fmt.Println(`start to commit? (yes/y/t to commit, retry/r to retry, enter to continue): `)
				_, err = fmt.Scanf("%s", &result)
				if err != nil {
					return err
				}
				switch strings.TrimSpace(strings.ToLower(result)) {
				case "retry", "r":
					continue
				case "yes", "y", "t":
					finalMessage = codec.AnyToString(commitMessage)
					break GENERATE
				}
			}
			fmt.Println("start to commit...")
			commit, err := worktree.Commit(finalMessage, &git.CommitOptions{
				All:       true,
				Author:    &headCommit.Author,
				Committer: &headCommit.Committer,
				Parents:   []plumbing.Hash{headCommit.ParentHashes[0]},
			})
			if err != nil {
				return utils.Wrapf(err, "amend commit failed: %v", finalMessage)
			}
			_, err = repo.CommitObject(commit)
			if err != nil {
				return utils.Wrap(err, `auto-commit changes failed`)
			}
			return nil
		},
	},
	{
		Name: "git-diff",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "repository,repo,r", Usage: `Setting Repository, default PWD`},
			cli.StringFlag{Name: "target", Usage: "Target Commit (Tag or Hash)"},
			cli.StringFlag{Name: "base", Usage: `Git Base Commit (Hash/Tag/Branch), default HEAD`, Value: "HEAD"},
			cli.StringFlag{
				Name:  "exclude-filenames",
				Usage: "ignore some filenames like protobuf or some...",
				Value: "*.pb.go,*.png,*.gif,*.jpg,*.jpeg," +
					"*_test.go,*.sum,bindata.*," +
					"*_lexer.go,*_parser.go,*.interp,*.tokens," +
					"*.yakdoc.yaml,embed.go,*.min.js,*.min.css"},

			cli.BoolFlag{Name: "debug", Usage: "Debug will show more log info"},
		},
		Action: func(c *cli.Context) error {
			debug := c.Bool("debug")

			var repoPath string = "."
			if ret := c.String("repository"); ret != "" {
				repoPath = ret
			}

			if debug {
				log.Infof("start to open plain repo: %v", repoPath)
			}
			repo, err := git.PlainOpen(repoPath)
			if err != nil {
				return err
			}

			var base = c.String("base")
			if base == "" || base == "HEAD" {
				if debug {
					log.Info("base: start to check HEAD hash...")
				}
				ref, err := repo.Head()
				if err != nil {
					return err
				}
				base = ref.Hash().String()
				if debug {
					log.Infof("found HEAD hash: %v", base)
				}
			}

			var target = c.String("target")
			if target == "" || target == "HEAD" {
				if debug {
					log.Info("target: start to check HEAD hash...")
				}
				ref, err := repo.Head()
				if err != nil {
					return err
				}
				target = ref.Hash().String()
				if debug {
					log.Infof("found HEAD hash: %v", base)
				}
			}

			excludeFiles := utils.PrettifyListFromStringSplitEx(c.String("exclude-filenames"))
			shouldFocus := func(pathName string) bool {
				_, filename := filepath.Split(pathName)
				if utils.MatchAnyOfGlob(filename, excludeFiles...) {
					return false
				}
				return true
			}

			if base == target {
				return utils.Errorf("base and target is the same hash(tag): %v", base)
			}

			err = yakdiff.GitHashDiffContext(
				context.Background(), repo,
				base, target,
				func(commit *object.Commit, change *object.Change, patch *object.Patch) error {
					if patch == nil {
						fmt.Println(patch.String())
						return nil
					}

					action, err := change.Action()
					if err != nil {
						log.Warnf("Change Action fetch failed: %s", err)
						return nil
					}

					fmt.Println(`---------------------------------------------------------------------------`)

					stats := patch.Stats()
					patches := patch.FilePatches()

					if len(patches) <= 0 {
						fmt.Println(patch.String())
						return nil
					}

					showStats := true
					if len(stats) != len(patches) {
						showStats = false
					}
					for idx, fp := range patches {
						fromFile, toFile := fp.Files()
						if fromFile == nil && toFile == nil {
							fmt.Println(fp)
							continue
						}

						var filename string
						if toFile == nil {
							filename = fromFile.Path()
						} else {
							filename = toFile.Path()
						}

						if fp.IsBinary() {
							if debug {
								log.Infof("skip binary change: %v", filename)
							}
							continue
						}

						if !shouldFocus(filename) {
							if debug {
								log.Infof("ignore file: %s by user config", filename)
							}
							continue
						}

						if showStats {
							fmt.Printf("%v: | %v\n", action, strings.TrimSpace(stats[idx].String()))
						} else {
							fmt.Printf("%v: | %v\n", action, filename)
						}

						chunks := fp.Chunks()
						if len(chunks) <= 0 {
							continue
						}

						dstFile, err := os.ReadFile(filename)
						if err != nil {
							continue
						}
						editor := memedit.NewMemEditor(string(dstFile))

						editorOffset := 0
						for _, chunked := range chunks {
							verbose := chunked.Content()
							verbose = utils.ShrinkString(verbose, 64)

							var action string = " "
							var suffix string
							switch chunked.Type() {
							case diff.Equal:
								action = " "
								editor.FindStringRangeIndexFirst(editorOffset, chunked.Content(), func(rangeIf memedit.RangeIf) {
									editorOffset = editor.GetOffsetByPosition(rangeIf.GetEnd())
								})
							case diff.Add:
								action = "+"
								editor.FindStringRangeIndexFirst(editorOffset, chunked.Content(), func(rangeIf memedit.RangeIf) {
									editorOffset = editor.GetOffsetByPosition(rangeIf.GetEnd())
									suffix = ` (` + fmt.Sprintf(
										"%v:%v-%v:%v",
										rangeIf.GetStart().GetLine(), rangeIf.GetStart().GetColumn(),
										rangeIf.GetEnd().GetLine(), rangeIf.GetEnd().GetColumn(),
									) + `)`
								})
							case diff.Delete:
								action = "-"
							}
							fmt.Printf("  |%v%-67s |%v\n", action, verbose, suffix)
						}
					}

					return nil
				},
			)
			if err != nil {
				return err
			}
			return nil
		},
	},
}
View Source
var JavaUtils = []*cli.Command{
	{
		Name:    "serialdumper",
		Usage:   "Java SerialDumper in Yaklang/Golang Implemented",
		Aliases: []string{"sd"},
		Action: func(c *cli.Context) {
			if len(c.Args()) > 0 {
				raw, err := codec.DecodeHex(c.Args()[0])
				if err != nil {
					log.Error(err)
					return
				}
				d := yserx.JavaSerializedDumper(raw)
				println(d)
			}
		},
	},
}
View Source
var PassiveCommands = cli.Command{
	Name:      "passive-scan",
	ShortName: "passive",
	Aliases: []string{
		"webscan",
	},
	Usage:       "yak passive-scan [options]",
	Description: "Passive Proxy(MITM) Scan.",
	Flags: []cli.Flag{
		cli.StringFlag{
			Name:  "listen",
			Value: "0.0.0.0:8084",
			Usage: "MITM on which addr:port?",
		},
	},
	Action: func(c *cli.Context) error {

		return nil
	},
}
View Source
var ProjectCommands = []*cli.Command{
	{
		Name:  "profile-export",
		Usage: "Export Yakit Profile Database to File",
		Action: func(c *cli.Context) {
			f := c.String("output")
			if utils.GetFirstExistedPath(f) != "" {
				log.Errorf("path[%s] is existed", f)
				return
			}

			if c.String("type") == "" {
				log.Error("export type cannot be emtpy")
				return
			}
			switch ret := strings.ToLower(c.String("type")); ret {
			case "plugin", "plugins":
				err := yakit.ExportYakScript(consts.GetGormProfileDatabase(), f)
				if err != nil {
					log.Error("output failed: %s", err)
				}
			default:
				log.Error("unsupported resource type: " + ret)
				return
			}
		}, Flags: []cli.Flag{
			cli.StringFlag{Name: "output"},
			cli.StringFlag{Name: "type"},
		}},
}
View Source
var SSACompilerCommands = []*cli.Command{
	{
		Name:    "ssa-remove",
		Aliases: []string{"ssa-rm"},
		Usage:   "Remove SSA OpCodes from database",
		Action: func(c *cli.Context) {
			for _, name := range c.Args() {
				if name == "*" {
					for _, i := range ssadb.AllPrograms(ssadb.GetDB()) {
						log.Infof("Start to delete program: %v", i)
						ssadb.DeleteProgram(ssadb.GetDB(), i)
					}
					break
				}
				log.Infof("Start to delete program: %v", name)
				ssadb.DeleteProgram(ssadb.GetDB(), name)
			}
		},
	},
	{
		Name:    "ssa-compile",
		Aliases: []string{"ssa"},
		Usage:   "Compile to SSA OpCodes from source code",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "log", Usage: "log level"},
			cli.StringFlag{
				Name: "language,l",
			},
			cli.StringFlag{
				Name:  "target,t",
				Usage: `target file or directory`,
			},
			cli.StringFlag{
				Name:  "program,p",
				Usage: `program name to save in database`,
			},
			cli.StringFlag{
				Name:  "entry",
				Usage: "Program Entry",
			},
			cli.BoolFlag{
				Name: "memory",
			},
			cli.StringFlag{
				Name:  "syntaxflow,sf",
				Usage: "syntax flow query language",
			},
			cli.StringFlag{
				Name:  "database,db",
				Usage: "database path",
			},
			cli.BoolFlag{
				Name:  "database-debug,dbdebug",
				Usage: "enable database debug mode",
			},
			cli.BoolFlag{
				Name:  "syntaxflow-debug,sfdebug",
				Usage: "enable syntax flow debug mode",
			},
			cli.BoolFlag{
				Name: "no-override", Usage: "no override existed database program(no delete)",
			},
			cli.BoolFlag{
				Name: "re-compile", Usage: "re-compile existed database program",
			},
			cli.BoolFlag{
				Name: "dot", Usage: "dot graph text for result",
			},
			cli.BoolFlag{
				Name: "with-code,code", Usage: "show code context",
			},
		},
		Action: func(c *cli.Context) error {
			if ret, err := log.ParseLevel(c.String("log")); err == nil {
				log.SetLevel(ret)
			}

			programName := c.String("program")
			if programName != "" {
				defer func() {
					ssa.ShowDatabaseCacheCost()
				}()
			}
			entry := c.String("entry")
			input_language := c.String("language")
			inMemory := c.Bool("memory")
			rawFile := c.String("target")
			target := utils.GetFirstExistedPath(rawFile)
			databaseFileRaw := c.String("database")
			noOverride := c.Bool("no-override")
			syntaxFlow := c.String("syntaxflow")
			dbDebug := c.Bool("database-debug")
			sfDebug := c.Bool("syntaxflow-debug")
			showDot := c.Bool("dot")
			withCode := c.Bool("with-code")

			if databaseFileRaw != "" {

				if target == "" &&
					utils.GetFirstExistedFile(databaseFileRaw) == "" {

					return utils.Errorf("database file not found: %v", databaseFileRaw)
				}
			}
			consts.SetSSADataBaseName(databaseFileRaw)

			if target == "" {
				return utils.Errorf("target file not found: %v", rawFile)
			}
			opt := make([]ssaapi.Option, 0, 3)
			log.Infof("start to compile file: %v ", target)
			if input_language != "" {
				input_language = strings.ToLower(input_language)
				var language consts.Language
				switch strings.ToLower(input_language) {
				case "javascript", "js":
					language = ssaapi.JS
				case "yak", "yaklang":
					language = ssaapi.Yak
				default:
					language = consts.Language(input_language)
				}
				log.Infof("start to use language: %v", language)
				opt = append(opt, ssaapi.WithLanguage(language))
			}
			if entry != "" {
				log.Infof("start to use entry file: %v", entry)
				opt = append(opt, ssaapi.WithFileSystemEntry(entry))
			}

			if inMemory {
				log.Infof("compile in memory mode, program-name will be ignored")
			} else {
				if programName == "" {
					programName = "default-" + ksuid.New().String()
				}
				log.Infof("compile save to database with program name: %v", programName)
				opt = append(opt, ssaapi.WithDatabaseProgramName(programName))
			}

			if !noOverride {
				ssadb.DeleteProgram(ssadb.GetDB(), programName)
			} else {
				log.Warnf("no-override flag is set, will not delete existed program: %v", programName)
			}

			proj, err := ssaapi.ParseProjectFromPath(target, opt...)
			if err != nil {
				return utils.Errorf("parse project [%v] failed: %v", target, err)
			}

			log.Infof("finished compiling..., results: %v", len(proj))
			if syntaxFlow != "" {
				log.Warn("Deprecated: syntax flow query language will be removed in ssa sub-command, please use `ssa-query(in short: sf/syntaxFlow)` instead")
				return SyntaxFlowQuery(programName, databaseFileRaw, syntaxFlow, dbDebug, sfDebug, showDot, withCode)
			}
			return nil
		},
	},
	{
		Name:    "syntaxflow-create",
		Aliases: []string{"create-sf", "csf"},
		Usage:   "create syntaxflow template file",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "language,l"},
			cli.StringFlag{Name: "keyword"},
			cli.BoolFlag{Name: "is-vuln,vuln,v", Usage: "set the current SyntaxFlow Rule is a vuln (in desc)"},
			cli.BoolFlag{Name: "audit-suggestion,audit,a", Usage: "set the current SyntaxFlow Rule is a suggestion"},
			cli.BoolFlag{Name: "sec-config,security-config,s", Usage: "set the current SyntaxFlow Rule is a suggestion"},
			cli.StringFlag{Name: "output,o,f", Usage: `set output filename`},
		},
		Action: func(c *cli.Context) error {
			var buf bytes.Buffer

			var typeStrs []string

			switch {
			case c.Bool("is-vuln"):
				typeStrs = append(typeStrs, "vuln")
			case c.Bool("audit-suggestion"):
				typeStrs = append(typeStrs, "audit")
			case c.Bool("security-config"):
				typeStrs = append(typeStrs, "sec-config")
			}

			if len(typeStrs) <= 0 {
				typeStrs = append(typeStrs, "audit")
			}

			buf.WriteString("desc(\n  ")
			buf.WriteString("title: 'checking []',\n  ")
			buf.WriteString("type: " + strings.Join(typeStrs, "|") + "\n)\n\n")
			buf.WriteString("// write your SyntaxFlow Rule, like:\n")
			buf.WriteString("//     DocumentBuilderFactory.newInstance()...parse(* #-> * as $source) as $sink; // find some call chain for parse\n")
			buf.WriteString("//     check $sink then 'find sink point' else 'No Found' // if not found sink, the rule will stop here and report error\n")
			buf.WriteString("//     alert $source // record $source\n\n\n")
			buf.WriteString("// the template is generate by yak.ssa.syntaxflow command line\n")

			filename := c.String("output")

			if l := c.String("language"); filename != "" && l != "" {
				l = strings.TrimSpace(strings.ToLower(l))
				dirname, filename := filepath.Split(filename)
				if !strings.HasPrefix(filename, l+"-") {
					filename = l + "-" + filename
				}
				filename = filepath.Join(dirname, filename)
			}

			if filename == "" {
				fmt.Println(buf.String())
				return nil
			}
			if !strings.HasSuffix(filename, ".sf") {
				filename += ".sf"
			}
			return os.WriteFile(filename, buf.Bytes(), 0o666)
		},
	},
	{
		Name:    "syntaxflow-export",
		Aliases: []string{"sf-export", "esf"},
		Usage:   "export SyntaxFlow rule to file system",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "output,o",
				Usage: "output file path",
			},
		},
		Action: func(c *cli.Context) error {
			if c.String("output") == "" {
				return utils.Error("output file is required")
			}
			local := filesys.NewLocalFs()
			results, _ := io.ReadAll(sfdb.Export())
			if len(results) <= 0 {
				return utils.Error("no rule found")
			}
			return local.WriteFile(c.String("output"), results, 0o666)
		},
	},
	{
		Name:    "syntaxflow-import",
		Usage:   "import SyntaxFlow rule from file system",
		Aliases: []string{"sf-import", "isf"},
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "file,f,i",
				Usage: "file path",
			},
		},
		Action: func(c *cli.Context) error {
			if c.String("file") == "" {
				return utils.Error("file is required")
			}
			file, err := os.Open(c.String("file"))
			if err != nil {
				return utils.Wrap(err, "open file failed")
			}
			err = sfdb.Import(file)
			if err != nil {
				return err
			}
			return nil
		},
	},
	{
		Name:    "syntaxflow-save",
		Aliases: []string{"save-syntaxflow", "ssf", "sfs"},
		Usage:   "save SyntaxFlow rule to database",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "filesystem,f",
				Usage: "file system for MVP",
			},
			cli.StringFlag{
				Name: "rule,r",
			},
		},
		Action: func(c *cli.Context) error {
			count := 0
			err := filesys.Recursive(c.String("filesystem"), filesys.WithFileStat(func(s string, info fs.FileInfo) error {
				count++
				if count > 50 {
					return utils.Error("too many files")
				}

				if info.Size() > 2*1024*1024 {
					return utils.Errorf("file %v size too large", s)
				}
				return nil
			}))
			if err != nil {
				return utils.Wrap(err, "read mvp file system failed")
			}

			memfs := filesys.NewVirtualFs()
			local := filesys.NewLocalFs()
			err = filesys.Recursive(c.String("filesystem"), filesys.WithFileStat(func(s string, info fs.FileInfo) error {
				raw, err := local.ReadFile(s)
				if err != nil {
					return nil
				}
				memfs.AddFile(s, string(raw))
				return nil
			}))
			if err != nil {
				return err
			}

			contentRaw, _ := local.ReadFile(c.String("rule"))
			if len(contentRaw) > 0 {
				err = sfdb.ImportValidRule(memfs, c.String("rule"), string(contentRaw))
				if err != nil {
					log.Warnf("import rule failed: %v", err)
				}
				return nil
			}

			entrys, err := utils.ReadDir(c.String("rule"))
			if err != nil {
				return err
			}
			for _, entry := range entrys {
				contentRaw, _ := local.ReadFile(entry.Path)
				if len(contentRaw) <= 0 {
					continue
				}
				err = sfdb.ImportValidRule(memfs, entry.Path, string(contentRaw))
				if err != nil {
					log.Warnf("import rule failed: %v", err)
					continue
				}
			}
			return nil
		},
	},
	{
		Name:    "ssa-query",
		Aliases: []string{"sf", "syntaxFlow", "sf-scan"},
		Usage:   "Use SyntaxFlow query SSA OpCodes from database",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "log", Usage: "log level"},
			cli.StringFlag{
				Name:  "program,p",
				Usage: `program name to save in database`,
			},
			cli.StringFlag{
				Name:  "syntaxflow,sf",
				Usage: "syntax flow query language code",
			},
			cli.StringFlag{
				Name:  "database,db",
				Usage: "database path",
			},
			cli.BoolFlag{
				Name:  "database-debug,dbdebug",
				Usage: "enable database debug mode",
			},
			cli.BoolFlag{
				Name:  "syntaxflow-debug,sfdebug",
				Usage: "enable syntax flow debug mode",
			},
			cli.BoolFlag{
				Name: "dot", Usage: "dot graph text for result",
			},
			cli.BoolFlag{
				Name: "with-code,code", Usage: "show code context",
			},
			cli.StringFlag{
				Name: "sarif,sarif-export,o", Usage: "export SARIF format to files",
			},
		},
		Action: func(c *cli.Context) error {
			if ret, err := log.ParseLevel(c.String("log")); err == nil {
				log.SetLevel(ret)
			}
			programName := c.String("program")
			databaseFileRaw := c.String("database")
			dbDebug := c.Bool("database-debug")
			sfDebug := c.Bool("syntaxflow-debug")
			syntaxFlow := c.String("syntaxflow")
			showDot := c.Bool("dot")
			withCode := c.Bool("with-code")

			sarifFile := c.String("sarif")
			if sarifFile != "" {
				if filepath.Ext(sarifFile) != ".sarif" {
					sarifFile += ".sarif"
				}
			}

			haveSarifRequired := false
			if sarifFile != "" {
				haveSarifRequired = true
			}
			var results []*sfvm.SFFrameResult
			var sarifCallback func(result *sfvm.SFFrameResult)
			if haveSarifRequired {
				sarifCallback = func(result *sfvm.SFFrameResult) {
					results = append(results, result)
				}
			} else {
				sarifCallback = func(result *sfvm.SFFrameResult) {

				}
			}

			defer func() {
				if len(results) > 0 && sarifFile != "" {
					log.Infof("fetch result: %v, exports sarif to %v", len(results), sarifFile)
					report, err := ssaapi.ConvertSyntaxFlowResultToSarif(results...)
					if err != nil {
						log.Errorf("convert SARIF failed: %v", err)
						return
					}
					if utils.GetFirstExistedFile(sarifFile) != "" {
						backup := sarifFile + ".bak"
						os.Rename(sarifFile, backup)
						os.RemoveAll(sarifFile)
					}
					err = report.WriteFile(sarifFile)
					if err != nil {
						log.Errorf("write SARIF failed: %v", err)
					}
				}
			}()

			if syntaxFlow != "" {
				return SyntaxFlowQuery(programName, databaseFileRaw, syntaxFlow, dbDebug, sfDebug, showDot, withCode, sarifCallback)
			}

			var dirChecking []string

			handleBySyntaxFlowContent := func(syntaxFlow string) error {
				err := SyntaxFlowQuery(programName, databaseFileRaw, syntaxFlow, dbDebug, sfDebug, showDot, withCode, sarifCallback)
				if err != nil {
					return err
				}
				fmt.Println()
				return nil
			}

			handleByFilename := func(filename string) error {
				log.Infof("start to use SyntaxFlow rule: %v", filename)
				raw, err := os.ReadFile(filename)
				if err != nil {
					return utils.Wrapf(err, "read %v failed", filename)
				}
				return handleBySyntaxFlowContent(string(raw))
			}

			var errs []error
			var cmdArgs []string = c.Args()
			for _, originName := range cmdArgs {
				name := utils.GetFirstExistedFile(originName)
				if name == "" {
					infos, _ := utils.ReadDir(originName)
					if len(infos) > 0 {
						dirChecking = append(dirChecking, originName)
						continue
					}

					if filepath.IsAbs(originName) {
						log.Warnf("cannot find rule as %v", originName)
					} else {
						absName, _ := filepath.Abs(originName)
						if absName != "" {
							log.Warnf("cannot find rule as %v(abs: %v)", originName, absName)
						} else {
							log.Warnf("cannot find rule as %v", originName)
						}
					}
					continue
				}
				err := handleByFilename(name)
				if err != nil {
					errs = append(errs, err)
				}
			}

			for _, dir := range dirChecking {
				log.Infof("start to read directory: %v", dir)
				err := filesys.Recursive(dir, filesys.WithRecursiveDirectory(true), filesys.WithFileStat(func(s string, info fs.FileInfo) error {
					fileExt := strings.ToLower(filepath.Ext(s))
					if strings.HasSuffix(fileExt, ".sf") {
						err := handleByFilename(s)
						if err != nil {
							errs = append(errs, err)
						}
					}
					return nil
				}))
				if err != nil {
					log.Warnf("read directory [%v] failed: %v", dir, err)
				}
			}

			if len(cmdArgs) <= 0 {
				prog, err := ssaapi.FromDatabase(programName)
				if err != nil {
					log.Errorf("load program [%v] from database failed: %v", programName, err)
					return err
				}

				db := consts.GetGormProfileDatabase()
				expected := []string{""}
				for _, l := range utils.PrettifyListFromStringSplitEx(prog.GetLanguage(), ",") {
					if l == "" {
						continue
					}
					expected = append(expected, l)
				}
				db = bizhelper.ExactQueryStringArrayOr(db, "language", expected)
				for result := range sfdb.YieldSyntaxFlowRules(db, context.Background()) {
					err := handleBySyntaxFlowContent(result.Content)
					if err != nil {
						errs = append(errs, err)
					}
				}
			}

			if len(errs) > 0 {
				var buf bytes.Buffer
				for i, e := range errs {
					buf.WriteString("  ")
					buf.WriteString(fmt.Sprintf("%-2d: ", i+1))
					buf.WriteString(e.Error())
					buf.WriteByte('\n')
				}
				return utils.Errorf("many error happened: \n%v", buf.String())
			}
			return nil
		},
	},
}
View Source
var ScanCommands = []*cli.Command{
	{
		Name:    "pull-plugins",
		Aliases: []string{"pull"},
		Usage:   "pull plugins from yaklang.io and nuclei-templates",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:   "proxy",
				Usage:  "Proxy Server(http/socks5...)",
				EnvVar: "http_proxy",
			},
			cli.StringFlag{
				Name:  "base-url,u",
				Usage: "yaklang / yakit plugin server url",
				Value: `https://www.yaklang.com/`,
			},
			cli.StringFlag{
				Name:  "nuclei-templates-url,n",
				Usage: "Nuclei Templates URL",
				Value: `https://github.com/projectdiscovery/nuclei-templates`,
			},
		},
		Action: func(c *cli.Context) error {
			client := yaklib.NewOnlineClient(c.String("base-url"))
			if c.String("proxy") != "" {
				consts.SetOnlineBaseUrlProxy(c.String("proxy"))
			}
			stream := client.DownloadYakitPluginAll(context.Background())
			count := 0
			for result := range stream.Chan {
				count++
				log.Infof("start to save plugin(%v/%v): %v", count, result.Total, result.Plugin.ScriptName)
				err := client.Save(consts.GetGormProfileDatabase(), result.Plugin)
				if err != nil {
					log.Errorf("save plugin failed: %s", err)
				}
			}

			tools.UpdatePoCWithUrl(c.String(`nuclei-templates-url`), c.String("proxy"))
			return nil
		},
	},
	{
		Name:  "update-nuclei-database",
		Usage: "Load Nuclei-Template into Local Yak Plugin Database",
		Flags: []cli.Flag{
			cli.BoolFlag{
				Name:  "no-cache",
				Usage: "do not use local file cache will not download from git",
			},
			cli.StringFlag{
				Name:  "url",
				Usage: "which url to download?",
				Value: `https://github.com/projectdiscovery/nuclei-templates`,
			},
		},
		Action: func(c *cli.Context) error {
			var err error
			err = yak.NewScriptEngine(1).ExecuteMain(`
loglevel("info")
log.info("start to load local database"); 
die(nuclei.UpdateDatabase())`, "main")
			if err != nil {
				log.Errorf("execute nuclei.UpdateDatabase() failed: %s", err)
				return err
			}
			return nil
		},
	},
	{
		Name: "remove-nuclei-database", Usage: "Remove Nuclei-Template from Local Yak Plugin Database",
		Action: func(c *cli.Context) error {
			err := tools.RemovePoCDatabase()
			if err != nil {
				log.Errorf("remove pocs failed: %s", err)
			}
			return nil
		},
	},
	&synscanCommand,
	&servicescanCommand,
	hybridScanCommand,
	&crawlerxCommand,
}
View Source
var TrafficUtilCommands = []*cli.Command{

	{
		Name:  "import-chaosmaker-json",
		Usage: "Import ChaosMaker Rules from JSON File",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "file,f"},
		},
		Action: func(c *cli.Context) error {
			file := utils.GetFirstExistedFile(c.String("file"))
			if file == "" {
				return utils.Errorf("file not found: %v", c.String("file"))
			}

			return rule.ImportRulesFromFile(consts.GetGormProfileDatabase(), file)
		},
	},
	{
		Name:  "export-chaosmaker-json",
		Usage: "Export ChaosMaker Rules to JSON File",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "file,f"},
		},
		Action: func(c *cli.Context) error {
			return rule.ExportRulesToFile(consts.GetGormProfileDatabase(), c.String("file"))
		},
	},
	&chaosMakerCommand,
	&suricataLoaderCommand,
	&pcapCommand,
}
View Source
var UpgradeCommand = cli.Command{
	Name:  "upgrade",
	Usage: "upgrade / reinstall newest or user-defined yak.",
	Flags: []cli.Flag{
		cli.IntFlag{
			Name:  "timeout",
			Usage: "Set Timeout for download yak binary, default 60s.",
			Value: 60,
		},
		cli.StringFlag{
			Name:  "version,v",
			Usage: "Set the version of yak to download, default latest.",
		},

		cli.BoolFlag{
			Name:  "list,l",
			Usage: "Show all active versions.",
		},

		cli.IntFlag{
			Name:  "n",
			Usage: "Show latest N active versions.",
			Value: 16,
		},
	},
	Action: func(c *cli.Context) error {
		const activeVersions = `https://aliyun-oss.yaklang.com/yak/version-info/active_versions.txt`
		if c.Bool("list") {
			rsp, _, err := poc.DoGET(activeVersions)
			if err != nil {
				log.Errorf("fetch active versions failed: %v", err)
				return err
			}
			versions := utils.PrettifyListFromStringSplitEx(string(rsp.GetBody()), "\n")
			if len(versions) == 0 {
				log.Errorf("fetch active versions failed: %v", err)
				return err
			}
			log.Infof("active versions: len: %v", len(versions))
			if c.Int("n") > 0 {
				log.Infof("show latest %v active versions", c.Int("n"))
				versions = versions[:c.Int("n")]
			}
			for _, ver := range versions {
				fmt.Println(ver)
			}
			return nil
		}

		exePath, err := os.Executable()
		exeDir := filepath.Dir(exePath)
		if err != nil {
			return utils.Errorf("cannot fetch os.Executable()...: %s", err)
		}

		version := c.String("version")
		if version == "" {
			rsp, _, err := poc.DoGET(`https://aliyun-oss.yaklang.com/yak/latest/version.txt`, poc.WithTimeout(10))
			if err != nil {
				log.Warnf("fetch latest yak version failed: %v", err)
				return err
			}
			raw := lowhttp.GetHTTPPacketBody(rsp.RawPacket)
			version = strings.TrimSpace(string(raw))
		}

		if version == "" {
			log.Warnf("fetch latest yak version failed: %v use latest", err)
			version = "latest"
		}

		fetchUrl := func(ver string) string {
			return fmt.Sprintf(`https://aliyun-oss.yaklang.com/yak/%v/yak_%v_%v`, ver, runtime.GOOS, runtime.GOARCH)
		}

		downloadUrl := fetchUrl(version)
		rsp, _, err := poc.DoHEAD(downloadUrl, poc.WithTimeout(10))
		if err != nil {
			log.Errorf("fetch yak binary failed: %v", err)
			return err
		}
		_ = rsp
		if rsp.GetStatusCode() >= 400 {
			log.Infof("fetch yak binary failed: %v", rsp.GetStatusCode())
			rsp, _, err := poc.DoGET(activeVersions)
			if err != nil {
				log.Errorf("fetch active versions failed: %v", err)
				return err
			}
			versions := utils.PrettifyListFromStringSplitEx(string(rsp.GetBody()), "\n")
			if len(versions) == 0 {
				log.Errorf("fetch active versions failed: %v", err)
				return err
			}
			log.Infof("active versions: len: %v", len(versions))
			for _, ver := range versions {
				fmt.Println(ver)
			}
			return nil
		}

		newFilePath := filepath.Join(exeDir, "yak.new")
		fd, err := os.OpenFile(newFilePath, os.O_RDWR|os.O_CREATE|os.O_TRUNC, 0o766)
		if err != nil {
			log.Errorf("create temp file failed: %v", err)
			return err
		}
		defer func() {
			fd.Sync()
			fd.Close()
		}()

		binary := downloadUrl
		log.Infof("start to download yak: %v", binary)
		timeout := c.Int("timeout")
		ctx, cancel := context.WithCancel(context.Background())
		_, _, err = poc.DoGET(binary, poc.WithConnectTimeout(float64(timeout)), poc.WithBodyStreamReaderHandler(func(header []byte, bodyReader io.ReadCloser) {
			defer func() {
				cancel()
			}()

			log.Infof("downloading yak binary...")
			contentLength := lowhttp.GetHTTPPacketHeader(header, "content-length")
			writer := progresswriter.New(uint64(codec.Atoi(contentLength)))
			writer.ShowProgress("downloading", ctx)
			_, err := io.Copy(io.MultiWriter(fd, writer), bodyReader)
			if err != nil && err != io.EOF {
				log.Errorf("download yak failed: %v", err)
				return
			}
		}))
		if err != nil {
			log.Errorf("download yak failed: %v", err)
			return err
		}

		destDir, _ := filepath.Split(exePath)
		backupPath := filepath.Join(destDir, fmt.Sprintf("yak_%s", consts.GetYakVersion()))
		if runtime.GOOS == "windows" {
			backupPath += ".exe"
		}
		log.Infof("backup yak old engine to %s", backupPath)

		log.Infof("origin binary: %s", exePath)

		if err := os.Rename(exePath, backupPath); err != nil {
			return utils.Errorf("backup old yak-engine failed: %s, retry re-Install with \n"+
				"    `bash <(curl -sS -L http://oss.yaklang.io/install-latest-yak.sh)`\n\n", err)
		}

		if err := os.Rename(newFilePath, exePath); err != nil {

			rerr := os.Rename(backupPath, exePath)
			if rerr != nil {
				return utils.Errorf("rename new yak-engine failed: %s, rollback failed: %s, retry re-Install with \n"+"    `bash <(curl -sS -L http://oss.yaklang.io/install-latest-yak.sh)`\n\n", err, rerr)
			}

			return utils.Errorf("rename new yak-engine failed: %s, retry re-Install with \n"+
				"    `bash <(curl -sS -L http://oss.yaklang.io/install-latest-yak.sh)`\n\n", err)
		}
		return nil
	},
}
View Source
var UtilsCommands = []*cli.Command{
	{
		Name:  "gzip",
		Usage: "gzip data or file",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "f,file",
				Usage: "input file",
			},
			cli.BoolFlag{Name: "d,decode"},
			cli.StringFlag{Name: "o,output"},
		}, Action: func(c *cli.Context) error {
			f := c.String("file")
			if utils.GetFirstExistedFile(f) == "" {
				return utils.Errorf("non-existed: %v", f)
			}
			originFp, err := os.Open(f)
			if err != nil {
				return err
			}
			defer originFp.Close()

			if c.Bool("decode") {
				outFile := c.String("output")
				if outFile == "" {
					return utils.Error("decode need output not empty")
				}
				log.Infof("start to d-gzip to %v", outFile)
				targetFp, err := os.OpenFile(outFile, os.O_CREATE|os.O_RDWR, 0o666)
				if err != nil {
					return err
				}
				defer targetFp.Close()
				r, err := gzip.NewReader(originFp)
				if err != nil {
					return err
				}
				defer r.Close()
				io.Copy(targetFp, r)
				log.Infof("finished")
				return nil
			}

			gf := f + ".gzip"
			fp, err := os.OpenFile(gf, os.O_CREATE|os.O_RDWR, 0o666)
			if err != nil {
				return err
			}
			defer fp.Close()
			gzipWriter := gzip.NewWriter(fp)
			io.Copy(gzipWriter, originFp)
			gzipWriter.Flush()
			gzipWriter.Close()
			return nil
		},
	},
	{
		Name: "hex",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "f,file",
				Usage: "input file",
			},
			cli.StringFlag{
				Name:  "d,data",
				Usage: "input data",
			},
		},
		Usage: "hex encode file or data to hex string",
		Action: func(c *cli.Context) {
			if c.String("file") != "" {
				raw, err := ioutil.ReadFile(c.String("file"))
				if err != nil {
					log.Error(err)
					return
				}
				println(codec.EncodeToHex(raw))
			}

			if c.String("data") != "" {
				println(codec.EncodeToHex(c.String("data")))
			}
		},
	},
	{
		Name:  "tag-stats",
		Usage: "Generate Tag Status(for Yakit)",
		Action: func(c *cli.Context) error {
			stats, err := yaklib.NewTagStat()
			if err != nil {
				return err
			}
			for _, v := range stats.All() {
				if v.Count <= 1 {
					continue
				}
				fmt.Printf("TAG:[%v]-%v\n", v.Name, v.Count)
			}
			return nil
		},
	},

	{
		Name:  "dap",
		Usage: "Start a server based on the Debug Adapter Protocol (DAP) to debug scripts.",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "host", Usage: "debugger adapter listen host"},
			cli.IntFlag{Name: "port", Usage: "debugger adapter listen port"},
			cli.BoolFlag{Name: "debug", Usage: "debug mode"},
			cli.BoolFlag{Name: "version,v", Usage: "show dap version"},
		},
		Action: func(c *cli.Context) error {
			host := c.String("host")
			port := c.Int("port")
			debug := c.Bool("debug")
			versionFlag := c.Bool("version")
			if versionFlag {
				fmt.Printf("Debugger Adapter version: %v\n", dap.DAVersion)
				return nil
			}

			if debug {
				log.SetLevel(log.DebugLevel)
			}

			server, stopChan, err := dap.StartDAPServer(host, port)
			if err != nil {
				return err
			}
			defer server.Stop()

			forceStop := make(chan struct{})
			select {
			case <-stopChan:
			case <-forceStop:
			}

			return nil
		},
	},

	{
		Name:  "fmt",
		Usage: "Formatter for Yaklang Code",
		Flags: []cli.Flag{
			cli.BoolFlag{Name: "version,v", Usage: "show formatter version"},
		},
		Action: func(c *cli.Context) error {
			if c.Bool("version") {
				fmt.Printf("Formatter version: %v\n", yakast.FormatterVersion)
				return nil
			}
			args := c.Args()
			file := args[0]
			if file != "" {
				var err error
				absFile := file
				if !filepath.IsAbs(absFile) {
					absFile, err = filepath.Abs(absFile)
					if err != nil {
						return utils.Errorf("fetch abs file path failed: %s", err)
					}
				}
				raw, err := os.ReadFile(file)
				if err != nil {
					return err
				}
				vt := yakast.NewYakCompiler()
				vt.Compiler(string(raw))
				fmt.Printf("%s", vt.GetFormattedCode())
			} else {
				return utils.Errorf("empty yak file")
			}
			return nil
		},
	},

	{
		Name:  "fuzz",
		Usage: "fuzztag short for fuzz tag, fuzz tag is a tool to generate fuzz string for fuzz testing",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "t,target",
				Usage: "Fuzztag Template, like: `{{int(1-10)}}`",
			},
		},
		Action: func(c *cli.Context) {
			for _, r := range mutate.MutateQuick(c.String("t")) {
				println(r)
			}
		},
	},

	{
		Name:  "sha256",
		Usage: "(Inner command) sha256 checksums for file and generate [filename].sha256.txt",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "file,f", Usage: "file to checksum"},
		},
		Action: func(c *cli.Context) error {
			filename := c.String("file")
			if filename == "" {
				return utils.Errorf("empty filename")
			}
			file, err := os.Open(filename)
			if err != nil {
				return err
			}
			defer func() {
				file.Close()
			}()
			hasher := sha256.New()
			if _, err := io.Copy(hasher, file); err != nil && err != io.EOF {
				return err
			}
			sum := hasher.Sum(nil)
			result := codec.EncodeToHex(sum)

			targetFile := filename + ".sha256.txt"
			err = os.WriteFile(targetFile, []byte(result), 0o644)
			if err != nil {
				return err
			}
			fmt.Printf("file[%s] Sha256 checksum: %s\nGenerate to %s", filename, result, targetFile)
			return nil
		},
	},
	{
		Name:  "repos-tag",
		Usage: "(Inner command) Get Current Git Repository Tag, if not found, generate a fallback tag with dev/{date}",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "output,o", Usage: "output file", Value: "tags.txt"},
		},
		Action: func(c *cli.Context) error {
			var err error
			fallback := func(suffix string) error {
				results := "dev/" + utils.DatePretty() + suffix
				return os.WriteFile(c.String("output"), []byte(results), 0o644)
			}
			rp, err := git.PlainOpen(".")
			if err != nil {
				return fallback("")
			}
			ref, err := rp.Head()
			if err != nil {
				return fallback("")
			}
			var suffix string
			if ref != nil && !ref.Hash().IsZero() {
				h := ref.Hash().String()
				if len(h) > 8 {
					suffix = "-" + h[:8]
				} else {
					suffix = "-" + h
				}
			}

			tags, err := rp.Tags()
			if err != nil {
				return fallback(suffix)
			}

			// 查找与当前 HEAD 提交相关联的标签
			var foundTags []string
			err = tags.ForEach(func(t *plumbing.Reference) error {
				if t.Hash() == ref.Hash() {
					foundTags = append(foundTags, t.Name().Short())
				}
				return nil
			})
			if err != nil {
				return fallback(suffix)
			}

			if len(foundTags) > 0 {
				return os.WriteFile(c.String("output"), []byte(strings.TrimLeft(foundTags[0], "v")), 0o644)
			}
			return fallback(suffix)
		},
	},

	{
		Name:  "upload-oss",
		Usage: "(Inner command) Upload File To Aliyun OSS",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "file,f", Usage: "local_file_path:remote_file_path, splited by ;"},
			cli.StringFlag{Name: "ak", Usage: "Aliyun Access Key"},
			cli.StringFlag{Name: "sk", Usage: "Aliyun Secret Key"},
			cli.StringFlag{Name: "endpoint", Usage: "Aliyun OSS Endpoint", Value: `oss-accelerate.aliyuncs.com`},
			cli.StringFlag{Name: "bucket, b", Usage: "Aliyun OSS Bucket", Value: "yaklang"},
			cli.IntFlag{Name: "times,t", Usage: "retry times", Value: 5},
		},
		Action: func(c *cli.Context) error {
			client, err := oss.New(c.String("endpoint"), c.String("ak"), c.String("sk"))
			if err != nil {
				return err
			}

			bucket, err := client.Bucket(c.String("bucket"))
			if err != nil {
				return err
			}
			for _, i := range strings.Split(c.String("file"), ";") {
				localFilePath, remoteFilePath, ok := strings.Cut(i, ":")
				if !ok {
					return utils.Errorf("invalid file path: %v", i)
				}
				localFilePath = strings.TrimSpace(localFilePath)
				remoteFilePath = strings.TrimSpace(strings.TrimLeft(remoteFilePath, "/"))

				_, _, err = lo.AttemptWithDelay(c.Int("times"), time.Second, func(index int, _ time.Duration) error {
					return bucket.PutObjectFromFile(remoteFilePath, localFilePath)
				})
				if err != nil {
					return utils.Wrap(err, "upload file to oss failed")
				}
			}

			return nil
		},
	},

	{
		Name:  "weight",
		Usage: "weight dir with depth",
		Flags: []cli.Flag{
			cli.StringFlag{Name: "dir,d", Usage: "dir to weight"},
			cli.IntFlag{Name: "depth", Usage: "depth to weight", Value: 1},
			cli.BoolFlag{Name: "asc", Usage: "sort asc"},
			cli.StringFlag{Name: "blacklist,exclude", Usage: "ignore blacklist", Value: "*_test.go|.git*|*testdata*"},
			cli.StringFlag{Name: "show-exclude", Usage: "filter result", Value: "*.md|*.yak|*.DS_Store|*License|*.g4"},
			cli.IntFlag{Name: "show-min-size", Usage: "show min size", Value: 100000},
		},
		Action: func(c *cli.Context) error {
			m := omap.NewOrderedMap(map[string]int64{})
			err := filesys.Recursive(c.String("dir"), filesys.WithFileStat(func(pathname string, info fs.FileInfo) error {
				if c.String("blacklist") != "" {
					if utils.MatchAnyOfGlob(pathname, utils.PrettifyListFromStringSplitEx(c.String("blacklist"), "|")...) {
						return nil
					}
				}
				log.Infof("path: %v, size: %v verbose: %v", pathname, info.Size(), utils.ByteSize(uint64(info.Size())))
				m.Set(pathname, info.Size())
				return nil
			}))
			if err != nil {
				return err
			}
			forest, err := utils.GeneratePathTrees(m.Keys()...)
			if err != nil {
				return err
			}

			results := omap.NewOrderedMap(make(map[string]int64))
			forest.Recursive(func(node2 *utils.PathNode) {
				if node2.GetDepth() > c.Int("depth") {
					return
				}
				count := int64(0)
				for _, child := range node2.AllChildren() {
					size, ok := m.Get(child.Path)
					if !ok {
						log.Warnf("path: %v, name: %v not found", child.Path, child.Name)
						continue
					}
					count += size
				}
				results.Set(node2.Path, count)
			})

			var desc []*sizeDescription
			results.ForEach(func(i string, v int64) bool {
				if c.String("show-exclude") != "" {
					if utils.MatchAnyOfGlob(i, utils.PrettifyListFromStringSplitEx(c.String("show-exclude"), "|")...) {
						return true
					}
				}
				desc = append(desc, &sizeDescription{Name: i, Size: uint64(v)})
				return true
			})

			sort.Slice(desc, func(i, j int) bool {
				if c.Bool("asc") {
					return desc[i].Size < desc[j].Size
				}
				return desc[i].Size > desc[j].Size
			})

			for _, i := range desc {
				fmt.Printf("[%6s]: %v\n", utils.ByteSize(i.Size), i.Name)
			}
			return nil
		},
	},

	{
		Name: "totp-forward",
		Flags: []cli.Flag{
			cli.StringFlag{
				Name:  "secret",
				Usage: "totp secret",
			},
			cli.StringFlag{
				Name:  "proxy-for",
				Usage: "which port for forwarding to",
			},
			cli.IntFlag{
				Name: "listen,l", Usage: "which port for listening", Value: 8084,
			},
		},
		Action: func(c *cli.Context) {
			var secret string = c.String("secret")
			var lisPort = c.Int("listen")
			if lisPort <= 0 {
				lisPort = 8084
			}
			if secret == "" {

			}

			if secret == "" {
				log.Warn("empty secret")
				return
			}

			for {
				err := twofa.NewOTPServer(secret, lisPort, c.String("proxy-for")).Serve()
				if err != nil {
					log.Errorf("failed to serve: %v", err)
					time.Sleep(time.Second)
					continue
				}
			}
		},
	},
}
View Source
var YsoCommands = []*cli.Command{}

Functions

func ShowHistoryHTTPFlowByRuntimeId added in v1.3.2

func ShowHistoryHTTPFlowByRuntimeId(last int64, runtimeId string) int64

func SyntaxFlowQuery

func SyntaxFlowQuery(
	programName, databaseFileRaw string,
	syntaxFlow string,
	dbDebug, sfDebug, showDot, withCode bool,
	callbacks ...func(*sfvm.SFFrameResult),
) error

Types

This section is empty.

Jump to

Keyboard shortcuts

? : This menu
/ : Search site
f or F : Jump to
y or Y : Canonical URL