Skip to content

Commit

Permalink
Refactor logging
Browse files Browse the repository at this point in the history
  • Loading branch information
v-byte-cpu committed Mar 10, 2021
1 parent e8843e4 commit 51e4c2b
Show file tree
Hide file tree
Showing 6 changed files with 249 additions and 35 deletions.
47 changes: 12 additions & 35 deletions command/arp.go
Original file line number Diff line number Diff line change
@@ -1,7 +1,6 @@
package command

import (
"bufio"
"context"
"errors"
"net"
Expand All @@ -12,11 +11,11 @@ import (
"time"

"github.com/spf13/cobra"
"github.com/v-byte-cpu/sx/command/log"
"github.com/v-byte-cpu/sx/pkg/ip"
"github.com/v-byte-cpu/sx/pkg/packet/afpacket"
"github.com/v-byte-cpu/sx/pkg/scan"
"github.com/v-byte-cpu/sx/pkg/scan/arp"
"go.uber.org/zap"
)

var errSrcIP = errors.New("invalid source IP")
Expand Down Expand Up @@ -77,44 +76,22 @@ var arpCmd = &cobra.Command{
}
}

r := &scan.Range{Subnet: dstSubnet, Interface: iface, SrcIP: srcIP.To4(), SrcMAC: srcMAC}
return startEngine(r)
},
}

func logResults(logger *zap.Logger, results <-chan *arp.ScanResult) {
bw := bufio.NewWriter(os.Stdout)
defer bw.Flush()
for result := range results {
// TODO refactor it using logger facade interface
var logger log.Logger
if jsonFlag {
data, err := result.MarshalJSON()
if err != nil {
logger.Error("arp", zap.Error(err))
}
_, err = bw.Write(data)
if err != nil {
logger.Error("arp", zap.Error(err))
}
logger, err = log.NewJSONLogger(os.Stdout, "arp")
} else {
_, err := bw.WriteString(result.String())
if err != nil {
logger.Error("arp", zap.Error(err))
}
logger, err = log.NewPlainLogger(os.Stdout, "arp")
}
err := bw.WriteByte('\n')
if err != nil {
logger.Error("arp", zap.Error(err))
return err
}
}
}

func startEngine(r *scan.Range) error {
logger, err := zap.NewProduction()
if err != nil {
return err
}
r := &scan.Range{Subnet: dstSubnet, Interface: iface, SrcIP: srcIP.To4(), SrcMAC: srcMAC}
return startEngine(logger, r)
},
}

func startEngine(logger log.Logger, r *scan.Range) error {
ctx, cancel := signal.NotifyContext(context.Background(), os.Interrupt)
defer cancel()

Expand All @@ -136,7 +113,7 @@ func startEngine(r *scan.Range) error {
wg.Add(1)
go func() {
defer wg.Done()
logResults(logger, m.Results())
logger.LogResults(m.Results())
}()

// start scan
Expand All @@ -153,7 +130,7 @@ func startEngine(r *scan.Range) error {
go func() {
defer wg.Done()
for err := range errc {
logger.Error("arp", zap.Error(err))
logger.Error(err)
}
}()
wg.Wait()
Expand Down
45 changes: 45 additions & 0 deletions command/log/json_logger.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,45 @@
package log

import (
"bufio"
"io"

"github.com/v-byte-cpu/sx/pkg/scan/arp"
"go.uber.org/zap"
)

type JSONLogger struct {
logger *zap.Logger
writer io.Writer
label string
}

func NewJSONLogger(w io.Writer, label string) (*JSONLogger, error) {
logger, err := zap.NewProduction()
if err != nil {
return nil, err
}
return &JSONLogger{logger, w, label}, nil
}

func (l *JSONLogger) Error(err error) {
l.logger.Error(l.label, zap.Error(err))
}

func (l *JSONLogger) LogResults(results <-chan *arp.ScanResult) {
bw := bufio.NewWriter(l.writer)
defer bw.Flush()
var err error
var data []byte
for result := range results {
if data, err = result.MarshalJSON(); err != nil {
l.Error(err)
}
if _, err = bw.Write(data); err != nil {
l.Error(err)
}
if err = bw.WriteByte('\n'); err != nil {
l.Error(err)
}
}
}
71 changes: 71 additions & 0 deletions command/log/json_logger_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
package log

import (
"bytes"
"net"
"strings"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/v-byte-cpu/sx/pkg/scan/arp"
)

func scanResultToJSON(t *testing.T, result *arp.ScanResult) string {
t.Helper()
data, err := result.MarshalJSON()
require.NoError(t, err)
return string(data)
}

func TestJSONLoggerResults(t *testing.T) {
t.Parallel()

tests := []struct {
name string
expected []byte
results []*arp.ScanResult
}{
{
name: "emptyResults",
expected: nil,
results: nil,
},
{
name: "oneResult",
expected: []byte(scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 3).To4())) + "\n"),
results: []*arp.ScanResult{
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
},
},
{
name: "twoResults",
expected: []byte(strings.Join([]string{
scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 3).To4())),
scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 5).To4())),
}, "\n") + "\n"),
results: []*arp.ScanResult{
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
newScanResult(net.IPv4(192, 168, 0, 5).To4()),
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {

var buf bytes.Buffer
logger, err := NewJSONLogger(&buf, "arp")
require.NoError(t, err)

resultCh := make(chan *arp.ScanResult, len(tt.results))
for _, result := range tt.results {
resultCh <- result
}
close(resultCh)
logger.LogResults(resultCh)

assert.Equal(t, tt.expected, buf.Bytes())
})
}
}
8 changes: 8 additions & 0 deletions command/log/logger.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,8 @@
package log

