yakcmds

package
v1.3.2 Latest Latest
Warning

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

Go to latest
Published: Apr 12, 2024 License: AGPL-3.0 Imports: 73 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 := gorm.Open("sqlite3", outputFile)
			if err != nil {
				return err
			}
			outputDB.AutoMigrate(&cveresources.CVE{}, &cveresources.CWE{})
			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 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 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 yak.",
	Flags: []cli.Flag{
		cli.IntFlag{
			Name:  "timeout",
			Usage: "连接超时时间",
			Value: 30,
		},
	},
	Action: func(c *cli.Context) error {
		exePath, err := os.Executable()
		exeDir := filepath.Dir(exePath)
		if err != nil {
			return utils.Errorf("cannot fetch os.Executable()...: %s", err)
		}

		binary := fmt.Sprintf(`https://yaklang.oss-accelerate.aliyuncs.com/yak/latest/yak_%v_%v`, runtime.GOOS, runtime.GOARCH)
		if runtime.GOOS == "darwin" && runtime.GOARCH == "arm64" {
			binary = fmt.Sprintf(`https://yaklang.oss-accelerate.aliyuncs.com/yak/latest/yak_%v_%v`, runtime.GOOS, "amd64")
		} else if runtime.GOOS == "windows" {
			binary = fmt.Sprintf(`https://yaklang.oss-accelerate.aliyuncs.com/yak/latest/yak_%v_%v.exe`, runtime.GOOS, "amd64")
		}

		versionUrl := `https://yaklang.oss-accelerate.aliyuncs.com/yak/latest/version.txt`
		timeout := float64(c.Int("timeout"))
		rspIns, _, err := poc.DoGET(versionUrl, poc.WithTimeout(timeout))
		if err != nil {
			log.Errorf("获取 yak 引擎最新版本失败:get yak latest version failed: %v", err)
			return err
		}
		if len(rspIns.RawPacket) > 0 {
			raw := lowhttp.GetHTTPPacketBody(rspIns.RawPacket)
			if len(utils.ParseStringToLines(string(raw))) <= 3 {
				log.Infof("当前 yak 核心引擎最新版本为 / current latest yak core engine version:%v", string(raw))
			}
		}

		log.Infof("start to download yak: %v", binary)
		rspIns, _, err = poc.DoGET(binary, poc.WithTimeout(timeout))
		if err != nil {
			log.Errorf("下载 yak 引擎失败:download yak failed: %v", err)
			return err
		}

		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
		}

		log.Infof("downloading for yak binary to local")
		_, err = io.Copy(fd, rspIns.MultiResponseInstances[0].Body)
		if err != nil && err != io.EOF {
			log.Errorf("download failed...: %v", err)
			return err
		}
		log.Infof("yak 核心引擎下载成功... / yak engine downloaded")
		fd.Sync()
		fd.Close()

		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:  "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 os.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
		},
	},
}
View Source
var YsoCommands = []*cli.Command{}

Functions

func ShowHistoryHTTPFlowByRuntimeId added in v1.3.2

func ShowHistoryHTTPFlowByRuntimeId(last int64, runtimeId string) int64

Types

This section is empty.

Jump to

Keyboard shortcuts

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