Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

optional obfuscate positional params, booleans and null #11

Merged
merged 3 commits into from
Sep 19, 2023
Merged
Show file tree
Hide file tree
Changes from 1 commit
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
40 changes: 38 additions & 2 deletions obfuscator.go
Original file line number Diff line number Diff line change
Expand Up @@ -5,8 +5,11 @@ import (
)

type obfuscatorConfig struct {
ReplaceDigits bool
DollarQuotedFunc bool
DollarQuotedFunc bool
ReplaceDigits bool
ReplacePositionalParameter bool
ReplaceBoolean bool
ReplaceNull bool
}

type obfuscatorOption func(*obfuscatorConfig)
Expand All @@ -17,6 +20,24 @@ func WithReplaceDigits(replaceDigits bool) obfuscatorOption {
}
}

func WithReplacePositionalParameter(replacePositionalParameter bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.ReplacePositionalParameter = replacePositionalParameter
}
}

func WithReplaceBoolean(replaceBoolean bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.ReplaceBoolean = replaceBoolean
}
}

func WithReplaceNull(replaceNull bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.ReplaceNull = replaceNull
}
}

func WithDollarQuotedFunc(dollarQuotedFunc bool) obfuscatorOption {
return func(c *obfuscatorConfig) {
c.DollarQuotedFunc = dollarQuotedFunc
Expand Down Expand Up @@ -69,7 +90,22 @@ func (o *Obfuscator) Obfuscate(input string, lexerOpts ...lexerOption) string {
}
case STRING, INCOMPLETE_STRING, DOLLAR_QUOTED_STRING:
obfuscatedSQL.WriteString(StringPlaceholder)
case POSITIONAL_PARAMETER:
if o.config.ReplacePositionalParameter {
obfuscatedSQL.WriteString(StringPlaceholder)
} else {
obfuscatedSQL.WriteString(token.Value)
}
case IDENT:
if o.config.ReplaceBoolean && isBoolean(token.Value) {
obfuscatedSQL.WriteString(StringPlaceholder)
continue
}
if o.config.ReplaceNull && isNull(token.Value) {
obfuscatedSQL.WriteString(StringPlaceholder)
continue
}

if o.config.ReplaceDigits {
obfuscatedSQL.WriteString(replaceDigits(token.Value, "?"))
} else {
Expand Down
50 changes: 45 additions & 5 deletions obfuscator_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -9,11 +9,14 @@ import (

func TestObfuscator(t *testing.T) {
tests := []struct {
input string
expected string
replaceDigits bool
dollarQuotedFunc bool
dbms DBMSType
input string
expected string
replaceDigits bool
replacePositionalParameter bool
replaceBoolean bool
replaceNull bool
dollarQuotedFunc bool
dbms DBMSType
}{
{
input: "SELECT * FROM users where id = 1",
Expand Down Expand Up @@ -55,6 +58,30 @@ func TestObfuscator(t *testing.T) {
expected: "SELECT * FROM users? where id = ?",
replaceDigits: true,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
replaceBoolean: false,
replaceNull: false,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is NULL and is_active = ? and is_admin = ?",
replaceBoolean: true,
replaceNull: false,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is ? and is_active = TRUE and is_admin = FALSE",
replaceBoolean: false,
replaceNull: true,
},
{
input: "SELECT * FROM users where id is NULL and is_active = TRUE and is_admin = FALSE",
expected: "SELECT * FROM users where id is ? and is_active = ? and is_admin = ?",
replaceBoolean: true,
replaceNull: true,
},
{
input: "SELECT * FROM users where id = 1 -- this is a comment",
expected: "SELECT * FROM users where id = ? -- this is a comment",
Expand Down Expand Up @@ -206,6 +233,16 @@ func TestObfuscator(t *testing.T) {
expected: `USING - SELECT`,
replaceDigits: true,
},
{
input: "SELECT * FROM users where id = $1",
expected: `SELECT * FROM users where id = $1`,
replacePositionalParameter: false,
},
{
input: "SELECT * FROM users where id = $1",
expected: `SELECT * FROM users where id = ?`,
replacePositionalParameter: true,
},
{
input: `SELECT * FROM "public"."users" where id = 2`,
expected: `SELECT * FROM "public"."users" where id = ?`,
Expand Down Expand Up @@ -335,6 +372,9 @@ func TestObfuscator(t *testing.T) {
t.Run("", func(t *testing.T) {
obfuscator := NewObfuscator(
WithReplaceDigits(tt.replaceDigits),
WithReplacePositionalParameter(tt.replacePositionalParameter),
WithReplaceBoolean(tt.replaceBoolean),
WithReplaceNull(tt.replaceNull),
WithDollarQuotedFunc(tt.dollarQuotedFunc),
)
got := obfuscator.Obfuscate(tt.input, WithDBMS(tt.dbms))
Expand Down
8 changes: 4 additions & 4 deletions sqllexer.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,7 +17,7 @@ const (
PUNCTUATION // punctuation
DOLLAR_QUOTED_FUNCTION // dollar quoted function
DOLLAR_QUOTED_STRING // dollar quoted string
NUMBERED_PARAMETER // numbered parameter
POSITIONAL_PARAMETER // numbered parameter
UNKNOWN // unknown token
)

Expand Down Expand Up @@ -121,7 +121,7 @@ func (s *Lexer) Scan() Token {
case ch == '$':
if isDigit(s.lookAhead(1)) {
// if the dollar sign is followed by a digit, then it's a numbered parameter
return s.scanNumberedParameter()
return s.scanPositionalParameter()
}
if s.config.DBMS == DBMSSQLServer && isLetter(s.lookAhead(1)) {
return s.scanIdentifier()
Expand Down Expand Up @@ -393,7 +393,7 @@ func (s *Lexer) scanDollarQuotedString() Token {
return Token{ERROR, s.src[s.start:s.cursor]}
}

func (s *Lexer) scanNumberedParameter() Token {
func (s *Lexer) scanPositionalParameter() Token {
s.start = s.cursor
ch := s.nextBy(2) // consume the dollar sign and the number
for {
Expand All @@ -402,7 +402,7 @@ func (s *Lexer) scanNumberedParameter() Token {
}
ch = s.next()
}
return Token{NUMBERED_PARAMETER, s.src[s.start:s.cursor]}
return Token{POSITIONAL_PARAMETER, s.src[s.start:s.cursor]}
}

func (s *Lexer) scanUnknown() Token {
Expand Down
2 changes: 1 addition & 1 deletion sqllexer_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -312,7 +312,7 @@ func TestLexer(t *testing.T) {
{WS, " "},
{OPERATOR, "="},
{WS, " "},
{NUMBERED_PARAMETER, "$1"},
{POSITIONAL_PARAMETER, "$1"},
},
},
{
Expand Down
8 changes: 8 additions & 0 deletions sqllexer_utils.go
Original file line number Diff line number Diff line change
Expand Up @@ -116,6 +116,14 @@ func isSQLKeyword(token Token) bool {
return token.Type == IDENT && keywordsRegex.MatchString(token.Value)
}

func isBoolean(ident string) bool {
return strings.ToUpper(ident) == "TRUE" || strings.ToUpper(ident) == "FALSE"
}

func isNull(ident string) bool {
return strings.ToUpper(ident) == "NULL"
}

func replaceDigits(input string, placeholder string) string {
var builder strings.Builder

Expand Down