From 1bc0d08a84ce557bf1168973ae4e984a702d747f Mon Sep 17 00:00:00 2001
From: =?UTF-8?q?Erik=20Hedenstr=C3=B6m?= <erik@hedenstroem.com>
Date: Fri, 4 Feb 2022 00:44:47 +0100
Subject: [PATCH] Added more commands
---
cmd/clients.go | 35 +++++++++++++++
cmd/devices.go | 27 ++++++++++++
cmd/dhcp_reservations.go | 92 ----------------------------------------
cmd/dump.go | 34 +++++++++++++++
cmd/list.go | 48 +++++++++++++++++++++
cmd/ping.go | 25 +++++++++++
cmd/root.go | 9 ++--
cmd/version.go | 21 +++++++++
ssh/tunnel.go | 4 +-
utils/document.go | 80 ++++++++++++++++++++++++++++++++++
10 files changed, 279 insertions(+), 96 deletions(-)
create mode 100644 cmd/clients.go
create mode 100644 cmd/devices.go
delete mode 100644 cmd/dhcp_reservations.go
create mode 100644 cmd/dump.go
create mode 100644 cmd/list.go
create mode 100644 cmd/ping.go
create mode 100644 cmd/version.go
create mode 100644 utils/document.go
diff --git a/cmd/clients.go b/cmd/clients.go
new file mode 100644
index 0000000..49439cd
--- /dev/null
+++ b/cmd/clients.go
@@ -0,0 +1,35 @@
+package cmd
+
+import (
+ "github.com/spf13/cobra"
+ "gitlab.hedenstroem.com/go/udm-query/utils"
+ "go.mongodb.org/mongo-driver/bson"
+)
+
+var fixedIp bool
+
+var clientsCmd = &cobra.Command{
+ Use: "clients",
+ Short: "clients",
+ Long: `clients`,
+ RunE: func(cmd *cobra.Command, args []string) error {
+ ace := client.Database("ace")
+ clients := ace.Collection("user")
+ filter := bson.D{}
+ if fixedIp {
+ filter = bson.D{{"use_fixedip", true}}
+ }
+ docs, err := utils.FindDocuments(clients, filter)
+ if err != nil {
+ return err
+ }
+ utils.IpSort(docs, "fixed_ip")
+ utils.DocumentsToTable(docs, []string{"fixed_ip", "name", "hostname", "mac"})
+ return nil
+ },
+}
+
+func init() {
+ clientsCmd.Flags().BoolVarP(&fixedIp, "fixed", "f", false, "Show debug output")
+ RootCmd.AddCommand(clientsCmd)
+}
diff --git a/cmd/devices.go b/cmd/devices.go
new file mode 100644
index 0000000..f901c30
--- /dev/null
+++ b/cmd/devices.go
@@ -0,0 +1,27 @@
+package cmd
+
+import (
+ "github.com/spf13/cobra"
+ "gitlab.hedenstroem.com/go/udm-query/utils"
+)
+
+var devicesCmd = &cobra.Command{
+ Use: "devices",
+ Short: "devices",
+ Long: `devices`,
+ RunE: func(cmd *cobra.Command, args []string) error {
+ ace := client.Database("ace")
+ devices := ace.Collection("device")
+ docs, err := utils.AllDocuments(devices)
+ if err != nil {
+ return err
+ }
+ utils.IpSort(docs, "ip")
+ utils.DocumentsToTable(docs, []string{"ip", "model", "version", "serial", "mac"})
+ return nil
+ },
+}
+
+func init() {
+ RootCmd.AddCommand(devicesCmd)
+}
diff --git a/cmd/dhcp_reservations.go b/cmd/dhcp_reservations.go
deleted file mode 100644
index ab7638f..0000000
--- a/cmd/dhcp_reservations.go
+++ /dev/null
@@ -1,92 +0,0 @@
-package cmd
-
-import (
- "context"
- "encoding/json"
- "fmt"
- "log"
- "os"
-
- "github.com/davecgh/go-spew/spew"
- "github.com/jedib0t/go-pretty/v6/table"
- "github.com/spf13/cobra"
- "go.mongodb.org/mongo-driver/bson"
- "go.mongodb.org/mongo-driver/mongo/options"
- "go.mongodb.org/mongo-driver/mongo/readpref"
-)
-
-var reservationsCmd = &cobra.Command{
- Use: "reservations",
- Short: "DHCP static reservations",
- Long: `Show all DHCP static reservations assignments`,
- RunE: func(cmd *cobra.Command, args []string) error {
-
- // Ping the primary
- if err := client.Ping(context.TODO(), readpref.Primary()); err != nil {
- panic(err)
- }
-
- ace := client.Database("ace")
- devices := ace.Collection("device")
- user := ace.Collection("user")
-
- /*
- filter := bson.D{
- {"$and",
- bson.A{
- bson.D{{"rating", bson.D{{"$gt", 5}}}},
- bson.D{{"rating", bson.D{{"$lt", 10}}}},
- }},
- }
- */
-
- // _filter := bson.D{{"use_fixedip", true}}
- opts := options.Find().SetSort(bson.D{{"ip", 1}})
- filter := bson.D{}
- cursor, _ := devices.Find(context.TODO(), filter)
- for cursor.Next(context.TODO()) {
- var result bson.D
- if err := cursor.Decode(&result); err != nil {
- log.Fatal(err)
- }
- b, _ := bson.MarshalExtJSONIndent(result, false, true, "", " ")
- fmt.Println("D----------------------------------")
- fmt.Println(string(b))
- }
-
- t := table.NewWriter()
- t.SetOutputMirror(os.Stdout)
- t.AppendHeader(table.Row{"name", "hostname", "ip"})
-
- opts = options.Find().SetSort(bson.D{{"fixed_ip", 1}})
- filter = bson.D{{"use_fixedip", true}}
- cursor, _ = user.Find(context.TODO(), filter, opts)
- for cursor.Next(context.TODO()) {
- var result bson.D
- if err := cursor.Decode(&result); err != nil {
- log.Fatal(err)
- }
- b, _ := bson.MarshalExtJSON(result, false, true)
- // b, _ := bson.MarshalExtJSONIndent(result, false, true, "", " ")
- var v map[string]interface{}
- json.Unmarshal(b, &v)
- t.AppendRows([]table.Row{{v["name"], v["hostname"], v["fixed_ip"]}})
- fmt.Println("U----------------------------------")
- fmt.Println(string(b))
- spew.Dump(v)
- }
- t.Render()
-
- // spew.Dump(docs)
- /*
- dbs, _ := client.ListDatabaseNames(context.TODO(), bson.D{}, nil)
- fmt.Println("Successfully connected and pinged.")
- spew.Dump(dbs)
- */
- return nil
- },
-}
-
-func init() {
- RootCmd.AddCommand(reservationsCmd)
-}
diff --git a/cmd/dump.go b/cmd/dump.go
new file mode 100644
index 0000000..54c2e71
--- /dev/null
+++ b/cmd/dump.go
@@ -0,0 +1,34 @@
+package cmd
+
+import (
+ "encoding/json"
+ "errors"
+ "fmt"
+
+ "github.com/spf13/cobra"
+ "gitlab.hedenstroem.com/go/udm-query/utils"
+)
+
+var dumpCmd = &cobra.Command{
+ Use: "dump [flags] db collection",
+ Short: "dump",
+ Long: `dump`,
+ RunE: func(cmd *cobra.Command, args []string) error {
+ if len(args) != 2 {
+ return errors.New("expected 2 arguments")
+ }
+ db := client.Database(args[0])
+ collection := db.Collection(args[1])
+ docs, err := utils.AllDocuments(collection)
+ if err != nil {
+ return err
+ }
+ b, _ := json.MarshalIndent(docs, "", " ")
+ fmt.Println(string(b))
+ return nil
+ },
+}
+
+func init() {
+ RootCmd.AddCommand(dumpCmd)
+}
diff --git a/cmd/list.go b/cmd/list.go
new file mode 100644
index 0000000..c9e67fd
--- /dev/null
+++ b/cmd/list.go
@@ -0,0 +1,48 @@
+package cmd
+
+import (
+ "context"
+ "fmt"
+ "os"
+
+ "github.com/jedib0t/go-pretty/v6/list"
+ "github.com/spf13/cobra"
+ "github.com/spf13/viper"
+ "go.mongodb.org/mongo-driver/bson"
+)
+
+var listCmd = &cobra.Command{
+ Use: "list",
+ Short: "list",
+ Long: `list`,
+ RunE: func(cmd *cobra.Command, args []string) error {
+ l := list.NewWriter()
+ l.AppendItem(viper.GetString("ADDRESS"))
+ l.Indent()
+ result, err := client.ListDatabases(context.TODO(), bson.D{})
+ if err != nil {
+ return err
+ }
+ for _, dbSpec := range result.Databases {
+ l.AppendItem(fmt.Sprintf("%s:%dMB", dbSpec.Name, dbSpec.SizeOnDisk/1000000))
+ l.Indent()
+ db := client.Database(dbSpec.Name)
+ collectionNames, err := db.ListCollectionNames(context.TODO(), bson.D{})
+ if err != nil {
+ return err
+ }
+ for _, collectionName := range collectionNames {
+ l.AppendItem(collectionName)
+ }
+ l.UnIndent()
+ }
+ l.SetStyle(list.StyleConnectedRounded)
+ l.SetOutputMirror(os.Stdout)
+ l.Render()
+ return nil
+ },
+}
+
+func init() {
+ RootCmd.AddCommand(listCmd)
+}
diff --git a/cmd/ping.go b/cmd/ping.go
new file mode 100644
index 0000000..8038136
--- /dev/null
+++ b/cmd/ping.go
@@ -0,0 +1,25 @@
+package cmd
+
+import (
+ "context"
+ "fmt"
+
+ "github.com/spf13/cobra"
+ "go.mongodb.org/mongo-driver/mongo/readpref"
+)
+
+var pingCmd = &cobra.Command{
+ Use: "ping",
+ Short: "Ping the database",
+ Long: `Sends no-op command to check if the database is responding to commands`,
+ RunE: func(cmd *cobra.Command, args []string) (err error) {
+ if err = client.Ping(context.TODO(), readpref.Primary()); err == nil {
+ fmt.Println("success")
+ }
+ return
+ },
+}
+
+func init() {
+ RootCmd.AddCommand(pingCmd)
+}
diff --git a/cmd/root.go b/cmd/root.go
index 4deffc2..be0cdbd 100644
--- a/cmd/root.go
+++ b/cmd/root.go
@@ -10,24 +10,26 @@ import (
"github.com/joho/godotenv"
"github.com/spf13/cobra"
"github.com/spf13/viper"
- "gitlab.hedenstroem.com/go/udm-query/constant"
"gitlab.hedenstroem.com/go/udm-query/ssh"
"go.mongodb.org/mongo-driver/mongo"
"go.mongodb.org/mongo-driver/mongo/options"
)
+var debug bool
var client *mongo.Client
var RootCmd = &cobra.Command{
Use: "udm-query",
- Long: `UDM query tool ` + constant.Version,
+ Long: `Tool to query Ubiquiti UniFi Dream Machines`,
PersistentPreRunE: func(cmd *cobra.Command, args []string) (err error) {
tunnel := ssh.NewSSHTunnel(
viper.GetString("ADDRESS"),
ssh.Password(viper.GetString("PASSWORD")),
"127.0.0.1:27117",
)
- tunnel.Log = log.New(os.Stdout, "", log.Ldate|log.Lmicroseconds)
+ if debug {
+ tunnel.Log = log.New(os.Stderr, "", log.Ldate|log.Lmicroseconds)
+ }
go tunnel.Start()
time.Sleep(100 * time.Millisecond)
uri := fmt.Sprintf("mongodb://127.0.0.1:%d/", tunnel.Local.Port)
@@ -50,6 +52,7 @@ func init() {
cobra.OnInitialize(initEnv)
RootCmd.PersistentFlags().StringP("address", "a", "192.168.1.1", "Address to the USM SSH server")
RootCmd.PersistentFlags().StringP("password", "p", "", "SSH password")
+ RootCmd.PersistentFlags().BoolVarP(&debug, "debug", "d", false, "Show debug output")
viper.SetEnvPrefix("SSH")
viper.BindPFlag("ADDRESS", RootCmd.PersistentFlags().Lookup("address"))
viper.BindPFlag("PASSWORD", RootCmd.PersistentFlags().Lookup("password"))
diff --git a/cmd/version.go b/cmd/version.go
new file mode 100644
index 0000000..841a3f3
--- /dev/null
+++ b/cmd/version.go
@@ -0,0 +1,21 @@
+package cmd
+
+import (
+ "fmt"
+
+ "github.com/spf13/cobra"
+ "gitlab.hedenstroem.com/go/udm-query/constant"
+)
+
+var versionCmd = &cobra.Command{
+ Use: "version",
+ Short: "Display version",
+ Long: `Displays the version of the udm-query tool`,
+ Run: func(cmd *cobra.Command, args []string) {
+ fmt.Println(constant.Version)
+ },
+}
+
+func init() {
+ RootCmd.AddCommand(versionCmd)
+}
diff --git a/ssh/tunnel.go b/ssh/tunnel.go
index 6be4768..f5f363f 100644
--- a/ssh/tunnel.go
+++ b/ssh/tunnel.go
@@ -30,12 +30,13 @@ func (tunnel *SSHTunnel) Start() error {
}
defer listener.Close()
tunnel.Local.Port = listener.Addr().(*net.TCPAddr).Port
+ tunnel.logf("accepting connections at %s", tunnel.Local.String())
for {
conn, err := listener.Accept()
if err != nil {
return err
}
- tunnel.logf("accepted connection")
+ tunnel.logf("accepted connection from %s", conn.RemoteAddr())
go tunnel.forward(conn)
}
}
@@ -59,6 +60,7 @@ func (tunnel *SSHTunnel) forward(localConn net.Conn) {
tunnel.logf("io.Copy error: %s", err)
}
}
+ tunnel.logf("%s -> %s -> %s -> %s\n", localConn.RemoteAddr(), localConn.LocalAddr(), tunnel.Server, tunnel.Remote)
go copyConn(localConn, remoteConn)
go copyConn(remoteConn, localConn)
}
diff --git a/utils/document.go b/utils/document.go
new file mode 100644
index 0000000..5e3ddb2
--- /dev/null
+++ b/utils/document.go
@@ -0,0 +1,80 @@
+package utils
+
+import (
+ "bytes"
+ "context"
+ "encoding/json"
+ "net"
+ "os"
+ "sort"
+
+ "github.com/jedib0t/go-pretty/v6/table"
+ "go.mongodb.org/mongo-driver/bson"
+ "go.mongodb.org/mongo-driver/mongo"
+ "go.mongodb.org/mongo-driver/mongo/options"
+)
+
+func AllDocuments(collection *mongo.Collection) ([]map[string]interface{}, error) {
+ return FindDocuments(collection, bson.D{})
+}
+
+func FindDocuments(collection *mongo.Collection, filter interface{}, opts ...*options.FindOptions) ([]map[string]interface{}, error) {
+ var docs []map[string]interface{}
+ cursor, err := collection.Find(context.TODO(), filter)
+ if err != nil {
+ return nil, err
+ }
+ defer cursor.Close(context.TODO())
+ for cursor.Next(context.TODO()) {
+ var result bson.D
+ if err := cursor.Decode(&result); err != nil {
+ return nil, err
+ }
+ b, err := bson.MarshalExtJSON(result, false, true)
+ if err != nil {
+ return nil, err
+ }
+ var v map[string]interface{}
+ json.Unmarshal(b, &v)
+ docs = append(docs, v)
+ }
+ return docs, nil
+}
+
+func DocumentsToTable(docs []map[string]interface{}, fields []string) {
+ t := table.NewWriter()
+ header := table.Row{}
+ for _, field := range fields {
+ header = append(header, field)
+ }
+ t.AppendHeader(header)
+ for _, doc := range docs {
+ row := table.Row{}
+ for _, field := range fields {
+ value, exists := doc[field]
+ if exists {
+ row = append(row, value)
+ } else {
+ row = append(row, "-")
+ }
+ }
+ t.AppendRows([]table.Row{row})
+ }
+ t.SetOutputMirror(os.Stdout)
+ t.SetStyle(table.StyleColoredBright)
+ t.Render()
+}
+
+func IpSort(docs []map[string]interface{}, field string) {
+ sort.SliceStable(docs, func(i, j int) bool {
+ ip_i := net.IPv4bcast
+ if str, exists := docs[i][field]; exists {
+ ip_i = net.ParseIP(str.(string))
+ }
+ ip_j := net.IPv4bcast
+ if str, exists := docs[j][field]; exists {
+ ip_j = net.ParseIP(str.(string))
+ }
+ return bytes.Compare(ip_i, ip_j) < 0
+ })
+}
--
GitLab