diff --git a/csv.go b/csv.go new file mode 100644 index 0000000..36bd8cc --- /dev/null +++ b/csv.go @@ -0,0 +1,353 @@ +package main + +import ( + "encoding/csv" + "fmt" + "math" + "os" + "time" +) + +type csvPrinter struct { + writer *csv.Writer + file *os.File + dataFilename string + headerDone bool + showTimestamp *bool + showLocalAddress *bool + statsWriter *csv.Writer + statsFile *os.File + statsFilename string + statsHeaderDone bool + cleanup func() +} + +const ( + colStatus = "Status" + colTimestamp = "Timestamp" + colHostname = "Hostname" + colIP = "IP" + colPort = "Port" + colTCPConn = "TCP_Conn" + colLatency = "Latency(ms)" + colLocalAddress = "Local Address" +) + +func ensureCSVExtension(filename string) string { + if len(filename) > 4 && filename[len(filename)-4:] == ".csv" { + return filename + } + return filename + ".csv" +} + +func newCSVPrinter(dataFilename string, showTimestamp *bool, showLocalAddress *bool) (*csvPrinter, error) { + // Ensure .csv extension for dataFilename + dataFilename = ensureCSVExtension(dataFilename) + + // Open the data file with the os.O_TRUNC flag to truncate it + file, err := os.OpenFile(dataFilename, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0644) + if err != nil { + return nil, fmt.Errorf("error creating data CSV file: %w", err) + } + + // Append _stats before the .csv extension for statsFilename + statsFilename := dataFilename[:len(dataFilename)-4] + "_stats.csv" + cp := &csvPrinter{ + writer: csv.NewWriter(file), + file: file, + dataFilename: dataFilename, + statsFilename: statsFilename, + showTimestamp: showTimestamp, + showLocalAddress: showLocalAddress, + } + + cp.cleanup = func() { + if cp.writer != nil { + cp.writer.Flush() + } + if cp.file != nil { + cp.file.Close() + } + if cp.statsWriter != nil { + cp.statsWriter.Flush() + } + if cp.statsFile != nil { + cp.statsFile.Close() + } + } + + return cp, nil +} + +func (cp *csvPrinter) writeHeader() error { + headers := []string{ + colStatus, + colHostname, + colIP, + colPort, + colTCPConn, + colLatency, + } + + if *cp.showLocalAddress { + headers = append(headers, colLocalAddress) + } + + if *cp.showTimestamp { + headers = append(headers, colTimestamp) + } + + if err := cp.writer.Write(headers); err != nil { + return fmt.Errorf("failed to write headers: %w", err) + } + + cp.writer.Flush() + return cp.writer.Error() +} + +func (cp *csvPrinter) writeRecord(record []string) error { + if _, err := os.Stat(cp.dataFilename); os.IsNotExist(err) { + file, err := os.OpenFile(cp.dataFilename, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + if err != nil { + return fmt.Errorf("failed to recreate data CSV file: %w", err) + } + cp.file = file + cp.writer = csv.NewWriter(file) + cp.headerDone = false + } + + if !cp.headerDone { + if err := cp.writeHeader(); err != nil { + return err + } + cp.headerDone = true + } + + if *cp.showTimestamp { + record = append(record, time.Now().Format(timeFormat)) + } + + if err := cp.writer.Write(record); err != nil { + return fmt.Errorf("failed to write record: %w", err) + } + + cp.writer.Flush() + return cp.writer.Error() +} + +func (cp *csvPrinter) printStart(hostname string, port uint16) { + fmt.Printf("TCPing results being written to: %s\n", cp.dataFilename) +} + +func (cp *csvPrinter) printProbeSuccess(localAddr string, userInput userInput, streak uint, rtt float32) { + hostname := userInput.hostname + if hostname == "" { + hostname = "-" + } + + record := []string{ + "Reply", + hostname, + userInput.ip.String(), + fmt.Sprint(userInput.port), + fmt.Sprint(streak), + fmt.Sprintf("%.3f", rtt), + } + + if *cp.showLocalAddress { + record = append(record, localAddr) + } + + if err := cp.writeRecord(record); err != nil { + cp.printError("Failed to write success record: %v", err) + } +} + +func (cp *csvPrinter) printProbeFail(userInput userInput, streak uint) { + hostname := userInput.hostname + if hostname == "" { + hostname = "-" + } + + record := []string{ + "No reply", + hostname, + userInput.ip.String(), + fmt.Sprint(userInput.port), + fmt.Sprint(streak), + "-", + } + + if err := cp.writeRecord(record); err != nil { + cp.printError("Failed to write failure record: %v", err) + } +} + +func (cp *csvPrinter) printRetryingToResolve(hostname string) { + record := []string{ + "Resolving", + hostname, + "-", + "-", + "-", + "-", + } + + if err := cp.writeRecord(record); err != nil { + cp.printError("Failed to write resolve record: %v", err) + } +} + +func (cp *csvPrinter) printError(format string, args ...any) { + fmt.Fprintf(os.Stderr, "CSV Error: "+format+"\n", args...) +} + +func (cp *csvPrinter) writeStatsHeader() error { + + headers := []string{ + "Metric", + "Value", + } + + if err := cp.statsWriter.Write(headers); err != nil { + return fmt.Errorf("failed to write statistics headers: %w", err) + } + + cp.statsWriter.Flush() + return cp.statsWriter.Error() +} + +func (cp *csvPrinter) writeStatsRecord(record []string) error { + if _, err := os.Stat(cp.statsFilename); os.IsNotExist(err) { + statsFile, err := os.OpenFile(cp.statsFilename, os.O_CREATE|os.O_WRONLY|os.O_APPEND, 0644) + if err != nil { + return fmt.Errorf("failed to recreate statistics CSV file: %w", err) + } + cp.statsFile = statsFile + cp.statsWriter = csv.NewWriter(statsFile) + cp.statsHeaderDone = false + } + + // Write header if not done + if !cp.statsHeaderDone { + if err := cp.writeStatsHeader(); err != nil { + return err + } + cp.statsHeaderDone = true + } + + if err := cp.statsWriter.Write(record); err != nil { + return fmt.Errorf("failed to write statistics record: %w", err) + } + + cp.statsWriter.Flush() + return cp.statsWriter.Error() +} + +func (cp *csvPrinter) printStatistics(t tcping) { + // Initialize stats file if not already done + if cp.statsFile == nil { + statsFile, err := os.OpenFile(cp.statsFilename, os.O_CREATE|os.O_WRONLY|os.O_APPEND|os.O_TRUNC, 0644) + if err != nil { + cp.printError("Failed to create statistics CSV file: %v", err) + return + } + cp.statsFile = statsFile + cp.statsWriter = csv.NewWriter(statsFile) + cp.statsHeaderDone = false + } + + totalPackets := t.totalSuccessfulProbes + t.totalUnsuccessfulProbes + packetLoss := (float32(t.totalUnsuccessfulProbes) / float32(totalPackets)) * 100 + if math.IsNaN(float64(packetLoss)) { + packetLoss = 0 + } + + // Collect statistics data + timestamp := time.Now().Format(time.RFC3339) + statistics := [][]string{ + {"Timestamp", timestamp}, + {"Total Packets", fmt.Sprint(totalPackets)}, + {"Successful Probes", fmt.Sprint(t.totalSuccessfulProbes)}, + {"Unsuccessful Probes", fmt.Sprint(t.totalUnsuccessfulProbes)}, + {"Packet Loss", fmt.Sprintf("%.2f%%", packetLoss)}, + } + + if t.lastSuccessfulProbe.IsZero() { + statistics = append(statistics, []string{"Last Successful Probe", "Never succeeded"}) + } else { + statistics = append(statistics, []string{"Last Successful Probe", t.lastSuccessfulProbe.Format(timeFormat)}) + } + + if t.lastUnsuccessfulProbe.IsZero() { + statistics = append(statistics, []string{"Last Unsuccessful Probe", "Never failed"}) + } else { + statistics = append(statistics, []string{"Last Unsuccessful Probe", t.lastUnsuccessfulProbe.Format(timeFormat)}) + } + + statistics = append(statistics, []string{"Total Uptime", durationToString(t.totalUptime)}) + statistics = append(statistics, []string{"Total Downtime", durationToString(t.totalDowntime)}) + + if t.longestUptime.duration != 0 { + statistics = append(statistics, + []string{"Longest Uptime Duration", durationToString(t.longestUptime.duration)}, + []string{"Longest Uptime From", t.longestUptime.start.Format(timeFormat)}, + []string{"Longest Uptime To", t.longestUptime.end.Format(timeFormat)}, + ) + } + + if t.longestDowntime.duration != 0 { + statistics = append(statistics, + []string{"Longest Downtime Duration", durationToString(t.longestDowntime.duration)}, + []string{"Longest Downtime From", t.longestDowntime.start.Format(timeFormat)}, + []string{"Longest Downtime To", t.longestDowntime.end.Format(timeFormat)}, + ) + } + + if !t.destIsIP { + statistics = append(statistics, []string{"Retried Hostname Lookups", fmt.Sprint(t.retriedHostnameLookups)}) + + if len(t.hostnameChanges) >= 2 { + for i := 0; i < len(t.hostnameChanges)-1; i++ { + statistics = append(statistics, + []string{"IP Change", t.hostnameChanges[i].Addr.String()}, + []string{"To", t.hostnameChanges[i+1].Addr.String()}, + []string{"At", t.hostnameChanges[i+1].When.Format(timeFormat)}, + ) + } + } + } + + if t.rttResults.hasResults { + statistics = append(statistics, + []string{"RTT Min", fmt.Sprintf("%.3f ms", t.rttResults.min)}, + []string{"RTT Avg", fmt.Sprintf("%.3f ms", t.rttResults.average)}, + []string{"RTT Max", fmt.Sprintf("%.3f ms", t.rttResults.max)}, + ) + } + + statistics = append(statistics, []string{"TCPing Started At", t.startTime.Format(timeFormat)}) + + if !t.endTime.IsZero() { + statistics = append(statistics, []string{"TCPing Ended At", t.endTime.Format(timeFormat)}) + } + + durationTime := time.Time{}.Add(t.totalDowntime + t.totalUptime) + statistics = append(statistics, []string{"Duration (HH:MM:SS)", durationTime.Format(hourFormat)}) + + // Write statistics to CSV + for _, record := range statistics { + if err := cp.writeStatsRecord(record); err != nil { + cp.printError("Failed to write statistics record: %v", err) + return + } + } + + // Print the message only if statistics are written + fmt.Printf("TCPing statistics written to: %s\n", cp.statsFilename) +} + +// Satisfying remaining printer interface methods +func (cp *csvPrinter) printTotalDownTime(downtime time.Duration) {} +func (cp *csvPrinter) printVersion() {} +func (cp *csvPrinter) printInfo(format string, args ...any) {} diff --git a/csv_test.go b/csv_test.go new file mode 100644 index 0000000..0650078 --- /dev/null +++ b/csv_test.go @@ -0,0 +1,132 @@ +package main + +import ( + "encoding/csv" + "os" + "testing" + "time" + + "github.com/stretchr/testify/assert" +) + +func TestNewCSVPrinter(t *testing.T) { + dataFilename := "test_data.csv" + showTimestamp := true + showLocalAddress := true + + cp, err := newCSVPrinter(dataFilename, &showTimestamp, &showLocalAddress) + assert.NoError(t, err) + assert.NotNil(t, cp) + assert.Equal(t, dataFilename, cp.dataFilename) + assert.Equal(t, dataFilename[:len(dataFilename)-4]+"_stats.csv", cp.statsFilename) + + cp.cleanup() + os.Remove(dataFilename) + os.Remove(cp.statsFilename) +} + +func TestWriteRecord(t *testing.T) { + dataFilename := "test_data.csv" + showTimestamp := false + showLocalAddress := true + + cp, err := newCSVPrinter(dataFilename, &showTimestamp, &showLocalAddress) + assert.NoError(t, err) + assert.NotNil(t, cp) + + record := []string{"Success", "hostname", "127.0.0.1", "80", "1", "10.123", "localAddr"} + err = cp.writeRecord(record) + assert.NoError(t, err) + + // Verify the record is written + file, err := os.Open(dataFilename) + assert.NoError(t, err) + defer file.Close() + + reader := csv.NewReader(file) + headers, err := reader.Read() + assert.NoError(t, err) + assert.Equal(t, []string{"Status", "Hostname", "IP", "Port", "TCP_Conn", "Latency(ms)", "Local Address"}, headers) + + readRecord, err := reader.Read() + assert.NoError(t, err) + assert.Equal(t, record, readRecord) + + // Cleanup + cp.cleanup() + os.Remove(dataFilename) + os.Remove(cp.statsFilename) +} + +func TestWriteStatistics(t *testing.T) { + dataFilename := "test_data.csv" + showTimestamp := true + showLocalAddress := false + + cp, err := newCSVPrinter(dataFilename, &showTimestamp, &showLocalAddress) + assert.NoError(t, err) + assert.NotNil(t, cp) + + tcping := tcping{ + totalSuccessfulProbes: 1, + totalUnsuccessfulProbes: 0, + lastSuccessfulProbe: time.Now(), + startTime: time.Now(), + } + + cp.printStatistics(tcping) + + statsFile, err := os.Open(cp.statsFilename) + assert.NoError(t, err) + defer statsFile.Close() + + reader := csv.NewReader(statsFile) + headers, err := reader.Read() + assert.NoError(t, err) + assert.Equal(t, []string{"Metric", "Value"}, headers) + + for { + record, err := reader.Read() + if err != nil { + break + } + assert.NotEmpty(t, record) + } + + cp.cleanup() + os.Remove(dataFilename) + os.Remove(cp.statsFilename) +} + +func TestCleanup(t *testing.T) { + dataFilename := "test_data.csv" + showTimestamp := true + showLocalAddress := false + + cp, err := newCSVPrinter(dataFilename, &showTimestamp, &showLocalAddress) + assert.NoError(t, err) + assert.NotNil(t, cp) + + // Call printStatistics to ensure the stats file is created + tcping := tcping{ + totalSuccessfulProbes: 1, + totalUnsuccessfulProbes: 0, + lastSuccessfulProbe: time.Now(), + startTime: time.Now(), + } + cp.printStatistics(tcping) + + // Perform cleanup + cp.cleanup() + + // Verify files are closed and flushed + _, err = os.Stat(dataFilename) + assert.NoError(t, err) + + _, err = os.Stat(cp.statsFilename) + assert.NoError(t, err) + + // Cleanup files + os.Remove(dataFilename) + os.Remove(cp.statsFilename) +} diff --git a/tcping.go b/tcping.go index 08c84f7..020edf7 100644 --- a/tcping.go +++ b/tcping.go @@ -14,7 +14,6 @@ import ( "strings" "syscall" "time" - "github.com/google/go-github/v45/github" ) @@ -199,11 +198,16 @@ func shutdown(tcping *tcping) { tcping.endTime = time.Now() tcping.printStats() - // if the printer type is `database`, close the it before exiting + // if the printer type is `database`, close it before exiting if db, ok := tcping.printer.(*database); ok { db.conn.Close() } + // if the printer type is `csvPrinter`, call the cleanup function before exiting + if cp, ok := tcping.printer.(*csvPrinter); ok { + cp.cleanup() + } + os.Exit(0) } @@ -230,7 +234,7 @@ func usage() { } // setPrinter selects the printer -func setPrinter(tcping *tcping, outputJSON, prettyJSON *bool, noColor *bool, timeStamp *bool, outputDb *string, args []string) { +func setPrinter(tcping *tcping, outputJSON, prettyJSON *bool, noColor *bool, timeStamp *bool, localAddress *bool, outputDb *string, outputCSV *string, args []string) { if *prettyJSON && !*outputJSON { colorRed("--pretty has no effect without the -j flag.") usage() @@ -239,7 +243,14 @@ func setPrinter(tcping *tcping, outputJSON, prettyJSON *bool, noColor *bool, tim tcping.printer = newJSONPrinter(*prettyJSON) } else if *outputDb != "" { tcping.printer = newDB(*outputDb, args) - } else if *noColor == true { + } else if *outputCSV != "" { + var err error + tcping.printer, err = newCSVPrinter(*outputCSV, timeStamp, localAddress) + if err != nil { + tcping.printError("Failed to create CSV file: %s", err) + os.Exit(1) + } + } else if *noColor { tcping.printer = newPlanePrinter(timeStamp) } else { tcping.printer = newColorPrinter(timeStamp) @@ -331,6 +342,7 @@ func processUserInput(tcping *tcping) { prettyJSON := flag.Bool("pretty", false, "use indentation when using json output format. No effect without the '-j' flag.") noColor := flag.Bool("no-color", false, "do not colorize output.") showTimestamp := flag.Bool("D", false, "show timestamp in output.") + saveToCSV := flag.String("csv", "", "path and file name to store tcping output to CSV file...If user prompts for stats, it will be saved to a file with the same name but _stats appended.") showVer := flag.Bool("v", false, "show version.") checkUpdates := flag.Bool("u", false, "check for updates and exit.") secondsBetweenProbes := flag.Float64("i", 1, "interval between sending probes. Real number allowed with dot as a decimal separator. The default is one second") @@ -341,6 +353,7 @@ func processUserInput(tcping *tcping) { showFailuresOnly := flag.Bool("show-failures-only", false, "Show only the failed probes.") showHelp := flag.Bool("h", false, "show help message.") + flag.CommandLine.Usage = usage permuteArgs(os.Args[1:]) @@ -351,7 +364,7 @@ func processUserInput(tcping *tcping) { // we need to set printers first, because they're used for // error reporting and other output. - setPrinter(tcping, outputJSON, prettyJSON, noColor, showTimestamp, outputDB, args) + setPrinter(tcping, outputJSON, prettyJSON, noColor, showTimestamp, showLocalAddress, outputDB, saveToCSV, args) // Handle -v flag if *showVer { @@ -423,6 +436,8 @@ func permuteArgs(args []string) { fallthrough case "i": fallthrough + case "csv": + fallthrough case "r": /* out of index */ if len(args) <= i+1 {