Skip to content

Commit 51e4c2b

Browse files
v-byte-cpuv-byte-cpu
committed
Refactor logging
1 parent e8843e4 commit 51e4c2b

File tree

6 files changed

+249
-35
lines changed

6 files changed

+249
-35
lines changed

‎command/arp.go‎

Lines changed: 12 additions & 35 deletions
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,6 @@
11
package command
22

33
import (
4-
"bufio"
54
"context"
65
"errors"
76
"net"
@@ -12,11 +11,11 @@ import (
1211
"time"
1312

1413
"github.com/spf13/cobra"
14+
"github.com/v-byte-cpu/sx/command/log"
1515
"github.com/v-byte-cpu/sx/pkg/ip"
1616
"github.com/v-byte-cpu/sx/pkg/packet/afpacket"
1717
"github.com/v-byte-cpu/sx/pkg/scan"
1818
"github.com/v-byte-cpu/sx/pkg/scan/arp"
19-
"go.uber.org/zap"
2019
)
2120

2221
var errSrcIP = errors.New("invalid source IP")
@@ -77,44 +76,22 @@ var arpCmd = &cobra.Command{
7776
}
7877
}
7978

80-
r := &scan.Range{Subnet: dstSubnet, Interface: iface, SrcIP: srcIP.To4(), SrcMAC: srcMAC}
81-
return startEngine(r)
82-
},
83-
}
84-
85-
func logResults(logger *zap.Logger, results <-chan *arp.ScanResult) {
86-
bw := bufio.NewWriter(os.Stdout)
87-
defer bw.Flush()
88-
for result := range results {
89-
// TODO refactor it using logger facade interface
79+
var logger log.Logger
9080
if jsonFlag {
91-
data, err := result.MarshalJSON()
92-
if err != nil {
93-
logger.Error("arp", zap.Error(err))
94-
}
95-
_, err = bw.Write(data)
96-
if err != nil {
97-
logger.Error("arp", zap.Error(err))
98-
}
81+
logger, err = log.NewJSONLogger(os.Stdout, "arp")
9982
} else {
100-
_, err := bw.WriteString(result.String())
101-
if err != nil {
102-
logger.Error("arp", zap.Error(err))
103-
}
83+
logger, err = log.NewPlainLogger(os.Stdout, "arp")
10484
}
105-
err := bw.WriteByte('\n')
10685
if err != nil {
107-
logger.Error("arp", zap.Error(err))
86+
return err
10887
}
109-
}
110-
}
11188

112-
func startEngine(r *scan.Range) error {
113-
logger, err := zap.NewProduction()
114-
if err != nil {
115-
return err
116-
}
89+
r := &scan.Range{Subnet: dstSubnet, Interface: iface, SrcIP: srcIP.To4(), SrcMAC: srcMAC}
90+
return startEngine(logger, r)
91+
},
92+
}
11793

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

@@ -136,7 +113,7 @@ func startEngine(r *scan.Range) error {
136113
wg.Add(1)
137114
go func() {
138115
defer wg.Done()
139-
logResults(logger, m.Results())
116+
logger.LogResults(m.Results())
140117
}()
141118

142119
// start scan
@@ -153,7 +130,7 @@ func startEngine(r *scan.Range) error {
153130
go func() {
154131
defer wg.Done()
155132
for err := range errc {
156-
logger.Error("arp", zap.Error(err))
133+
logger.Error(err)
157134
}
158135
}()
159136
wg.Wait()

‎command/log/json_logger.go‎

Lines changed: 45 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,45 @@
1+
package log
2+
3+
import (
4+
"bufio"
5+
"io"
6+
7+
"github.com/v-byte-cpu/sx/pkg/scan/arp"
8+
"go.uber.org/zap"
9+
)
10+
11+
type JSONLogger struct {
12+
logger *zap.Logger
13+
writer io.Writer
14+
label string
15+
}
16+
17+
func NewJSONLogger(w io.Writer, label string) (*JSONLogger, error) {
18+
logger, err := zap.NewProduction()
19+
if err != nil {
20+
return nil, err
21+
}
22+
return &JSONLogger{logger, w, label}, nil
23+
}
24+
25+
func (l *JSONLogger) Error(err error) {
26+
l.logger.Error(l.label, zap.Error(err))
27+
}
28+
29+
func (l *JSONLogger) LogResults(results <-chan *arp.ScanResult) {
30+
bw := bufio.NewWriter(l.writer)
31+
defer bw.Flush()
32+
var err error
33+
var data []byte
34+
for result := range results {
35+
if data, err = result.MarshalJSON(); err != nil {
36+
l.Error(err)
37+
}
38+
if _, err = bw.Write(data); err != nil {
39+
l.Error(err)
40+
}
41+
if err = bw.WriteByte('\n'); err != nil {
42+
l.Error(err)
43+
}
44+
}
45+
}

‎command/log/json_logger_test.go‎

Lines changed: 71 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,71 @@
1+
package log
2+
3+
import (
4+
"bytes"
5+
"net"
6+
"strings"
7+
"testing"
8+
9+
"github.com/stretchr/testify/assert"
10+
"github.com/stretchr/testify/require"
11+
"github.com/v-byte-cpu/sx/pkg/scan/arp"
12+
)
13+
14+
func scanResultToJSON(t *testing.T, result *arp.ScanResult) string {
15+
t.Helper()
16+
data, err := result.MarshalJSON()
17+
require.NoError(t, err)
18+
return string(data)
19+
}
20+
21+
func TestJSONLoggerResults(t *testing.T) {
22+
t.Parallel()
23+
24+
tests := []struct {
25+
name string
26+
expected []byte
27+
results []*arp.ScanResult
28+
}{
29+
{
30+
name: "emptyResults",
31+
expected: nil,
32+
results: nil,
33+
},
34+
{
35+
name: "oneResult",
36+
expected: []byte(scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 3).To4())) + "\n"),
37+
results: []*arp.ScanResult{
38+
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
39+
},
40+
},
41+
{
42+
name: "twoResults",
43+
expected: []byte(strings.Join([]string{
44+
scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 3).To4())),
45+
scanResultToJSON(t, newScanResult(net.IPv4(192, 168, 0, 5).To4())),
46+
}, "\n") + "\n"),
47+
results: []*arp.ScanResult{
48+
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
49+
newScanResult(net.IPv4(192, 168, 0, 5).To4()),
50+
},
51+
},
52+
}
53+
54+
for _, tt := range tests {
55+
t.Run(tt.name, func(t *testing.T) {
56+
57+
var buf bytes.Buffer
58+
logger, err := NewJSONLogger(&buf, "arp")
59+
require.NoError(t, err)
60+
61+
resultCh := make(chan *arp.ScanResult, len(tt.results))
62+
for _, result := range tt.results {
63+
resultCh <- result
64+
}
65+
close(resultCh)
66+
logger.LogResults(resultCh)
67+
68+
assert.Equal(t, tt.expected, buf.Bytes())
69+
})
70+
}
71+
}

