Skip to content

Commit

Permalink
add whitelist support
Browse files Browse the repository at this point in the history
  • Loading branch information
dmachard committed Dec 4, 2024
1 parent 40a752d commit 2b24b62
Show file tree
Hide file tree
Showing 5 changed files with 86 additions and 12 deletions.
24 changes: 24 additions & 0 deletions docs/transformers/transform_newdomaintracker.md
Original file line number Diff line number Diff line change
Expand Up @@ -25,11 +25,16 @@ The **New Domain Tracker** transformer identifies domains that are newly observe
* `cache-size` (integer)
> Maximum number of domains to track
* `white-domains-file` (string)
> path file to domain white list, domains list can be a partial domain name with regexp expression

```yaml
transforms:
new-domain-tracker:
ttl: 3600
cache-size: 100000
white-domains-file: ""
```
## Cache
Expand All @@ -38,6 +43,25 @@ The New Domain Tracker uses an **LRU Cache** to manage memory consumption effici
The LRU Cache ensures finite memory usage but may cause some domains to be forgotten if the cache size is too small.
## Whitelist
Example of configuration to load a whitelist of domains to ignore.
```yaml
transforms:
new-domain-tracker:
ttl: 3600
cache-size: 100000
white-domains-file: /tmp/whitelist_domain.txt
```
Example of content for the file `/tmp/whitelist_domain.txt`

```
(mail|wwww).google.com
github.com
```
## Persistence
To ensure continuity across application restarts, you can enable the persistence feature by specifying a file path (persistence). The transformer will save the domain cache to this file and reload it on startup.
Expand Down
2 changes: 2 additions & 0 deletions tests/testsdata/newdomain_whitelist_regex.txt
Original file line number Diff line number Diff line change
@@ -0,0 +1,2 @@
.*\.google\.com
github\.com
39 changes: 29 additions & 10 deletions transformers/newdomain.go
Original file line number Diff line number Diff line change
Expand Up @@ -16,14 +16,15 @@ import (

// NewDomainTracker transformer to detect newly observed domains
type NewDomainTracker struct {
ttl time.Duration // Time window to consider a domain as "new"
cache *lru.Cache // LRU Cache to store observed domains
logInfo func(msg string, v ...interface{})
logError func(msg string, v ...interface{})
ttl time.Duration // Time window to consider a domain as "new"
cache *lru.Cache // LRU Cache to store observed domains
whitelist map[string]*regexp.Regexp // Whitelisted domains
logInfo func(msg string, v ...interface{})
logError func(msg string, v ...interface{})
}

// NewNewDomainTracker initializes the NewDomainTracker transformer
func NewNewDomainTracker(ttl time.Duration, maxSize int, logInfo, logError func(msg string, v ...interface{})) (*NewDomainTracker, error) {
func NewNewDomainTracker(ttl time.Duration, maxSize int, whitelist map[string]*regexp.Regexp, logInfo, logError func(msg string, v ...interface{})) (*NewDomainTracker, error) {
cache, err := lru.New(maxSize)
if err != nil {
return nil, err
Expand All @@ -34,20 +35,38 @@ func NewNewDomainTracker(ttl time.Duration, maxSize int, logInfo, logError func(
}

return &NewDomainTracker{
ttl: ttl,
cache: cache,
logInfo: logInfo,
logError: logError,
ttl: ttl,
cache: cache,
whitelist: whitelist,
logInfo: logInfo,
logError: logError,
}, nil
}

// isWhitelisted checks if a domain or its subdomain is in the whitelist
func (ndt *NewDomainTracker) isWhitelisted(domain string) bool {
for _, d := range ndt.whitelist {
if d.MatchString(domain) {
return true
}
}
return false
}

// IsNewDomain checks if the domain is newly observed
func (ndt *NewDomainTracker) IsNewDomain(domain string) bool {
// Check if the domain is whitelisted
if ndt.isWhitelisted(domain) {
return false
}

now := time.Now()

// Check if the domain exists in the cache
if lastSeen, exists := ndt.cache.Get(domain); exists {
fmt.Println("exists")
if now.Sub(lastSeen.(time.Time)) < ndt.ttl {
fmt.Println(now.Sub(lastSeen.(time.Time)), ndt.ttl)
// Domain was recently seen, not new
return false
}
Expand Down Expand Up @@ -91,7 +110,7 @@ func (t *NewDomainTrackerTransform) GetTransforms() ([]Subtransform, error) {
// Initialize the domain tracker
ttl := time.Duration(t.config.NewDomainTracker.TTL) * time.Second
maxSize := t.config.NewDomainTracker.CacheSize
tracker, err := NewNewDomainTracker(ttl, maxSize, t.LogInfo, t.LogError)
tracker, err := NewNewDomainTracker(ttl, maxSize, t.listDomainsRegex, t.LogInfo, t.LogError)
if err != nil {
return nil, err
}
Expand Down
31 changes: 30 additions & 1 deletion transformers/newdomain_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,7 +9,7 @@ import (
"github.com/dmachard/go-logger"
)

func TestNewDomainTracker(t *testing.T) {
func TestNewDomainTracker_IsNew(t *testing.T) {
// config
config := pkgconfig.GetFakeConfigTransformers()
config.NewDomainTracker.Enable = true
Expand Down Expand Up @@ -43,5 +43,34 @@ func TestNewDomainTracker(t *testing.T) {
if result, _ := tracker.trackNewDomain(&dm); result != ReturnKeep {
t.Errorf("3. this domain should be new!!")
}
}

func TestNewDomainTracker_Whitelist(t *testing.T) {
// config
config := pkgconfig.GetFakeConfigTransformers()
config.NewDomainTracker.Enable = true
config.NewDomainTracker.TTL = 2
config.NewDomainTracker.CacheSize = 10
config.NewDomainTracker.WhiteDomainsFile = "../tests/testsdata/newdomain_whitelist_regex.txt"

// init subproccesor
outChans := []chan dnsutils.DNSMessage{}
tracker := NewNewDomainTrackerTransform(config, logger.New(false), "test", 0, outChans)
_, err := tracker.GetTransforms()
if err != nil {
t.Error("fail to init transform", err)
}

// first test, check domain in whilist
dm := dnsutils.GetFakeDNSMessage()
dm.DNS.Qname = testURL1
if result, _ := tracker.trackNewDomain(&dm); result != ReturnDrop {
t.Errorf("2. this domain should NOT be new!!")
}

// second test, check domain in whilist
dm = dnsutils.GetFakeDNSMessage()
if result, _ := tracker.trackNewDomain(&dm); result != ReturnKeep {
t.Errorf("2. this domain should be new!!")
}
}
2 changes: 1 addition & 1 deletion transformers/transformers.go
Original file line number Diff line number Diff line change
Expand Up @@ -108,7 +108,7 @@ func (p *Transforms) Prepare() error {
for _, transform := range p.availableTransforms {
subtransforms, err := transform.GetTransforms()
if err != nil {
p.LogError("error on init subtransforms:", err)
p.LogError("error on init subtransforms: %v", err)
continue
}
for _, subtransform := range subtransforms {
Expand Down

0 comments on commit 2b24b62

Please sign in to comment.