-
Notifications
You must be signed in to change notification settings - Fork 7
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
* Add aslookup segment * Clarify documentation * Add config sanity checking, make db default type, add tests * Add example config * Add example lookup file, use correct file path --------- Co-authored-by: Sebastian Schnorbus <[email protected]>
- Loading branch information
Showing
8 changed files
with
184 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Binary file not shown.
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,110 @@ | ||
package aslookup | ||
|
||
import ( | ||
"os" | ||
"log" | ||
"net" | ||
"sync" | ||
|
||
"github.com/banviktor/asnlookup/pkg/database" | ||
"github.com/bwNetFlow/flowpipeline/segments" | ||
) | ||
|
||
type AsLookup struct { | ||
segments.BaseSegment | ||
FileName string | ||
Type string | ||
|
||
asDatabase database.Database | ||
} | ||
|
||
func (segment AsLookup) New(config map[string]string) segments.Segment { | ||
|
||
newSegment := &AsLookup{} | ||
|
||
// parse options | ||
if config["filename"] == "" { | ||
log.Println("[error] AsLookup: This segment requires a 'filename' parameter.") | ||
return nil | ||
} | ||
newSegment.FileName = config["filename"] | ||
|
||
if config["type"] == "db" { | ||
newSegment.Type = "db" | ||
} else if config["type"] == "mrt" { | ||
newSegment.Type = "mrt" | ||
} else { | ||
log.Println("[info] AsLookup: 'type' set to default 'db'.") | ||
newSegment.Type = "db" | ||
} | ||
|
||
// open lookup file | ||
lookupfile, err := os.OpenFile(config["filename"], os.O_RDONLY, 0) | ||
if err != nil { | ||
log.Printf("[error] AsLookup: Error opening lookup file: %s", err) | ||
return nil | ||
} | ||
defer lookupfile.Close() | ||
|
||
// lookup file can either be an MRT file or a lookup database generated with asnlookup | ||
// see: https://github.com/banviktor/asnlookup | ||
if newSegment.Type == "db" { | ||
// open lookup db | ||
db, err := database.NewFromDump(lookupfile) | ||
if err != nil { | ||
log.Printf("[error] AsLookup: Error parsing database file: %s", err) | ||
} | ||
newSegment.asDatabase = db | ||
} else { | ||
// parse with asnlookup | ||
builder := database.NewBuilder() | ||
if err = builder.ImportMRT(lookupfile); err != nil { | ||
log.Printf("[error] AsLookup: Error parsing MRT file: %s", err) | ||
} | ||
|
||
// build lookup database | ||
db, err := builder.Build() | ||
if err != nil { | ||
log.Printf("[error] AsLookup: Error building lookup database: %s", err) | ||
} | ||
newSegment.asDatabase = db | ||
} | ||
|
||
|
||
return newSegment | ||
} | ||
|
||
func (segment *AsLookup) Run(wg *sync.WaitGroup) { | ||
defer func() { | ||
close(segment.Out) | ||
wg.Done() | ||
}() | ||
for msg := range segment.In { | ||
// Look up destination AS | ||
dstIp := net.ParseIP(msg.DstAddrObj().String()) | ||
dstAs, err := segment.asDatabase.Lookup(dstIp) | ||
if err != nil { | ||
log.Printf("[warning] AsLookup: Failed to look up ASN for %s: %s", msg.DstAddrObj().String(), err) | ||
segment.Out <- msg | ||
continue | ||
} | ||
msg.DstAS = dstAs.Number | ||
|
||
// Look up source AS | ||
srcIp := net.ParseIP(msg.SrcAddrObj().String()) | ||
srcAs, err := segment.asDatabase.Lookup(srcIp) | ||
if err != nil { | ||
log.Printf("[warning] AsLookup: Failed to look up ASN for %s: %s", msg.SrcAddrObj().String(), err) | ||
segment.Out <- msg | ||
continue | ||
} | ||
msg.SrcAS = srcAs.Number | ||
|
||
segment.Out <- msg | ||
} | ||
} | ||
|
||
func init() { | ||
segment := &AsLookup{} | ||
segments.RegisterSegment("aslookup", segment) | ||
} |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Original file line number | Diff line number | Diff line change |
---|---|---|
@@ -0,0 +1,31 @@ | ||
package aslookup | ||
|
||
import ( | ||
"testing" | ||
|
||
"github.com/bwNetFlow/flowpipeline/pb" | ||
"github.com/bwNetFlow/flowpipeline/segments" | ||
) | ||
|
||
// TODO: write tests for this | ||
func TestSegment_AsLookup_existingIp(t *testing.T) { | ||
result := segments.TestSegment("aslookup", map[string]string{"filename": "../../../examples/enricher/lookup.db", "type": "db"}, | ||
&pb.EnrichedFlow{SrcAddr: []byte{192, 168, 1, 10}, DstAddr: []byte{192, 168, 1, 10}}) | ||
if result.SrcAS != 65015 { | ||
t.Error("Segment AsLookup is not setting the source AS when the corresponding IP exists in the lookup database") | ||
} | ||
if result.DstAS != 65015 { | ||
t.Error("Segment AsLookup is not setting the destination AS when the corresponding IP exists in the lookup database") | ||
} | ||
} | ||
|
||
func TestSegment_AsLookup_nonexistingIp(t *testing.T) { | ||
result := segments.TestSegment("aslookup", map[string]string{"filename": "../../../examples/enricher/lookup.db", "type": "db"}, | ||
&pb.EnrichedFlow{SrcAddr: []byte{2, 125, 160, 218}, DstAddr: []byte{2, 125, 160, 218}}) | ||
if result.SrcAS != 0 { | ||
t.Error("Segment AsLookup is setting the source AS when the corresponding IP does not exist in the lookup database.") | ||
} | ||
if result.DstAS != 0 { | ||
t.Error("Segment AsLookup is setting the destination AS when the corresponding IP does not exist in the lookup database.") | ||
} | ||
} |