import "github.com/v-byte-cpu/sx/pkg/scan/arp"

type Logger interface {
Error(err error)
LogResults(results <-chan *arp.ScanResult)
}
41 changes: 41 additions & 0 deletions command/log/plain_logger.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,41 @@
package log

import (
"bufio"
"io"

"github.com/v-byte-cpu/sx/pkg/scan/arp"
"go.uber.org/zap"
)

type PlainLogger struct {
logger *zap.Logger
writer io.Writer
label string
}

func NewPlainLogger(w io.Writer, label string) (*PlainLogger, error) {
logger, err := zap.NewProduction()
if err != nil {
return nil, err
}
return &PlainLogger{logger, w, label}, nil
}

func (l *PlainLogger) Error(err error) {
l.logger.Error(l.label, zap.Error(err))
}

func (l *PlainLogger) LogResults(results <-chan *arp.ScanResult) {
bw := bufio.NewWriter(l.writer)
defer bw.Flush()
var err error
for result := range results {
if _, err = bw.WriteString(result.String()); err != nil {
l.Error(err)
}
if err = bw.WriteByte('\n'); err != nil {
l.Error(err)
}
}
}
72 changes: 72 additions & 0 deletions command/log/plain_logger_test.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,72 @@
package log

import (
"bytes"
"net"
"strings"
"testing"

"github.com/stretchr/testify/assert"
"github.com/stretchr/testify/require"
"github.com/v-byte-cpu/sx/pkg/scan/arp"
)

func newScanResult(ip net.IP) *arp.ScanResult {
return &arp.ScanResult{
IP: ip.String(),
MAC: net.HardwareAddr{0x11, 0x22, 0x33, 0x44, 0x55, 0x66}.String(),
Vendor: "Sunny Industries",
}
}

func TestPlainLoggerResults(t *testing.T) {
t.Parallel()

tests := []struct {
name string
expected []byte
results []*arp.ScanResult
}{
{
name: "emptyResults",
expected: nil,
results: nil,
},
{
name: "oneResult",
expected: []byte(newScanResult(net.IPv4(192, 168, 0, 3).To4()).String() + "\n"),
results: []*arp.ScanResult{
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
},
},
{
name: "twoResults",
expected: []byte(strings.Join([]string{
newScanResult(net.IPv4(192, 168, 0, 3).To4()).String(),
newScanResult(net.IPv4(192, 168, 0, 5).To4()).String(),
}, "\n") + "\n"),
results: []*arp.ScanResult{
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
newScanResult(net.IPv4(192, 168, 0, 5).To4()),
},
},
}

for _, tt := range tests {
t.Run(tt.name, func(t *testing.T) {

var buf bytes.Buffer
logger, err := NewPlainLogger(&buf, "arp")
require.NoError(t, err)

resultCh := make(chan *arp.ScanResult, len(tt.results))
for _, result := range tt.results {
resultCh <- result
}
close(resultCh)
logger.LogResults(resultCh)

assert.Equal(t, tt.expected, buf.Bytes())
})
}
}

0 comments on commit 51e4c2b

Please sign in to comment.