From 6c5b090ed845d8e7b1164c15f61d2a1c97b9773d Mon Sep 17 00:00:00 2001
From: Tim Foerster <tim.foerster@hetzner.de>
Date: Tue, 19 Dec 2017 23:23:31 +0100
Subject: [PATCH] Update rendering and add json option

Signed-off-by: Tim Foerster <tim.foerster@hetzner.de>
---
 cli.go | 18 +++++++++++++++++-
 mtr.go |  9 +++++++--
 2 files changed, 24 insertions(+), 3 deletions(-)

diff --git a/cli.go b/cli.go
index a4f99df..eee6d81 100644
--- a/cli.go
+++ b/cli.go
@@ -7,6 +7,7 @@ import (
 	"time"
 
 	tm "github.com/buger/goterm"
+	pj "github.com/hokaccha/go-prettyjson"
 	"github.com/spf13/cobra"
 )
 
@@ -16,6 +17,7 @@ var (
 	INTERVAL         = 100 * time.Millisecond
 	MAX_HOPS         = 64
 	RING_BUFFER_SIZE = 50
+	jsonFmt          = false
 )
 
 // rootCmd represents the root command
@@ -25,8 +27,21 @@ var RootCmd = &cobra.Command{
 		if len(args) != 1 {
 			return errors.New("No target provided")
 		}
-		fmt.Println("Start:", time.Now())
 		m, ch := NewMTR(args[0], TIMEOUT, INTERVAL)
+		if jsonFmt {
+			go func(ch chan struct{}) {
+				for {
+					<-ch
+				}
+			}(ch)
+			for i := 0; i < COUNT; i++ {
+				m.Run(ch)
+			}
+			s, _ := pj.Marshal(m)
+			fmt.Println(string(s))
+			return nil
+		}
+		fmt.Println("Start:", time.Now())
 		tm.Clear()
 		mu := &sync.Mutex{}
 		go func(ch chan struct{}) {
@@ -60,4 +75,5 @@ func init() {
 	RootCmd.Flags().DurationVarP(&INTERVAL, "interval", "i", INTERVAL, "Wait time between icmp packets before sending new one")
 	RootCmd.Flags().IntVar(&MAX_HOPS, "max-hops", MAX_HOPS, "Maximal TTL count")
 	RootCmd.Flags().IntVar(&RING_BUFFER_SIZE, "buffer-size", RING_BUFFER_SIZE, "Cached packet buffer size")
+	RootCmd.Flags().BoolVar(&jsonFmt, "json", jsonFmt, "Print json results")
 }
diff --git a/mtr.go b/mtr.go
index d878328..0f4e09d 100644
--- a/mtr.go
+++ b/mtr.go
@@ -33,6 +33,8 @@ func (m *MTR) registerStatistic(ttl int, r ICMPReturn) {
 	if m.Statistic[ttl] == nil {
 		m.Statistic[ttl] = &HopStatistic{
 			Sent:       1,
+			TTL:        ttl,
+			Target:     r.Addr,
 			Last:       r,
 			Best:       r,
 			Worst:      r,
@@ -47,7 +49,11 @@ func (m *MTR) registerStatistic(ttl int, r ICMPReturn) {
 		return
 	}
 	s := m.Statistic[ttl]
+	s.Packets = s.Packets.Prev()
 	s.Packets.Value = r
+	if s.Target == "" {
+		s.Target = r.Addr
+	}
 	s.Sent++
 	s.SumElapsed = r.Elapsed + s.SumElapsed
 	if !r.Success {
@@ -60,7 +66,6 @@ func (m *MTR) registerStatistic(ttl int, r ICMPReturn) {
 	if s.Worst.Elapsed < r.Elapsed {
 		s.Worst = r
 	}
-	s.Packets = s.Packets.Next()
 }
 
 func (m *MTR) Render(offset int) {
@@ -70,7 +75,7 @@ func (m *MTR) Render(offset int) {
 	for i := 1; i <= len(m.Statistic); i++ {
 		gm.MoveCursor(1, offset+i)
 		m.mutex.RLock()
-		m.Statistic[i].Render(i)
+		m.Statistic[i].Render()
 		m.mutex.RUnlock()
 	}
 	return