From 257ce4aa9206c1a5781dd795fff996a8a5b63b16 Mon Sep 17 00:00:00 2001 From: Khalid Nowaf Date: Tue, 2 Jul 2024 22:21:03 +0300 Subject: [PATCH] feat(CLI): new flags(output-format, drop-keys, split-ip-versions, fill-empty-priority,fip-rank-priority) --- pkg/cli/cli.go | 2 + pkg/cli/reslove.go | 163 +++++++++++++++++++++------------------------ 2 files changed, 79 insertions(+), 86 deletions(-) diff --git a/pkg/cli/cli.go b/pkg/cli/cli.go index 41e144a..139da5e 100644 --- a/pkg/cli/cli.go +++ b/pkg/cli/cli.go @@ -19,7 +19,9 @@ var cli struct { func NewCLI(super *supernet.Supernet) { ctx := kong.Parse(&cli, kong.UsageOnError()) + if cli.Log { + // configure supernet to use simple logger super = supernet.WithSimpleLogger()(super) } if err := ctx.Run(&Context{super: super}); err != nil { diff --git a/pkg/cli/reslove.go b/pkg/cli/reslove.go index f631d4b..8e1b407 100644 --- a/pkg/cli/reslove.go +++ b/pkg/cli/reslove.go @@ -1,119 +1,110 @@ package cli import ( - "encoding/csv" - "encoding/json" "fmt" - "os" + "path/filepath" + "time" "github.com/khalid-nowaf/supernet/pkg/supernet" ) +type Stats struct { + Input int + Output int + Conflicted int + StartInsertTime time.Time + EndInsertTime time.Time + StartOutputTime time.Time + EndOutputTime time.Time +} + type ResolveCmd struct { - Files []string `arg:"" type:"existingfile" help:"Input file containing CIDRs in CSV or JSON format"` - CidrKey string `help:"Index of the CIDRs in the file" default:"cidr"` - PriorityKey string `help:"Index of the CIDRs priorities" default:"priority"` - PriorityDel string `help:"Delimiter for priorities in the field" default:" "` - Report bool `help:"Report only conflicted CIDRs"` + Files []string `arg:"" type:"existingfile" help:"Input file containing CIDRs in CSV or JSON format"` + CidrKey string `help:"Key/Colum of the CIDRs in the file" default:"cidr"` + PriorityKeys []string `help:"Keys/Columns to be used as CIDRs priorities" default:""` + FillEmptyPriority bool `help:"Replace empty/null priority with zero value" default:"true"` + FlipRankPriority bool `help:"Make low value priority mean higher priority" default:"false"` + Report bool `help:"Report only conflicted CIDRs"` + + OutputFormat string `enum:"json,csv,tsv" default:"csv" help:"Output file format" default:"csv"` + DropKeys []string `help:"Keys/Columns to be dropped" default:""` + SplitIpVersions bool `help:"Split the results in to separate files based on the CIDR IP version" default:"false"` + Stats Stats `kong:"-"` } // Run executes the resolve command. func (cmd *ResolveCmd) Run(ctx *Context) error { + cmd.Stats.StartInsertTime = time.Now() + + // we read each record and insert it in supernet for _, file := range cmd.Files { if err := parseAndInsertCidrs(ctx.super, cmd, file); err != nil { return err } - if err := writeCsvResults(ctx.super, ".", cmd.CidrKey); err != nil { - - } } - - return nil -} - -// parseAndInsertCidrs parses a file and inserts CIDRs into the supernet. -func parseAndInsertCidrs(super *supernet.Supernet, cmd *ResolveCmd, file string) error { - return parseCsv(cmd, file, func(cidr *CIDR) error { - super.InsertCidr(cidr.cidr, cidr.Metadata) - return nil - }) -} - -// writeResults writes the results of CIDR resolution to a JSON file. -func writeJsonResults(super *supernet.Supernet, directory string, cidrCol string) error { - filePath := directory + "/resolved.json" - file, err := os.Create(filePath) - if err != nil { - return err + cmd.Stats.EndInsertTime = time.Now() + + // write back the resolved cidrs to file + var writer Writer + switch cmd.OutputFormat { + case "csv": + writer = &CsvWriter{splitIpVersions: cmd.SplitIpVersions, Stats: &cmd.Stats} + case "tsv": + writer = &CsvWriter{splitIpVersions: cmd.SplitIpVersions, isTSV: true, Stats: &cmd.Stats} + case "json": + writer = &JsonWriter{splitIpVersions: cmd.SplitIpVersions, Stats: &cmd.Stats} + default: + return fmt.Errorf("--output-format %s is not supported, please uses one of the following: [json,csv,tsv]", cmd.OutputFormat) } - defer file.Close() - encoder := json.NewEncoder(file) - cidrs := super.AllCIDRS(false) + cmd.Stats.StartOutputTime = time.Now() - fmt.Println("Starting to write resolved CIDRs...") - if _, err = file.Write([]byte("[")); err != nil { - return err - } + var err error + if cmd.SplitIpVersions { + err = writer.IsIpV6(false).Write(ctx.super, ".", cmd.CidrKey, cmd.DropKeys) + err = writer.IsIpV6(true).Write(ctx.super, ".", cmd.CidrKey, cmd.DropKeys) + } else { - for i, cidr := range cidrs { - cidr.Metadata().Attributes[cidrCol] = supernet.NodeToCidr(cidr) - if i > 0 { - if _, err = file.Write([]byte(",")); err != nil { - return err - } - } - if err = encoder.Encode(cidr.Metadata().Attributes); err != nil { - return err - } + err = writer.Write(ctx.super, ".", cmd.CidrKey, cmd.DropKeys) } - - if _, err = file.Write([]byte("]")); err != nil { + if err != nil { return err } - fmt.Println("Writing complete.") + cmd.Stats.EndOutputTime = time.Now() + printStats(cmd.Stats) return nil } -// writeResults writes the results of CIDR resolution to a CSV file. -func writeCsvResults(super *supernet.Supernet, directory string, cidrCol string) error { - filePath := directory + "/resolved.csv" - file, err := os.Create(filePath) - if err != nil { - return err - } - defer file.Close() - - // Create a CSV writer - writer := csv.NewWriter(file) - defer writer.Flush() - - cidrs := super.AllCIDRS(false) - - fmt.Println("Starting to write resolved CIDRs...") - - // Optional: Write headers to the CSV file - headers := []string{} - for key := range cidrs[0].Metadata().Attributes { - headers = append(headers, key) - } - if err := writer.Write(headers); err != nil { - return err +// parseAndInsertCidrs parses a file and inserts CIDRs into the supernet. +func parseAndInsertCidrs(super *supernet.Supernet, cmd *ResolveCmd, file string) error { + var parser CidrParser + extension := filepath.Ext(file) + switch extension { + case ".json": + parser = &JsonParser{} + case ".csv": + parser = &CsvCidrParser{} + case ".tsv": + parser = &CsvCidrParser{isTSV: true} + default: + { + return fmt.Errorf("File type %s is not supported, please use one of the following [json,csv,tsv]", extension) + } } - // Write data to the CSV file - for _, cidr := range cidrs { - cidr.Metadata().Attributes[cidrCol] = supernet.NodeToCidr(cidr) - record := make([]string, 0, len(cidr.Metadata().Attributes)) - // Ensure the fields are written in the same order as headers - for _, header := range headers { - record = append(record, cidr.Metadata().Attributes[header]) - } - if err := writer.Write(record); err != nil { - return err + return parser.Parse(cmd, file, func(cidr *CIDR) error { + result := super.InsertCidr(cidr.cidr, cidr.Metadata) + if _, noConflict := result.ConflictType.(supernet.NoConflict); noConflict { + cmd.Stats.Conflicted++ } - } + cmd.Stats.Input++ + return nil + }) +} - fmt.Println("Writing complete.") - return nil +func printStats(stats Stats) { + fmt.Printf("CIDRs Inserted:\t\t\t\t%d\nCIDRs With Conflicts:\t\t\t%d\nTotal CIDRs After Conflict Resolution:\t%d\n", stats.Input, stats.Conflicted, stats.Output) + fmt.Printf("Conflict Resolution Duration:\t\t%f Sec\n", stats.EndInsertTime.Sub(stats.StartInsertTime).Seconds()) + fmt.Printf("Writing Results Duration:\t\t%f Sec\n", stats.EndOutputTime.Sub(stats.StartOutputTime).Seconds()) + fmt.Printf("Total Time:\t\t\t\t%f Sec\n", stats.EndOutputTime.Sub(stats.StartInsertTime).Seconds()) }