‎command/log/logger.go‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,8 @@
1+
package log
2+
3+
import "github.com/v-byte-cpu/sx/pkg/scan/arp"
4+
5+
type Logger interface {
6+
Error(err error)
7+
LogResults(results <-chan *arp.ScanResult)
8+
}

‎command/log/plain_logger.go‎

Lines changed: 41 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,41 @@
1+
package log
2+
3+
import (
4+
"bufio"
5+
"io"
6+
7+
"github.com/v-byte-cpu/sx/pkg/scan/arp"
8+
"go.uber.org/zap"
9+
)
10+
11+
type PlainLogger struct {
12+
logger *zap.Logger
13+
writer io.Writer
14+
label string
15+
}
16+
17+
func NewPlainLogger(w io.Writer, label string) (*PlainLogger, error) {
18+
logger, err := zap.NewProduction()
19+
if err != nil {
20+
return nil, err
21+
}
22+
return &PlainLogger{logger, w, label}, nil
23+
}
24+
25+
func (l *PlainLogger) Error(err error) {
26+
l.logger.Error(l.label, zap.Error(err))
27+
}
28+
29+
func (l *PlainLogger) LogResults(results <-chan *arp.ScanResult) {
30+
bw := bufio.NewWriter(l.writer)
31+
defer bw.Flush()
32+
var err error
33+
for result := range results {
34+
if _, err = bw.WriteString(result.String()); err != nil {
35+
l.Error(err)
36+
}
37+
if err = bw.WriteByte('\n'); err != nil {
38+
l.Error(err)
39+
}
40+
}
41+
}

‎command/log/plain_logger_test.go‎

Lines changed: 72 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,72 @@
1+
package log
2+
3+
import (
4+
"bytes"
5+
"net"
6+
"strings"
7+
"testing"
8+
9+
"github.com/stretchr/testify/assert"
10+
"github.com/stretchr/testify/require"
11+
"github.com/v-byte-cpu/sx/pkg/scan/arp"
12+
)
13+
14+
func newScanResult(ip net.IP) *arp.ScanResult {
15+
return &arp.ScanResult{
16+
IP: ip.String(),
17+
MAC: net.HardwareAddr{0x11, 0x22, 0x33, 0x44, 0x55, 0x66}.String(),
18+
Vendor: "Sunny Industries",
19+
}
20+
}
21+
22+
func TestPlainLoggerResults(t *testing.T) {
23+
t.Parallel()
24+
25+
tests := []struct {
26+
name string
27+
expected []byte
28+
results []*arp.ScanResult
29+
}{
30+
{
31+
name: "emptyResults",
32+
expected: nil,
33+
results: nil,
34+
},
35+
{
36+
name: "oneResult",
37+
expected: []byte(newScanResult(net.IPv4(192, 168, 0, 3).To4()).String() + "\n"),
38+
results: []*arp.ScanResult{
39+
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
40+
},
41+
},
42+
{
43+
name: "twoResults",
44+
expected: []byte(strings.Join([]string{
45+
newScanResult(net.IPv4(192, 168, 0, 3).To4()).String(),
46+
newScanResult(net.IPv4(192, 168, 0, 5).To4()).String(),
47+
}, "\n") + "\n"),
48+
results: []*arp.ScanResult{
49+
newScanResult(net.IPv4(192, 168, 0, 3).To4()),
50+
newScanResult(net.IPv4(192, 168, 0, 5).To4()),
51+
},
52+
},
53+
}
54+
55+
for _, tt := range tests {
56+
t.Run(tt.name, func(t *testing.T) {
57+
58+
var buf bytes.Buffer
59+
logger, err := NewPlainLogger(&buf, "arp")
60+
require.NoError(t, err)
61+
62+
resultCh := make(chan *arp.ScanResult, len(tt.results))
63+
for _, result := range tt.results {
64+
resultCh <- result
65+
}
66+
close(resultCh)
67+
logger.LogResults(resultCh)
68+
69+
assert.Equal(t, tt.expected, buf.Bytes())
70+
})
71+
}
72+
}

0 commit comments

Comments
 (0)