Merge pull request #2 from meyskens/master

Split into packages to make importable

.gitignore

@@ -1,2 +1 @@
-mtr
 *.swp

cli.go → cli/cli.go

@@ -1,4 +1,4 @@
-package main
+package cli
 
 import (
 	"errors"
@@ -9,6 +9,7 @@ import (
 	tm "github.com/buger/goterm"
 	pj "github.com/hokaccha/go-prettyjson"
 	"github.com/spf13/cobra"
+	"github.com/tonobo/mtr/mtr"
 )
 
 var (
@@ -28,7 +29,7 @@ var RootCmd = &cobra.Command{
 		if len(args) != 1 {
 			return errors.New("No target provided")
 		}
-		m, ch := NewMTR(args[0], TIMEOUT, INTERVAL, HOP_SLEEP)
+		m, ch := mtr.NewMTR(args[0], TIMEOUT, INTERVAL, HOP_SLEEP, MAX_HOPS, RING_BUFFER_SIZE)
 		if jsonFmt {
 			go func(ch chan struct{}) {
 				for {
@@ -60,7 +61,7 @@ var RootCmd = &cobra.Command{
 	},
 }
 
-func render(m *MTR) {
+func render(m *mtr.MTR) {
 	tm.MoveCursor(1, 1)
 	m.Render(1)
 	tm.Flush() // Call it every time at the end of rendering

hop.go → hop/hop.go

@@ -1,4 +1,4 @@
-package main
+package hop
 
 import (
 	"container/ring"
@@ -8,21 +8,23 @@ import (
 	"time"
 
 	gm "github.com/buger/goterm"
+	"github.com/tonobo/mtr/imcp"
 )
 
 type HopStatistic struct {
-	dest       *net.IPAddr
-	timeout    time.Duration
-	pid        int
-	Sent       int
-	TTL        int
-	Target     string
-	Last       ICMPReturn
-	Best       ICMPReturn
-	Worst      ICMPReturn
-	SumElapsed time.Duration
-	Lost       int
-	Packets    *ring.Ring
+	Dest           *net.IPAddr
+	Timeout        time.Duration
+	PID            int
+	Sent           int
+	TTL            int
+	Target         string
+	Last           imcp.ICMPReturn
+	Best           imcp.ICMPReturn
+	Worst          imcp.ICMPReturn
+	SumElapsed     time.Duration
+	Lost           int
+	Packets        *ring.Ring
+	RingBufferSize int
 }
 
 type packet struct {
@@ -31,7 +33,7 @@ type packet struct {
 }
 
 func (s *HopStatistic) Next() {
-	r, _ := Icmp("0.0.0.0", s.dest, s.TTL, s.pid, s.timeout)
+	r, _ := imcp.SendIMCP("0.0.0.0", s.Dest, s.TTL, s.PID, s.Timeout)
 	s.Packets = s.Packets.Prev()
 	s.Packets.Value = r
 	if s.Target == "" {
@@ -68,7 +70,7 @@ func (h *HopStatistic) MarshalJSON() ([]byte, error) {
 		TTL:              h.TTL,
 		Loss:             h.Loss(),
 		Target:           h.Target,
-		PacketBufferSize: RING_BUFFER_SIZE,
+		PacketBufferSize: h.RingBufferSize,
 		Last:             h.Last.Elapsed.Seconds() * 1000,
 		Best:             h.Best.Elapsed.Seconds() * 1000,
 		Worst:            h.Worst.Elapsed.Seconds() * 1000,
@@ -90,15 +92,15 @@ func (h *HopStatistic) Loss() float64 {
 }
 
 func (h *HopStatistic) packets() []*packet {
-	v := make([]*packet, RING_BUFFER_SIZE)
+	v := make([]*packet, h.RingBufferSize)
 	i := 0
 	h.Packets.Do(func(f interface{}) {
 		if f == nil {
 			v[i] = nil
 			i++
 			return
 		}
-		x := f.(ICMPReturn)
+		x := f.(imcp.ICMPReturn)
 		if x.Success {
 			v[i] = &packet{
 				Success:      true,
@@ -116,12 +118,12 @@ func (h *HopStatistic) packets() []*packet {
 }
 
 func (h *HopStatistic) Render() {
-	packets := make([]byte, RING_BUFFER_SIZE)
-	i := RING_BUFFER_SIZE - 1
+	packets := make([]byte, h.RingBufferSize)
+	i := h.RingBufferSize - 1
 	h.Packets.Do(func(f interface{}) {
 		if f == nil {
 			packets[i] = ' '
-		} else if !f.(ICMPReturn).Success {
+		} else if !f.(imcp.ICMPReturn).Success {
 			packets[i] = '?'
 		} else {
 			packets[i] = '.'
@@ -132,7 +134,7 @@ func (h *HopStatistic) Render() {
 	if h.Target != "" {
 		addr = h.Target
 	}
-	l := fmt.Sprintf("%d", RING_BUFFER_SIZE)
+	l := fmt.Sprintf("%d", h.RingBufferSize)
 	gm.Printf("%3d:|-- %-20s  %5.1f%%  %4d  %6.1f  %6.1f  %6.1f  %6.1f  %"+l+"s\n",
 		h.TTL,
 		addr,

icmp.go → imcp/icmp.go

@@ -1,4 +1,4 @@
-package main
+package imcp
 
 import (
 	"net"
@@ -8,13 +8,15 @@ import (
 	"golang.org/x/net/ipv4"
 )
 
+// ICMPReturn contains the info for a returned IMCP
 type ICMPReturn struct {
 	Success bool
 	Addr    string
 	Elapsed time.Duration
 }
 
-func Icmp(localAddr string, dst net.Addr, ttl, pid int, timeout time.Duration) (hop ICMPReturn, err error) {
+// SendIMCP sends a IMCP to a given destination
+func SendIMCP(localAddr string, dst net.Addr, ttl, pid int, timeout time.Duration) (hop ICMPReturn, err error) {
 	hop.Success = false
 	start := time.Now()
 	c, err := icmp.ListenPacket("ip4:icmp", localAddr)

main.go

@@ -1,5 +1,14 @@
 package main
 
+import (
+	"fmt"
+
+	"github.com/tonobo/mtr/cli"
+)
+
 func main() {
-	RootCmd.Execute()
+	err := cli.RootCmd.Execute()
+	if err != nil {
+		fmt.Println(err)
+	}
 }

mtr.go → mtr/mtr.go

@@ -1,4 +1,4 @@
-package main
+package mtr
 
 import (
 	"container/ring"
@@ -9,40 +9,47 @@ import (
 	"time"
 
 	gm "github.com/buger/goterm"
+	"github.com/tonobo/mtr/hop"
+	"github.com/tonobo/mtr/imcp"
 )
 
 type MTR struct {
-	mutex     *sync.RWMutex
-	timeout   time.Duration
-	interval  time.Duration
-	Address   string `json:"destination"`
-	hopsleep  time.Duration
-	Statistic map[int]*HopStatistic `json:"statistic"`
+	mutex          *sync.RWMutex
+	timeout        time.Duration
+	interval       time.Duration
+	Address        string `json:"destination"`
+	hopsleep       time.Duration
+	Statistic      map[int]*hop.HopStatistic `json:"statistic"`
+	ringBufferSize int
+	maxHops        int
 }
 
-func NewMTR(addr string, timeout time.Duration, interval time.Duration, hopsleep time.Duration) (*MTR, chan struct{}) {
+func NewMTR(addr string, timeout time.Duration, interval time.Duration, hopsleep time.Duration, maxHops, ringBufferSize int) (*MTR, chan struct{}) {
 	return &MTR{
-		interval:  interval,
-		timeout:   timeout,
-		hopsleep:  hopsleep,
-		Address:   addr,
-		mutex:     &sync.RWMutex{},
-		Statistic: map[int]*HopStatistic{},
+		interval:       interval,
+		timeout:        timeout,
+		hopsleep:       hopsleep,
+		Address:        addr,
+		mutex:          &sync.RWMutex{},
+		Statistic:      map[int]*hop.HopStatistic{},
+		maxHops:        maxHops,
+		ringBufferSize: ringBufferSize,
 	}, make(chan struct{})
 }
 
-func (m *MTR) registerStatistic(ttl int, r ICMPReturn) *HopStatistic {
-	m.Statistic[ttl] = &HopStatistic{
-		Sent:       1,
-		TTL:        ttl,
-		Target:     r.Addr,
-		timeout:    m.timeout,
-		Last:       r,
-		Best:       r,
-		Worst:      r,
-		Lost:       0,
-		SumElapsed: r.Elapsed,
-		Packets:    ring.New(RING_BUFFER_SIZE),
+func (m *MTR) registerStatistic(ttl int, r imcp.ICMPReturn) *hop.HopStatistic {
+	m.Statistic[ttl] = &hop.HopStatistic{
+		Sent:           1,
+		TTL:            ttl,
+		Target:         r.Addr,
+		Timeout:        m.timeout,
+		Last:           r,
+		Best:           r,
+		Worst:          r,
+		Lost:           0,
+		SumElapsed:     r.Elapsed,
+		Packets:        ring.New(m.ringBufferSize),
+		RingBufferSize: m.ringBufferSize,
 	}
 	if !r.Success {
 		m.Statistic[ttl].Lost++
@@ -53,7 +60,7 @@ func (m *MTR) registerStatistic(ttl int, r ICMPReturn) *HopStatistic {
 
 func (m *MTR) Render(offset int) {
 	gm.MoveCursor(1, offset)
-	l := fmt.Sprintf("%d", RING_BUFFER_SIZE)
+	l := fmt.Sprintf("%d", m.ringBufferSize)
 	gm.Printf("HOP:    %-20s  %5s%%  %4s  %6s  %6s  %6s  %6s  %"+l+"s\n", "Address", "Loss", "Sent", "Last", "Avg", "Best", "Worst", "Packets")
 	for i := 1; i <= len(m.Statistic); i++ {
 		gm.MoveCursor(1, offset+i)
@@ -86,17 +93,17 @@ func (m *MTR) discover(ch chan struct{}) {
 	ipAddr := net.IPAddr{IP: net.ParseIP(m.Address)}
 	pid := os.Getpid() & 0xffff
 	ttlDoubleBump := false
-	for ttl := 1; ttl < MAX_HOPS; ttl++ {
+	for ttl := 1; ttl < m.maxHops; ttl++ {
 		time.Sleep(m.hopsleep)
-		hopReturn, err := Icmp("0.0.0.0", &ipAddr, ttl, pid, m.timeout)
+		hopReturn, err := imcp.SendIMCP("0.0.0.0", &ipAddr, ttl, pid, m.timeout)
 		if err != nil || !hopReturn.Success {
 			if ttlDoubleBump {
 				break
 			}
 			m.mutex.Lock()
 			s := m.registerStatistic(ttl, hopReturn)
-			s.dest = &ipAddr
-			s.pid = pid
+			s.Dest = &ipAddr
+			s.PID = pid
 			m.mutex.Unlock()
 			ch <- struct{}{}
 			ttlDoubleBump = true
@@ -105,8 +112,8 @@ func (m *MTR) discover(ch chan struct{}) {
 		ttlDoubleBump = false
 		m.mutex.Lock()
 		s := m.registerStatistic(ttl, hopReturn)
-		s.dest = &ipAddr
-		s.pid = pid
+		s.Dest = &ipAddr
+		s.PID = pid
 		m.mutex.Unlock()
 		ch <- struct{}{}
 		if hopReturn.Addr == m.Address {