diff --git a/cmd/root.go b/cmd/root.go index 8ca105d..7b7e822 100644 --- a/cmd/root.go +++ b/cmd/root.go @@ -3,13 +3,14 @@ package cmd import ( "bufio" "fmt" + "log" + "os" + "strings" + "github.com/fatih/color" "github.com/spf13/cobra" "golang.org/x/text/encoding/simplifiedchinese" "golang.org/x/text/transform" - "log" - "os" - "strings" "github.com/zu1k/nali/internal/constant" "github.com/zu1k/nali/pkg/common" @@ -61,7 +62,6 @@ Find document on: https://github.com/zu1k/nali Run: func(cmd *cobra.Command, args []string) { gbk, _ := cmd.Flags().GetBool("gbk") isJson, _ := cmd.Flags().GetBool("json") - isJsonline, _ := cmd.Flags().GetBool("jsonline") if len(args) == 0 { stdin := bufio.NewScanner(os.Stdin) @@ -75,17 +75,14 @@ Find document on: https://github.com/zu1k/nali return } if isJson { - _, _ = fmt.Fprintf(color.Output, "%s", entity.ParseLine(line).Json()) + _, _ = fmt.Fprintf(color.Output, "%s\n", entity.ParseLine(line).Json()) } else { - _, _ = fmt.Fprintf(color.Output, "%s", entity.ParseLine(line).ColorString()) + _, _ = fmt.Fprintf(color.Output, "%s\n", entity.ParseLine(line).ColorString()) } - } } else { if isJson { _, _ = fmt.Fprintf(color.Output, "%s\n", entity.ParseLine(strings.Join(args, " ")).Json()) - } else if isJsonline { - _, _ = fmt.Fprintf(color.Output, "%s\n", entity.ParseLine(strings.Join(args, " ")).JsonLine()) } else { _, _ = fmt.Fprintf(color.Output, "%s\n", entity.ParseLine(strings.Join(args, " ")).ColorString()) } @@ -102,6 +99,5 @@ func Execute() { func init() { rootCmd.Flags().Bool("gbk", false, "Use GBK decoder") - rootCmd.PersistentFlags().BoolP("json", "j", false, "Output in JSON format") - rootCmd.PersistentFlags().BoolP("jsonline", "l", false, "JSON output with newlines") + rootCmd.Flags().BoolP("json", "j", false, "Output in JSON format") } diff --git a/internal/db/db.go b/internal/db/db.go index 24fcad3..f8a4c2a 100644 --- a/internal/db/db.go +++ b/internal/db/db.go @@ -1,11 +1,12 @@ package db import ( - "github.com/spf13/viper" - "github.com/zu1k/nali/pkg/wry" "log" + "github.com/spf13/viper" + "github.com/zu1k/nali/pkg/cdn" + "github.com/zu1k/nali/pkg/common" "github.com/zu1k/nali/pkg/dbif" "github.com/zu1k/nali/pkg/geoip" "github.com/zu1k/nali/pkg/qqwry" @@ -68,15 +69,15 @@ func GetDB(typ dbif.QueryType) (db dbif.DB) { return } -func Find(typ dbif.QueryType, query string) *wry.Result { +func Find(typ dbif.QueryType, query string) common.Result { if result, found := queryCache.Load(query); found { - return result.(*wry.Result) + return result.(common.Result) } result, err := GetDB(typ).Find(query) if err != nil { return nil } - res := result.(*wry.Result) + res := result queryCache.Store(query, res) return res } diff --git a/pkg/common/result.go b/pkg/common/result.go new file mode 100644 index 0000000..c66e6bc --- /dev/null +++ b/pkg/common/result.go @@ -0,0 +1,5 @@ +package common + +type Result interface { + String() string +} diff --git a/pkg/entity/entity.go b/pkg/entity/entity.go index fee458b..618cc62 100644 --- a/pkg/entity/entity.go +++ b/pkg/entity/entity.go @@ -22,9 +22,9 @@ type Entity struct { Loc [2]int `json:"-"` // s[Loc[0]:Loc[1]] Type EntityType `json:"type"` - Text string `json:"ip"` - Info string `json:"tag"` - GEO string `json:"geo"` + Text string `json:"ip"` + InfoText string `json:"text"` + Info interface{} `json:"info"` } func (e Entity) ParseInfo() error { @@ -57,8 +57,8 @@ func (es Entities) String() string { var result strings.Builder for _, entity := range es { result.WriteString(entity.Text) - if entity.Type != TypePlain && len(entity.Info) > 0 { - result.WriteString("[" + entity.Info + "] ") + if entity.Type != TypePlain && len(entity.InfoText) > 0 { + result.WriteString("[" + entity.InfoText + "] ") } } return result.String() @@ -67,35 +67,30 @@ func (es Entities) String() string { func (es Entities) ColorString() string { var line strings.Builder for _, e := range es { - s := e.Text switch e.Type { case TypeIPv4: - s = color.GreenString(e.Text) + line.WriteString(color.GreenString(e.Text)) case TypeIPv6: - s = color.BlueString(e.Text) + line.WriteString(color.BlueString(e.Text)) case TypeDomain: - s = color.YellowString(e.Text) + line.WriteString(color.YellowString(e.Text)) + default: + line.WriteString(e.Text) } - if e.Type != TypePlain && len(e.Info) > 0 { - s += " [" + color.RedString(e.Info) + "] " + if e.Type != TypePlain { + if len(e.InfoText) > 0 { + line.WriteString(" [" + color.RedString(e.InfoText) + "] ") + } } - line.WriteString(s + "\n") + line.WriteString("\n") } - return line.String() + return strings.TrimSpace(line.String()) } func (es Entities) Json() string { - jsonResult, err := json.Marshal(es) - if err != nil { - log.Fatal(err.Error()) - } - return string(jsonResult) -} - -func (es Entities) JsonLine() string { var s strings.Builder for _, e := range es { s.WriteString(e.Json() + "\n") } - return s.String() + return strings.TrimSpace(s.String()) } diff --git a/pkg/entity/parse.go b/pkg/entity/parse.go index 17c8855..93bd99f 100644 --- a/pkg/entity/parse.go +++ b/pkg/entity/parse.go @@ -48,27 +48,15 @@ func ParseLine(line string) Entities { for _, e := range tmp { start := e.Loc[0] if start >= idx { - //if start > idx { - // es = append(es, &Entity{ - // Loc: [2]int{idx, start}, - // Type: TypePlain, - // Text: line[idx:start], - // }) - //} res := db.Find(dbif.QueryType(e.Type), e.Text) - e.Info = res.Area - e.GEO = res.Country - es = append(es, e) - idx = e.Loc[1] + if res != nil { + e.InfoText = res.String() + e.Info = res + es = append(es, e) + idx = e.Loc[1] + } } } - if total := len(line); idx < total { - es = append(es, &Entity{ - Loc: [2]int{idx, total}, - Type: TypePlain, - Text: line[idx:total], - }) - } return es }