JSON Streaming

cmd/api/src/daemons/datapipe/azure_convertors.go

@@ -36,20 +36,6 @@ const (
 	PrincipalTypeUser             = "User"
 )
 
-func convertAzureData(data []json.RawMessage) ConvertedAzureData {
-	converted := CreateConvertedAzureData(0)
-	for _, bytes := range data {
-		var data AzureBase
-		if err := json.Unmarshal(bytes, &data); err != nil {
-			log.Error().Fault(err).Msg("Failed to convert Azure data")
-		} else {
-			convert := getKindConverter(data.Kind)
-			convert(data.Data, &converted)
-		}
-	}
-	return converted
-}
-
 func getKindConverter(kind enums.Kind) func(json.RawMessage, *ConvertedAzureData) {
 	switch kind {
 	case enums.KindAZApp:

cmd/api/src/daemons/datapipe/decoders.go

@@ -0,0 +1,151 @@
+// Copyright 2024 Specter Ops, Inc.
+//
+// Licensed under the Apache License, Version 2.0
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package datapipe
+
+import (
+	"github.com/specterops/bloodhound/dawgs/graph"
+	"github.com/specterops/bloodhound/ein"
+	"github.com/specterops/bloodhound/log"
+	"io"
+)
+
+type ConversionFunc[T any] func(decoded T, converted *ConvertedData)
+
+func decodeBasicData[T any](batch graph.Batch, reader io.ReadSeeker, conversionFunc ConversionFunc[T]) error {
+	decoder, err := CreateIngestDecoder(reader)
+	if err != nil {
+		return err
+	}
+
+	var (
+		count         = 0
+		convertedData ConvertedData
+	)
+
+	for decoder.More() {
+		var decodeTarget T
+		if err := decoder.Decode(&decodeTarget); err != nil {
+			log.Errorf("Error decoding %T object: %v", decodeTarget, err)
+		} else {
+			count++
+			conversionFunc(decodeTarget, &convertedData)
+		}
+
+		if count == IngestCountThreshold {
+			IngestBasicData(batch, convertedData)
+			convertedData.Clear()
+			count = 0
+		}
+	}
+
+	if count > 0 {
+		IngestBasicData(batch, convertedData)
+	}
+
+	return nil
+}
+
+func decodeGroupData(batch graph.Batch, reader io.ReadSeeker) error {
+	decoder, err := CreateIngestDecoder(reader)
+	if err != nil {
+		return err
+	}
+
+	convertedData := ConvertedGroupData{}
+	var group ein.Group
+	count := 0
+	for decoder.More() {
+		if err := decoder.Decode(&group); err != nil {
+			log.Errorf("Error decoding group object: %v", err)
+		} else {
+			count++
+			convertGroupData(group, &convertedData)
+			if count == IngestCountThreshold {
+				IngestGroupData(batch, convertedData)
+				convertedData.Clear()
+				count = 0
+			}
+		}
+	}
+
+	if count > 0 {
+		IngestGroupData(batch, convertedData)
+	}
+
+	return nil
+}
+
+func decodeSessionData(batch graph.Batch, reader io.ReadSeeker) error {
+	decoder, err := CreateIngestDecoder(reader)
+	if err != nil {
+		return err
+	}
+
+	convertedData := ConvertedSessionData{}
+	var session ein.Session
+	count := 0
+	for decoder.More() {
+		if err := decoder.Decode(&session); err != nil {
+			log.Errorf("Error decoding session object: %v", err)
+		} else {
+			count++
+			convertSessionData(session, &convertedData)
+			if count == IngestCountThreshold {
+				IngestSessions(batch, convertedData.SessionProps)
+				convertedData.Clear()
+				count = 0
+			}
+		}
+	}
+
+	if count > 0 {
+		IngestSessions(batch, convertedData.SessionProps)
+	}
+
+	return nil
+}
+
+func decodeAzureData(batch graph.Batch, reader io.ReadSeeker) error {
+	decoder, err := CreateIngestDecoder(reader)
+	if err != nil {
+		return err
+	}
+
+	convertedData := ConvertedAzureData{}
+	var data AzureBase
+	count := 0
+	for decoder.More() {
+		if err := decoder.Decode(&data); err != nil {
+			log.Errorf("Error decoding azure object: %v", err)
+		} else {
+			convert := getKindConverter(data.Kind)
+			convert(data.Data, &convertedData)
+			count++
+			if count == IngestCountThreshold {
+				IngestAzureData(batch, convertedData)
+				convertedData.Clear()
+				count = 0
+			}
+		}
+	}
+
+	if count > 0 {
+		IngestAzureData(batch, convertedData)
+	}
+
+	return nil
+}

cmd/api/src/daemons/datapipe/fileops.go

@@ -0,0 +1,83 @@
+// Copyright 2024 Specter Ops, Inc.
+//
+// Licensed under the Apache License, Version 2.0
+// you may not use this file except in compliance with the License.
+// You may obtain a copy of the License at
+//
+//     http://www.apache.org/licenses/LICENSE-2.0
+//
+// Unless required by applicable law or agreed to in writing, software
+// distributed under the License is distributed on an "AS IS" BASIS,
+// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+// See the License for the specific language governing permissions and
+// limitations under the License.
+//
+// SPDX-License-Identifier: Apache-2.0
+
+package datapipe
+
+import (
+	"encoding/json"
+	"errors"
+	"fmt"
+	"io"
+)
+
+const (
+	delimOpenBracket        = json.Delim('{')
+	delimCloseBracket       = json.Delim('}')
+	delimOpenSquareBracket  = json.Delim('[')
+	delimCloseSquareBracket = json.Delim(']')
+)
+
+func SeekToDataTag(decoder *json.Decoder) error {
+	depth := 0
+	dataTagFound := false
+	for {
+		if token, err := decoder.Token(); err != nil {
+			if errors.Is(err, io.EOF) {
+				return ErrDataTagNotFound
+			}
+
+			return fmt.Errorf("%w: %w", ErrJSONDecoderInternal, err)
+		} else {
+			//Break here to allow for one more token read, which should take us to the "[" token, exactly where we need to be
+			if dataTagFound {
+				//Do some extra checks
+				if typed, ok := token.(json.Delim); !ok {
+					return ErrInvalidDataTag
+				} else if typed != delimOpenSquareBracket {
+					return ErrInvalidDataTag
+				}
+				//Break out of our loop if we're in a good spot
+				return nil
+			}
+			switch typed := token.(type) {
+			case json.Delim:
+				switch typed {
+				case delimCloseBracket, delimCloseSquareBracket:
+					depth--
+				case delimOpenBracket, delimOpenSquareBracket:
+					depth++
+				}
+			case string:
+				if !dataTagFound && depth == 1 && typed == "data" {
+					dataTagFound = true
+				}
+			}
+		}
+	}
+}
+
+func CreateIngestDecoder(reader io.ReadSeeker) (*json.Decoder, error) {
+	if _, err := reader.Seek(0, io.SeekStart); err != nil {
+		return nil, fmt.Errorf("error seeking to start of file: %w", err)
+	} else {
+		decoder := json.NewDecoder(reader)
+		if err := SeekToDataTag(decoder); err != nil {
+			return nil, fmt.Errorf("error seeking to data tag: %w", err)
+		} else {
+			return decoder, nil
+		}
+	}
+}

cmd/api/src/daemons/datapipe/fileops_test.go

@@ -0,0 +1,64 @@
+package datapipe_test
+
+import (
+	"encoding/json"
+	"strings"
+	"testing"
+
+	"github.com/specterops/bloodhound/src/daemons/datapipe"
+	"github.com/stretchr/testify/assert"
+)
+
+type dataTagAssertion struct {
+	rawString string
+	err       error
+}
+
+func TestSeekToDataTag(t *testing.T) {
+	assertions := []dataTagAssertion{
+		{
+			rawString: "{\"data\": []}",
+			err:       nil,
+		},
+		{
+			rawString: "{\"data\": {}}",
+			err:       datapipe.ErrInvalidDataTag,
+		},
+		{
+			rawString: "{\"data\": ]}",
+			err:       datapipe.ErrJSONDecoderInternal,
+		},
+		{
+			rawString: "",
+			err:       datapipe.ErrDataTagNotFound,
+		},
+		{
+			rawString: "{[]}",
+			err:       datapipe.ErrJSONDecoderInternal,
+		},
+		{
+			rawString: "{\"data\": \"oops\"}",
+			err:       datapipe.ErrInvalidDataTag,
+		},
+		{
+			rawString: "{\"nothing\": [}",
+			err:       datapipe.ErrJSONDecoderInternal,
+		},
+		{
+			rawString: `{"meta": {"methods": 0, "type": "sessions", "count": 0, "version": 5}, "data": []}`,
+			err:       nil,
+		},
+		{
+			rawString: `{"test": {"data": {}}, "meta": {"methods": 0, "type": "sessions", "count": 0, "version": 5}, "data": []}`,
+			err:       nil,
+		},
+	}
+
+	for _, assertion := range assertions {
+		r := strings.NewReader(assertion.rawString)
+		j := json.NewDecoder(r)
+
+		err := datapipe.SeekToDataTag(j)
+		assert.ErrorIs(t, err, assertion.err)
+	}
+}