feat: add count func (#162)

Signed-off-by: Dwi Siswanto <git@dw1.io>
projectdiscovery · Aug 20, 2024 · d714c40 · d714c40
1 parent 758c17e
commit d714c40
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 4 deletions.
diff --git a/dsl.go b/dsl.go
@@ -910,7 +910,6 @@ func init() {
 		}
 	}))
 
-
 	MustAddFunction(NewWithSingleSignature("print_debug",
 		"(args ...interface{})",
 		false,
@@ -1245,6 +1244,21 @@ func init() {
 		return jarm.HashWithDialer(nil, hostname, port, 10)
 	}))
 
+	MustAddFunction(NewWithSingleSignature("count",
+		"(str, substr string) int",
+		false,
+		func(args ...interface{}) (interface{}, error) {
+			if len(args) < 2 {
+				return nil, ErrInvalidDslFunction
+			}
+
+			str := toString(args[0])
+			substr := toString(args[1])
+
+			return strings.Count(str, substr), nil
+		},
+	))
+
 	DefaultHelperFunctions = HelperFunctions()
 	FunctionNames = GetFunctionNames(DefaultHelperFunctions)
 }

diff --git a/dsl_test.go b/dsl_test.go
@@ -110,6 +110,7 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
 	contains(arg1, arg2 interface{}) interface{}
 	contains_all(body interface{}, substrs ...string) bool
 	contains_any(body interface{}, substrs ...string) bool
+	count(str, substr string) int
 	date_time(dateTimeFormat string, optionalUnixTime interface{}) string
 	dec_to_hex(arg1 interface{}) interface{}
 	deflate(arg1 interface{}) interface{}
@@ -235,9 +236,9 @@ func TestDslExpressions(t *testing.T) {
 		`trim_suffix("aaHelloaa", "aa")`:                 "aaHello",
 		`url_decode("https:%2F%2Fprojectdiscovery.io%3Ftest=1")`: "https://projectdiscovery.io?test=1",
 		`url_encode("https://projectdiscovery.io/test?a=1")`:     "https%3A%2F%2Fprojectdiscovery.io%2Ftest%3Fa%3D1",
-		`gzip("Hello")`: "\x1f\x8b\b\x00\x00\x00\x00\x00\x00\xff\xf2H\xcd\xc9\xc9\a\x04\x00\x00\xff\xff\x82\x89\xd1\xf7\x05\x00\x00\x00",
-                `zip("aaa.txt","abcd")`: ([]byte("PK\x03\x04\x14\x00\x08\x00\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x07\x00\x00\x00aaa.txtJLJN\x01\x04\x00\x00\xff\xffPK\x07\x08\x11\xcd\x82\xed\n\x00\x00\x00\x04\x00\x00\x00PK\x01\x02\x14\x00\x14\x00\x08\x00\x08\x00\x00\x00\x00\x00\x11\xcd\x82\xed\n\x00\x00\x00\x04\x00\x00\x00\x07\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00aaa.txtPK\x05\x06\x00\x00\x00\x00\x01\x00\x01\x005\x00\x00\x00?\x00\x00\x00\x00\x00")),
-		`zlib("Hello")`: "\x78\x9c\xf2\x48\xcd\xc9\xc9\x07\x04\x00\x00\xff\xff\x05\x8c\x01\xf5",
+		`gzip("Hello")`:         "\x1f\x8b\b\x00\x00\x00\x00\x00\x00\xff\xf2H\xcd\xc9\xc9\a\x04\x00\x00\xff\xff\x82\x89\xd1\xf7\x05\x00\x00\x00",
+		`zip("aaa.txt","abcd")`: ([]byte("PK\x03\x04\x14\x00\x08\x00\x08\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x07\x00\x00\x00aaa.txtJLJN\x01\x04\x00\x00\xff\xffPK\x07\x08\x11\xcd\x82\xed\n\x00\x00\x00\x04\x00\x00\x00PK\x01\x02\x14\x00\x14\x00\x08\x00\x08\x00\x00\x00\x00\x00\x11\xcd\x82\xed\n\x00\x00\x00\x04\x00\x00\x00\x07\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00\x00aaa.txtPK\x05\x06\x00\x00\x00\x00\x01\x00\x01\x005\x00\x00\x00?\x00\x00\x00\x00\x00")),
+		`zlib("Hello")`:         "\x78\x9c\xf2\x48\xcd\xc9\xc9\x07\x04\x00\x00\xff\xff\x05\x8c\x01\xf5",
 		`zlib_decode(hex_decode("789cf248cdc9c907040000ffff058c01f5"))`: "Hello",
 		`deflate("Hello")`:                              "\xf2\x48\xcd\xc9\xc9\x07\x04\x00\x00\xff\xff",
 		`inflate(hex_decode("f348cdc9c90700"))`:         "Hello",
@@ -304,6 +305,7 @@ func TestDslExpressions(t *testing.T) {
 		`ip_format('127.0.1.0', '11')`:                                         "127.0.256",
 		"unpack('>I', '\xac\xd7\t\xd0')":                                       -272646673,
 		"xor('\x01\x02', '\x02\x01')":                                          []uint8([]byte{0x3, 0x3}),
+		`count("projectdiscovery", "e")`:                                       2,
 	}
 
 	testDslExpressions(t, dslExpressions)