From 8ecfe02e42f6d794e2d71a57ec671f33cc9b205e Mon Sep 17 00:00:00 2001
From: seph <seph@directionless.org>
Date: Mon, 22 Feb 2021 00:57:46 -0500
Subject: [PATCH] db stats table

---
 pkg/agent/dbdump.go                   | 66 +++++++++++++++++++++++++++
 pkg/osquery/table/launcher_db_info.go | 56 +++++++++++++++++++++++
 pkg/osquery/table/table.go            |  1 +
 3 files changed, 123 insertions(+)
 create mode 100644 pkg/agent/dbdump.go
 create mode 100644 pkg/osquery/table/launcher_db_info.go

diff --git a/pkg/agent/dbdump.go b/pkg/agent/dbdump.go
new file mode 100644
index 0000000000..5ad6788649
--- /dev/null
+++ b/pkg/agent/dbdump.go
@@ -0,0 +1,66 @@
+package agent
+
+import (
+	"github.com/pkg/errors"
+	"go.etcd.io/bbolt"
+)
+
+type bucketStatsHolder struct {
+	Stats        bbolt.BucketStats
+	FillPercent  float64
+	NumberOfKeys uint
+	Size         uint
+}
+
+type dbStatsHolder struct {
+	Stats bbolt.TxStats
+	Size  int64
+}
+
+type Stats struct {
+	DB      dbStatsHolder
+	Buckets map[string]bucketStatsHolder
+}
+
+func GetStats(db *bbolt.DB) (*Stats, error) {
+	stats := &Stats{
+		Buckets: make(map[string]bucketStatsHolder),
+	}
+
+	if err := db.View(func(tx *bbolt.Tx) error {
+		stats.DB.Stats = tx.Stats()
+		stats.DB.Size = tx.Size()
+
+		if err := tx.ForEach(bucketStatsFunc(stats)); err != nil {
+			return errors.Wrap(err, "dumping bucket")
+		}
+		return nil
+	}); err != nil {
+		return nil, errors.Wrap(err, "creating view tx")
+	}
+
+	return stats, nil
+}
+
+func bucketStatsFunc(stats *Stats) func([]byte, *bbolt.Bucket) error {
+	return func(name []byte, b *bbolt.Bucket) error {
+		numKeys := uint(0)
+		sizeValues := uint(0)
+		if err := b.ForEach(func(_, v []byte) error {
+			numKeys = numKeys + 1
+			sizeValues = sizeValues + uint(len(v))
+			return nil
+		}); err != nil {
+			return errors.Wrap(err, "foreach kv")
+		}
+
+		stats.Buckets[string(name)] = bucketStatsHolder{
+			Stats:        b.Stats(),
+			FillPercent:  b.FillPercent,
+			NumberOfKeys: numKeys,
+			Size:         sizeValues,
+		}
+
+		return nil
+	}
+}
diff --git a/pkg/osquery/table/launcher_db_info.go b/pkg/osquery/table/launcher_db_info.go
new file mode 100644
index 0000000000..05f40769e9
--- /dev/null
+++ b/pkg/osquery/table/launcher_db_info.go
@@ -0,0 +1,56 @@
+package table
+
+import (
+	"context"
+	"encoding/json"
+	"strings"
+
+	"github.com/kolide/launcher/pkg/agent"
+	"github.com/kolide/launcher/pkg/dataflatten"
+	"github.com/kolide/launcher/pkg/osquery/tables/dataflattentable"
+	"github.com/kolide/launcher/pkg/osquery/tables/tablehelpers"
+	"github.com/kolide/osquery-go/plugin/table"
+	"github.com/pkg/errors"
+	"go.etcd.io/bbolt"
+)
+
+func LauncherDbInfo(db *bbolt.DB) *table.Plugin {
+	columns := dataflattentable.Columns()
+	return table.NewPlugin("kolide_launcher_db_info", columns, generateLauncherDbInfo(db))
+}
+
+func generateLauncherDbInfo(db *bbolt.DB) table.GenerateFunc {
+	return func(ctx context.Context, queryContext table.QueryContext) ([]map[string]string, error) {
+		stats, err := agent.GetStats(db)
+		if err != nil {
+			return nil, err
+		}
+
+		statsJson, err := json.Marshal(stats)
+		if err != nil {
+			return nil, err
+		}
+
+		var results []map[string]string
+
+		for _, dataQuery := range tablehelpers.GetConstraints(queryContext, "query", tablehelpers.WithDefaults("*")) {
+			flattenOpts := []dataflatten.FlattenOpts{
+				dataflatten.WithQuery(strings.Split(dataQuery, "/")),
+			}
+
+			rows, err := dataflatten.Json(statsJson, flattenOpts...)
+
+			if err != nil {
+				// In other places, we'd log the error
+				// and continue. We don't have logs
+				// here, so we'll just return it.
+				return results, errors.Wrapf(err, "flattening with query %s", dataQuery)
+
+			}
+
+			results = append(results, dataflattentable.ToMap(rows, dataQuery, nil)...)
+		}
+
+		return results, nil
+	}
+}
diff --git a/pkg/osquery/table/table.go b/pkg/osquery/table/table.go
index 121d6c5aeb..b4ac2d3e14 100644
--- a/pkg/osquery/table/table.go
+++ b/pkg/osquery/table/table.go
@@ -14,6 +14,7 @@ import (
 func LauncherTables(db *bbolt.DB, opts *launcher.Options) []osquery.OsqueryPlugin {
 	return []osquery.OsqueryPlugin{
 		LauncherConfigTable(db),
+		LauncherDbInfo(db),
 		LauncherIdentifierTable(db),
 		TargetMembershipTable(db),
 		LauncherAutoupdateConfigTable(opts),