readBolt.go

package main

import (
	"bytes"
	"errors"
	"flag"
	"fmt"
	"log"
	"os"

	"github.com/boltdb/bolt"
	"github.com/groob/plist"
	"github.com/micromdm/micromdm/mdm/mdm"
	"github.com/micromdm/micromdm/platform/queue"
	"github.com/olekukonko/tablewriter"
	"go.mozilla.org/pkcs7"
)

var (
	dbPath string
)

type myboltDB struct {
	*bolt.DB
}

type profPayload struct {
	PayloadIdentifier string
}

// DB is defined in github.com/micromdm/micromdm/platform/device

// amount each type of payload per device
// device uuid
// average length of queue for a device and outliers
// unmarshal InstallProfile (profile identifier and uuid)
// api endpoint for ^
// add micro flag for param to drop commands after X days

func requestTypeTable(requestTypeMap map[string]*requestTypeMetrics) *tablewriter.Table {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"Request Type", "Count", "Highest Length"})
	for k, v := range requestTypeMap {
		table.Append([]string{k, fmt.Sprint(v.Count), fmt.Sprint(v.HighestLength)})
	}
	return table
}

func profileIdentifierTable(profileIdentifierMap map[string]*profileIdentifierMetrics) *tablewriter.Table {
	table := tablewriter.NewWriter(os.Stdout)
	table.SetHeader([]string{"ProfileIdentifier", "Count"})
	for k, v := range profileIdentifierMap {
		table.Append([]string{k, fmt.Sprint(v.Count)})
	}
	return table
}

func pkcs7strip(data []byte, blockSize int) ([]byte, error) {
	length := len(data)
	if length == 0 {
		return nil, errors.New("pkcs7: Data is empty")
	}
	if length%blockSize != 0 {
		return nil, errors.New("pkcs7: Data is not block-aligned")
	}
	padLen := int(data[length-1])
	ref := bytes.Repeat([]byte{byte(padLen)}, padLen)
	if padLen > blockSize || padLen == 0 || !bytes.HasSuffix(data, ref) {
		return nil, errors.New("pkcs7: Invalid padding")
	}
	return data[:length-padLen], nil
}

func numCmdsPerDeviceTable(numCmdsPerDevice map[string]int) *tablewriter.Table {
	table := tablewriter.NewWriter(os.Stdout)
	outlierThreshold := 10
	outlierTitle := fmt.Sprintf("Devices w/ %v more commands than avg", outlierThreshold)
	table.SetHeader([]string{"Avg # Commands Per Device", outlierTitle, "# of Devices"})
	totalNumDevices := len(numCmdsPerDevice)
	totalNumCmds := 0
	for _, numForSingleDevice := range numCmdsPerDevice {
		totalNumCmds += numForSingleDevice
	}
	// Average number of Commands Per device
	avg := totalNumCmds / totalNumDevices

	numOfOutliers := 0
	for _, numForSingleDevice := range numCmdsPerDevice {
		if numForSingleDevice > avg+outlierThreshold {
			numOfOutliers++
		}
	}
	table.Append([]string{fmt.Sprint(avg), fmt.Sprint(numOfOutliers), fmt.Sprint(totalNumDevices)})
	return table
}

type requestTypeMetrics struct {
	Count              int
	HighestLength      int
	LargestProfileType string
}

type profileIdentifierMetrics struct {
	Count int
}

func (db myboltDB) listCommandsViaTables(bucketName string) []*tablewriter.Table {
	allTheTables := []*tablewriter.Table{}
	requestTypeMap := map[string]*requestTypeMetrics{}
	profileIdentifierMap := map[string]*profileIdentifierMetrics{}
	numCmdsPerDevice := map[string]int{}

	var commands []queue.DeviceCommand
	err := db.View(func(tx *bolt.Tx) error {
		b := tx.Bucket([]byte(bucketName))
		return b.ForEach(func(k, v []byte) error {
			var deviceCommand queue.DeviceCommand
			if err := queue.UnmarshalDeviceCommand(v, &deviceCommand); err != nil {
				return err
			}
			//fmt.Println("deviceUUID: ", deviceCommand.DeviceUDID)
			if _, ok := numCmdsPerDevice[deviceCommand.DeviceUDID]; ok {
				numCmdsPerDevice[deviceCommand.DeviceUDID]++
			} else {
				numCmdsPerDevice[deviceCommand.DeviceUDID] = 1
			}
			for _, c := range deviceCommand.Commands {
				numCmdsPerDevice[deviceCommand.DeviceUDID]++
				var commandPayload mdm.CommandPayload
				if c.Payload != nil {
					if err := plist.Unmarshal(c.Payload, &commandPayload); err != nil {
						return err
					}
					if commandPayload.Command.RequestType == "InstallProfile" {
						signedPayload, err := pkcs7.Parse(commandPayload.Command.InstallProfile.Payload)
						if err != nil {
							return err
						}
						//fmt.Println(string(signedPayload.Content))
						var profPay profPayload
						if err := plist.Unmarshal(signedPayload.Content, &profPay); err != nil {
							return err
						}
						if _, ok := profileIdentifierMap[profPay.PayloadIdentifier]; ok {
							count := profileIdentifierMap[profPay.PayloadIdentifier].Count
							profileIdentifierMap[profPay.PayloadIdentifier] = &profileIdentifierMetrics{
								Count: count + 1,
							}
						} else {
							profileIdentifierMap[profPay.PayloadIdentifier] = &profileIdentifierMetrics{
								Count: 1,
							}
						}
						//fmt.Println(profPay.PayloadIdentifier)
					} //else {
					// 	fmt.Println(string(c.Payload))
					// }
					if _, ok := requestTypeMap[commandPayload.Command.RequestType]; ok {
						currentCount := requestTypeMap[commandPayload.Command.RequestType].Count
						highestLength := requestTypeMap[commandPayload.Command.RequestType].HighestLength
						if len(c.Payload) > highestLength {
							highestLength = len(c.Payload)
						}
						requestTypeMap[commandPayload.Command.RequestType] = &requestTypeMetrics{
							Count:         currentCount + 1,
							HighestLength: highestLength,
						}

					} else {
						requestTypeMap[commandPayload.Command.RequestType] = &requestTypeMetrics{
							Count:         1,
							HighestLength: len(c.Payload),
						}
					}
				}
			}
			commands = append(commands, deviceCommand)
			//fmt.Println("")
			return nil
		})
	})
	if err != nil {
		fmt.Println(err)
	}
	allTheTables = append(allTheTables, requestTypeTable(requestTypeMap))
	allTheTables = append(allTheTables, numCmdsPerDeviceTable(numCmdsPerDevice))
	allTheTables = append(allTheTables, profileIdentifierTable(profileIdentifierMap))

	return allTheTables
}

func main() {

	flag.StringVar(&dbPath, "db_path", "", "help message for flagname")
	flag.Parse()
	db, err := bolt.Open(dbPath, 0600, nil)
	if err != nil {
		log.Fatal(err)
	}
	// embed *bolt.DB in myBoltDB so we can add methods
	a := myboltDB{db}
	tables := a.listCommandsViaTables(queue.DeviceCommandBucket)
	for _, table := range tables {
		table.Render()
	}
	defer db.Close()
}