From b996df31b23282dc99c6ee8a9e71d322777e7d13 Mon Sep 17 00:00:00 2001 From: Kevin Pham Date: Fri, 1 Dec 2023 12:44:12 -0600 Subject: [PATCH] add rulesetmap type for efficient ruleset lookup for proxychain impl --- internal/cli/art.go | 45 ++++++++++++++++++++++++++------ pkg/ruleset/ruleset.go | 37 ++++++++++++++++++++++++++ pkg/ruleset/ruleset_test.go | 52 +++++++++++++++++++++++++++++++++++++ 3 files changed, 126 insertions(+), 8 deletions(-) diff --git a/internal/cli/art.go b/internal/cli/art.go index ccae8e7..145be31 100644 --- a/internal/cli/art.go +++ b/internal/cli/art.go @@ -2,6 +2,8 @@ package cli import ( "fmt" + "golang.org/x/term" + "os" "strings" ) @@ -20,20 +22,32 @@ var art string = ` ` func StartupMessage(version string, port string, ruleset string) string { + isTerm := term.IsTerminal(int(os.Stdout.Fd())) version = strings.Trim(version, " ") version = strings.Trim(version, "\n") - link := createHyperlink("http://localhost:" + port) - buf := fmt.Sprintf(art, version, link) - if ruleset == "" { - buf += "\n ! no ruleset specified.\n > for better performance, use a ruleset using --ruleset\n" + + var link string + if isTerm { + link = createHyperlink("http://localhost:" + port) } else { - buf += fmt.Sprintf("\n > using ruleset: %s\n", ruleset) + link = "http://localhost:" + port } - return colorizeNonASCII(buf) + + buf := fmt.Sprintf(art, version, link) + if isTerm { + buf = blinkChars(buf, '.', '•', '·', '▪') + } + + if ruleset == "" { + buf += "\n [!] no ruleset specified.\n [!] for better performance, use a ruleset using --ruleset\n" + } + if isTerm { + buf = colorizeNonASCII(buf) + } + return buf } func createHyperlink(url string) string { - //return fmt.Sprintf("\033]8;;%s\a%s\033]8;;\a", url, url) return fmt.Sprintf("\033[4m%s\033[0m", url) } @@ -42,7 +56,7 @@ func colorizeNonASCII(input string) string { for _, r := range input { if r > 127 { // If the character is non-ASCII, color it blue - result += fmt.Sprintf("\033[94m%c\033[0m", r) + result += fmt.Sprintf("\033[34m%c\033[0m", r) } else { // ASCII characters remain unchanged result += string(r) @@ -50,3 +64,18 @@ func colorizeNonASCII(input string) string { } return result } + +func blinkChars(input string, chars ...rune) string { + result := "" +MAIN: + for _, x := range input { + for _, y := range chars { + if x == y { + result += fmt.Sprintf("\033[5m%s\033[0m", string(x)) + continue MAIN + } + } + result += fmt.Sprintf("%s", string(x)) + } + return result +} diff --git a/pkg/ruleset/ruleset.go b/pkg/ruleset/ruleset.go index 7f3079f..24fa577 100644 --- a/pkg/ruleset/ruleset.go +++ b/pkg/ruleset/ruleset.go @@ -308,3 +308,40 @@ func debugPrintRule(rule string, err error) { fmt.Println(rule) fmt.Println("------------------------------ END DEBUG RULESET -------------------------------") } + +// ======================= RuleSetMap implementation ================================================= + +// RuleSetMap: A map with domain names as keys and pointers to the corresponding Rules as values. +// This type is used to efficiently access rules based on domain names. +type RuleSetMap map[string]*Rule + +// ToMap converts a RuleSet into a RuleSetMap. It transforms each Rule in the RuleSet +// into a map entry where the key is the Rule's domain (lowercase) +// and the value is a pointer to the Rule. This method is used to +// efficiently access rules based on domain names. +// The RuleSetMap may be accessed with or without a "www." prefix in the domain. +func (rs *RuleSet) ToMap() RuleSetMap { + rsm := make(RuleSetMap) + + addMapEntry := func(d string, rule *Rule) { + d = strings.ToLower(d) + rsm[d] = rule + if strings.HasPrefix(d, "www.") { + d = strings.TrimPrefix(d, "www.") + rsm[d] = rule + } else { + d = fmt.Sprintf("www.%s", d) + rsm[d] = rule + } + } + + for i, rule := range *rs { + rulePtr := &(*rs)[i] + addMapEntry(rule.Domain, rulePtr) + for _, domain := range rule.Domains { + addMapEntry(domain, rulePtr) + } + } + + return rsm +} diff --git a/pkg/ruleset/ruleset_test.go b/pkg/ruleset/ruleset_test.go index 42f115c..99bd663 100644 --- a/pkg/ruleset/ruleset_test.go +++ b/pkg/ruleset/ruleset_test.go @@ -171,3 +171,55 @@ func TestLoadRulesFromLocalDir(t *testing.T) { assert.Equal(t, rule.RegexRules[0].Replace, "https:") } } + +func TestToMap(t *testing.T) { + // Prepare a ruleset with multiple rules, including "www." prefixed domains + rules := RuleSet{ + { + Domain: "Example.com", + RegexRules: []Regex{{Match: "match1", Replace: "replace1"}}, + }, + { + Domain: "www.AnotherExample.com", + RegexRules: []Regex{{Match: "match2", Replace: "replace2"}}, + }, + { + Domain: "www.foo.bAr.baz.bOol.quX.com", + RegexRules: []Regex{{Match: "match3", Replace: "replace3"}}, + }, + } + + // Convert to RuleSetMap + rsm := rules.ToMap() + + // Test for correct number of entries + if len(rsm) != 6 { + t.Errorf("Expected 6 entries in RuleSetMap, got %d", len(rsm)) + } + + // Test for correct mapping + testDomains := []struct { + domain string + expectedMatch string + }{ + {"example.com", "match1"}, + {"www.example.com", "match1"}, + {"anotherexample.com", "match2"}, + {"www.anotherexample.com", "match2"}, + {"foo.bar.baz.bool.qux.com", "match3"}, + {"no.ruleset.domain.com", ""}, + } + + for _, test := range testDomains { + if test.domain == "no.ruleset.domain.com" { + assert.Empty(t, test.expectedMatch) + continue + } + rule, exists := rsm[test.domain] + if !exists { + t.Errorf("Expected domain %s to exist in RuleSetMap", test.domain) + } else if rule.RegexRules[0].Match != test.expectedMatch { + t.Errorf("Expected match for %s to be %s, got %s", test.domain, test.expectedMatch, rule.RegexRules[0].Match) + } + } +}