Skip to content

Commit

Permalink
feat(cli): LQLv2 support (#441)
Browse files Browse the repository at this point in the history
ALLY-504
  • Loading branch information
hazedav committed Jun 10, 2021
1 parent 7a0d77f commit 8ad4abf
Show file tree
Hide file tree
Showing 9 changed files with 113 additions and 120 deletions.
47 changes: 35 additions & 12 deletions api/lql.go
Original file line number Diff line number Diff line change
Expand Up @@ -22,8 +22,8 @@ import (
"encoding/json"
"fmt"
"net/url"
"regexp"
"strconv"
"strings"
"time"

"github.com/pkg/errors"
Expand All @@ -32,15 +32,14 @@ import (
)

const (
reLQL string = `(?ms)^(\w+)\([^)]+\)\s*{`
LQLQueryTranslateError string = "unable to translate query blob"
)

type LQLQuery struct {
ID string `json:"LQL_ID,omitempty"`
StartTimeRange string `json:"START_TIME_RANGE,omitempty"`
EndTimeRange string `json:"END_TIME_RANGE,omitempty"`
QueryText string `json:"QUERY_TEXT"`
ID string `json:"lql_id,omitempty"`
StartTimeRange string `json:"start_time_range,omitempty"`
EndTimeRange string `json:"end_time_range,omitempty"`
QueryText string `json:"query_text"`
// QueryBlob is a special string that supports type conversion
// back and forth from LQL to JSON
QueryBlob string `json:"-"`
Expand Down Expand Up @@ -84,13 +83,12 @@ func (q *LQLQuery) Translate() error {
}

func (q *LQLQuery) TranslateQuery() error {
// empty
// if query text is already populated
if q.QueryText != "" {
return nil
}
// json
// valid json
var t LQLQuery

if err := json.Unmarshal([]byte(q.QueryBlob), &t); err == nil {
if q.StartTimeRange == "" {
q.StartTimeRange = t.StartTimeRange
Expand All @@ -101,11 +99,23 @@ func (q *LQLQuery) TranslateQuery() error {
q.QueryText = t.QueryText
return err
}
// lql
if matched, _ := regexp.MatchString(reLQL, q.QueryBlob); matched {
// invalid json
qblob := strings.ToLower(q.QueryBlob)
if strings.Contains(qblob, "start_time_range") ||
strings.Contains(qblob, "end_time_range") ||
strings.Contains(qblob, "lql_id") ||
strings.Contains(qblob, "query_text") {

return errors.New(LQLQueryTranslateError)
}
// valid lql text
if strings.Contains(q.QueryBlob, "{") &&
strings.Contains(q.QueryBlob, "}") {
q.QueryText = q.QueryBlob

return nil
}
// invalid lql text
return errors.New(LQLQueryTranslateError)
}

Expand Down Expand Up @@ -175,7 +185,7 @@ type LQLService struct {
}

func (svc *LQLService) CreateQuery(query string) (
response LQLQueryResponse,
response LQLQuery,
err error,
) {
lqlQuery := LQLQuery{QueryBlob: query}
Expand All @@ -187,6 +197,19 @@ func (svc *LQLService) CreateQuery(query string) (
return
}

func (svc *LQLService) UpdateQuery(query string) (
response LQLQuery,
err error,
) {
lqlQuery := LQLQuery{QueryBlob: query}
if err = lqlQuery.Validate(true); err != nil {
return
}

err = svc.client.RequestEncoderDecoder("PATCH", apiLQL, lqlQuery, &response)
return
}

func (svc *LQLService) GetQueries() (LQLQueryResponse, error) {
return svc.GetQueryByID("")
}
Expand Down
75 changes: 41 additions & 34 deletions api/lql_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -32,14 +32,12 @@ import (
)

var (
lqlQueryID = "my_lql"
lqlQueryStr = "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
lqlCreateData = `[
{
"lql_id": "my_lql",
"query_text": "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
}
]`
lqlQueryID = "my_lql"
lqlQueryStr = "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
lqlCreateResponse = `{
"lql_id": "my_lql",
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
}`
lqlRunData = `[
{
"INSERT_ID": "35308423"
Expand Down Expand Up @@ -285,24 +283,35 @@ var lqlQueryTypeTests = []LQLQueryTest{
QueryBlob: `this is junk`,
},
},
LQLQueryTest{
Name: "partial-blob",
Input: &api.LQLQuery{
QueryBlob: `{`,
},
Return: api.LQLQueryTranslateError,
Expected: &api.LQLQuery{
QueryText: ``,
QueryBlob: `{`,
},
},
LQLQueryTest{
Name: "json-blob",
Input: &api.LQLQuery{
QueryBlob: `{
"START_TIME_RANGE": "678910",
"END_TIME_RANGE": "111213141516",
"QUERY_TEXT": "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
"start_time_range": "678910",
"end_time_range": "111213141516",
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
}`,
},
Return: nil,
Expected: &api.LQLQuery{
StartTimeRange: "1970-01-01T00:11:18Z",
EndTimeRange: "1973-07-11T04:32:21Z",
QueryText: "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }",
QueryText: "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }",
QueryBlob: `{
"START_TIME_RANGE": "678910",
"END_TIME_RANGE": "111213141516",
"QUERY_TEXT": "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
"start_time_range": "678910",
"end_time_range": "111213141516",
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
}`,
},
},
Expand All @@ -312,30 +321,30 @@ var lqlQueryTypeTests = []LQLQueryTest{
QueryBlob: `{
"start_time_range": "678910",
"end_time_range": "111213141516",
"query_text": "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
}`,
},
Return: nil,
Expected: &api.LQLQuery{
StartTimeRange: "1970-01-01T00:11:18Z",
EndTimeRange: "1973-07-11T04:32:21Z",
QueryText: "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }",
QueryText: "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }",
QueryBlob: `{
"start_time_range": "678910",
"end_time_range": "111213141516",
"query_text": "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
}`,
},
},
LQLQueryTest{
Name: "lql-blob",
Input: &api.LQLQuery{
QueryBlob: "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }",
QueryBlob: "--a comment\nmy_lql { source { CloudTrailRawEvents } return { INSERT_ID } }",
},
Return: nil,
Expected: &api.LQLQuery{
QueryText: "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }",
QueryBlob: "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }",
QueryText: "--a comment\nmy_lql { source { CloudTrailRawEvents } return { INSERT_ID } }",
QueryBlob: "--a comment\nmy_lql { source { CloudTrailRawEvents } return { INSERT_ID } }",
},
},
LQLQueryTest{
Expand All @@ -345,9 +354,9 @@ var lqlQueryTypeTests = []LQLQueryTest{
EndTimeRange: "1",
QueryText: "should not overwrite",
QueryBlob: `{
"START_TIME_RANGE": "678910",
"END_TIME_RANGE": "111213141516",
"QUERY_TEXT": "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
"start_time_range": "678910",
"end_time_range": "111213141516",
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
}`,
},
Return: nil,
Expand All @@ -356,9 +365,9 @@ var lqlQueryTypeTests = []LQLQueryTest{
EndTimeRange: "1970-01-01T00:00:00Z",
QueryText: "should not overwrite",
QueryBlob: `{
"START_TIME_RANGE": "678910",
"END_TIME_RANGE": "111213141516",
"QUERY_TEXT": "my_lql(CloudTrailRawEvents e) { SELECT INSERT_ID LIMIT 10 }"
"start_time_range": "678910",
"end_time_range": "111213141516",
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID } }"
}`,
},
},
Expand Down Expand Up @@ -435,13 +444,11 @@ func TestLQLCreateBadInput(t *testing.T) {
}

func TestLQLCreateOK(t *testing.T) {
mockResponse := mockLQLDataResponse(lqlCreateData)

fakeServer := lacework.MockServer()
fakeServer.MockAPI(
"external/lql",
func(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, mockResponse)
fmt.Fprint(w, lqlCreateResponse)
},
)
defer fakeServer.Close()
Expand All @@ -452,10 +459,10 @@ func TestLQLCreateOK(t *testing.T) {
)
assert.Nil(t, err)

createExpected := api.LQLQueryResponse{}
_ = json.Unmarshal([]byte(mockResponse), &createExpected)
createExpected := api.LQLQuery{}
_ = json.Unmarshal([]byte(lqlCreateResponse), &createExpected)

var createActual api.LQLQueryResponse
var createActual api.LQLQuery
createActual, err = c.LQL.CreateQuery(lqlQueryStr)
assert.Nil(t, err)

Expand Down Expand Up @@ -510,7 +517,7 @@ func TestLQLGetQueriesMethod(t *testing.T) {
}

func TestLQLGetQueryByIDOK(t *testing.T) {
mockResponse := mockLQLDataResponse(lqlCreateData)
mockResponse := mockLQLDataResponse("[" + lqlCreateResponse + "]")

fakeServer := lacework.MockServer()
fakeServer.MockAPI(
Expand Down
41 changes: 0 additions & 41 deletions api/lql_update.go

This file was deleted.

20 changes: 11 additions & 9 deletions api/lql_update_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,13 @@ import (
"github.com/stretchr/testify/assert"
)

var (
lqlUpdateResponse = `{
"lql_id": "my_lql",
"query_text": "my_lql { source { CloudTrailRawEvents } return { INSERT_ID, INSERT_TIME } }"
}`
)

func TestLQLUpdateMethod(t *testing.T) {
fakeServer := lacework.MockServer()
fakeServer.MockAPI(
Expand Down Expand Up @@ -71,16 +78,11 @@ func TestLQLUpdateBadInput(t *testing.T) {
}

func TestLQLUpdateOK(t *testing.T) {
mockResponse := mockLQLMessageResponse(
fmt.Sprintf(`"lqlUpdated": "%s"`, lqlQueryID),
"true",
)

fakeServer := lacework.MockServer()
fakeServer.MockAPI(
"external/lql",
func(w http.ResponseWriter, r *http.Request) {
fmt.Fprint(w, mockResponse)
fmt.Fprint(w, lqlUpdateResponse)
},
)
defer fakeServer.Close()
Expand All @@ -91,10 +93,10 @@ func TestLQLUpdateOK(t *testing.T) {
)
assert.Nil(t, err)

updateExpected := api.LQLUpdateResponse{}
_ = json.Unmarshal([]byte(mockResponse), &updateExpected)
updateExpected := api.LQLQuery{}
_ = json.Unmarshal([]byte(lqlUpdateResponse), &updateExpected)

var updateActual api.LQLUpdateResponse
var updateActual api.LQLQuery
updateActual, err = c.LQL.UpdateQuery(lqlQueryStr)
assert.Nil(t, err)

Expand Down
Loading

0 comments on commit 8ad4abf

Please sign in to comment.