Added padding function (#75)

* Added padding function * Added test cases * rename to padding + minor improvements * add example in README.md --------- Co-authored-by: Tarun Koyalwar <tarun@projectdiscovery.io>
projectdiscovery · Aug 11, 2023 · 4bde0f5 · 4bde0f5
1 parent 516390e
commit 4bde0f5
Show file tree

Hide file tree

Showing 3 changed files with 59 additions and 0 deletions.
diff --git a/README.md b/README.md
@@ -131,6 +131,7 @@ func main() {
 | xor(sequences ...strings/bytes)                                                | Perform xor on sequences of same length                                                                | `xor('abc','def')`                                                                                                                                       | `50705` in hex                                                                                                                                                                                                                                                                                                                                                                                     |
 | zlib(input string) string                                             | Compresses the input using Zlib                                                                                     | `base64(zlib("Hello"))`                                                                                                                              | `eJzySM3JyQcEAAD//wWMAfU=`                                                                                                                                                                                                                                                                                                                                                                 |
 | zlib_decode(input string) string                                      | Decompresses the input using Zlib                                                                                   | `zlib_decode(hex_decode("789cf248cdc9c907040000ffff058c01f5"))`                                                                                      | `Hello`                                                                                                                                                                                                                                                                                                                                                                                    |
+| padding(input string, padding string,length int) string               | Adds padding char to input string until it reaches desired length                              					  | `padding("A","b",3)`                                                                                                                          | `Abb`                                                                                                                                                                                                                                                                                                                                                                               |
 
 **Supported encodings:**
 

diff --git a/dsl.go b/dsl.go
@@ -44,6 +44,7 @@ import (
 	"github.com/projectdiscovery/mapcidr"
 	jarm "github.com/projectdiscovery/utils/crypto/jarm"
 	errors "github.com/projectdiscovery/utils/errors"
+	errorutil "github.com/projectdiscovery/utils/errors"
 	maputils "github.com/projectdiscovery/utils/maps"
 	randint "github.com/projectdiscovery/utils/rand"
 	stringsutil "github.com/projectdiscovery/utils/strings"
@@ -818,6 +819,61 @@ func init() {
 			result := constraint.Check(firstParsed)
 			return result, nil
 		}))
+	MustAddFunction(NewWithPositionalArgs("padding", 3, false, func(args ...interface{}) (interface{}, error) {
+		// padding('Test String','A',50) // will pad "Test String" up to 50 characters with "A" as padding byte.
+		bLen := 0
+		switch value := args[2].(type) {
+		case float64:
+			bLen = int(value)
+		case int:
+			bLen = value
+		default:
+			strLen := toString(args[2])
+			floatVal, err := strconv.ParseFloat(strLen, 64)
+			if err != nil {
+				return nil, err
+			}
+			bLen = int(floatVal)
+		}
+		if bLen == 0 {
+			return nil, errorutil.New("invalid padding length")
+		}
+		bByte := []byte(toString(args[1]))
+		if len(bByte) == 0 {
+			return nil, errorutil.New("invalid padding byte")
+		}
+		bData := []byte(toString(args[0]))
+		dataLen := len(bData)
+		if dataLen >= bLen {
+			return toString(bData), nil // Note: if given string is longer than the desired length, it will not be truncated
+		}
+		if dataLen == 0 {
+			// If the initial string is empty, simply create a padded array with the specified length
+			paddedData := make([]byte, bLen)
+			for i := 0; i < bLen; i++ {
+				paddedData[i] = bByte[i%len(bByte)]
+			}
+			return toString(paddedData), nil
+		}
+
+		// Calculate the number of bytes needed for padding
+		paddingLen := (bLen - (dataLen % bLen)) % bLen
+
+		// Create a new byte array with the desired length
+		paddedData := make([]byte, dataLen+paddingLen)
+
+		// Copy the original data into the padded array
+		copy(paddedData, bData)
+
+		// Add padding bytes with the specified padding byte
+		for i := dataLen; i < len(paddedData); i++ {
+			paddedData[i] = bByte[i%len(bByte)]
+		}
+
+		return toString(paddedData), nil
+
+	}))
+
 	MustAddFunction(NewWithSingleSignature("print_debug",
 		"(args ...interface{})",
 		false,

diff --git a/dsl_test.go b/dsl_test.go
@@ -139,6 +139,7 @@ func TestGetPrintableDslFunctionSignatures(t *testing.T) {
 	md5(arg1 interface{}) interface{}
 	mmh3(arg1 interface{}) interface{}
 	oct_to_dec(arg1 interface{}) interface{}
+	padding(arg1, arg2, arg3 interface{}) interface{}
 	print_debug(args ...interface{})
 	public_ip() string
 	rand_base(length uint, optionalCharSet string) string
@@ -253,6 +254,7 @@ func TestDslExpressions(t *testing.T) {
 		"line_ends_with('Hii\nHello', 'ii')":                      true, // back quotes do not support escape sequences
 		`regex("H([a-z]+)o", "Hello")`:                            true,
 		`wait_for(1)`:                                             nil,
+		`padding("A","b",3)`:                                      "Abb",
 		`print_debug(1+2, "Hello")`:                               nil,
 		`to_number('4')`:                                          float64(4),
 		`to_string(4)`:                                            "4",